sequence-js/src/random/index.ts

import { BitArray } from "../bitarray/index.js";
import { asArray } from "../utils.js";
import { AsyncRandomOptions, ElementPredicate, ElementWeight, RandomGenerator, RandomOptions } from "./types.js";

export const alwaysTrue: ElementPredicate = () => true;
export const weightOfOne: ElementWeight = () => 1.0;
export const mathRandom: RandomGenerator = () => Math.random();

const defaultOptions = Object.freeze<Required<RandomOptions>>({
  predicate: alwaysTrue,
  weight: weightOfOne,
  random: mathRandom
});

function mergeOptions<T>(first: RandomOptions<T> | undefined, second: RandomOptions<T> | undefined): RandomOptions<T> | undefined {
  if (!first) {
    return second;
  }

  if (!second) {
    return first;
  }

  const firstPredicate = first.predicate;
  const secondPredicate = second.predicate;

  return {
    predicate: firstPredicate ? secondPredicate ? (i, o) => firstPredicate(i, o) && secondPredicate(i, o) : firstPredicate : secondPredicate,
    weight: first.weight ?? second.weight,
    random: first.random ?? second.random
  };
}

function withDefaultOptions<T>(options: RandomOptions<T> | undefined): Required<RandomOptions<T>>;
function withDefaultOptions<T>(options: AsyncRandomOptions<T> | undefined): Required<AsyncRandomOptions<T>>;
function withDefaultOptions<T>(options: RandomOptions<T> | AsyncRandomOptions<T> | undefined): Required<RandomOptions<T> | AsyncRandomOptions<T>> {
  if (!options || options === defaultOptions) {
    return defaultOptions;
  }

  return {
    predicate: options.predicate ?? defaultOptions.predicate,
    weight: options.weight ?? defaultOptions.weight,
    random: options.random ?? defaultOptions.random
  };
}

function _getRandomElement<T>(sequence: Iterable<T>, options: Required<RandomOptions<T>>) {
  const { predicate, weight, random } = options;

  let result: T | undefined = undefined;
  let resultIndex = -1;
  let index = 0;
  let weightAcc = 0.0;

  for (const element of sequence) {
    const currentIndex = index++;

    if (predicate(currentIndex, element)) {
      const w = weight(currentIndex, element);

      if (w <= 0.0) {
        continue;
      }

      weightAcc += w;

      if (random() * weightAcc < w) {
        result = element;
        resultIndex = currentIndex;
      }
    }
  }

  return { element: result, index: resultIndex };
}

export function getRandomElement<T>(sequence: Iterable<T>, options?: RandomOptions<T>) {
  return _getRandomElement(sequence, withDefaultOptions(options));
}

async function _getRandomElementAsync<T>(sequence: AsyncIterable<T>, options: Required<AsyncRandomOptions<T>>) {
  const { predicate, weight, random } = options;

  let result: T | undefined = undefined;
  let resultIndex = -1;
  let index = 0;
  let weightAcc = 0.0;

  for await (const element of sequence) {
    const currentIndex = index++;

    if (await predicate(currentIndex, element)) {
      const w = await weight(currentIndex, element);

      if (w <= 0.0) {
        continue;
      }

      weightAcc += w;

      if (random() * weightAcc < w) {
        result = element;
        resultIndex = currentIndex;
      }
    }
  }

  return { element: result, index: resultIndex };
}

export async function getRandomElementAsync<T>(sequence: AsyncIterable<T>, options?: AsyncRandomOptions<T>) {
  return await _getRandomElementAsync(sequence, withDefaultOptions(options));
}

export class RandomPicker<T> {
  readonly #elements: Iterable<T>;
  readonly #flags: BitArray;
  readonly #options: Required<RandomOptions<T>>;

  public constructor(elements: Iterable<T>, length: number, options?: RandomOptions<T>) {
    this.#elements = elements;
    this.#flags = BitArray.create(length);
    this.#options = withDefaultOptions(mergeOptions({ predicate: i => this.#flags.get(i) }, options));

    this.reset();
  }

  public static from<T>(iterable: Iterable<T>, options?: RandomOptions<T>) {
    const arr = asArray(iterable);
    return new this<T>(arr, arr.length, options);
  }

  public static of<T>(options?: RandomOptions<T>, ...values: T[]) {
    return new this<T>(values, values.length, options);
  }

  public get state() {
    return this.#flags;
  }

  public reset() {
    this.#flags.fill(true);
  }

  public next() {
    const result = _getRandomElement(this.#elements, this.#options);

    if (result.index < 0) {
      this.reset();
    } else {
      this.#flags.set(result.index, false);

      if (this.#flags.isEmpty()) {
        this.reset();
      }
    }

    return result.element;
  }
}