src/lib/utils/rng.ts

/**
 * Deterministic PRNG utilities shared by all watermark generators and detectors.
 *
 * We deliberately do NOT replicate torch's MT19937 / randperm bit-for-bit
 * (impractical in JS). A watermark only requires that the generator and the
 * detector share the same PRF/PRNG stack. Deviations from the reference
 * PyTorch implementations are documented in the README.
 */

const MASK64 = (1n << 64n) - 1n;

/** splitmix64: high-quality 64-bit mixer. Returns the next state and output. */
export function splitmix64(state: bigint): { state: bigint; out: bigint } {
  let s = (state + 0x9e3779b97f4a7c15n) & MASK64;
  let z = s;
  z = ((z ^ (z >> 30n)) * 0xbf58476d1ce4e5b9n) & MASK64;
  z = ((z ^ (z >> 27n)) * 0x94d049bb133111ebn) & MASK64;
  z = z ^ (z >> 31n);
  return { state: s, out: z & MASK64 };
}

/** One-shot 64-bit hash of arbitrary bigint input (stateless). */
export function mix64(x: bigint): bigint {
  return splitmix64(x & MASK64).out;
}

/** Combine multiple 64-bit values into one (order-sensitive). */
export function hashCombine(...values: bigint[]): bigint {
  let h = 0x51_7c_c1_b7_27_22_0a_95n; // arbitrary non-zero start
  for (const v of values) {
    h = mix64((h ^ (v & MASK64)) & MASK64);
  }
  return h;
}

/** Map a 64-bit value to a float in [0, 1) using the top 53 bits. */
export function toUnitFloat(x: bigint): number {
  return Number((x & MASK64) >> 11n) / 2 ** 53;
}

/**
 * Seedable uniform stream (counter-based on splitmix64).
 * Counter-based design means identical seeds yield identical streams
 * regardless of platform - this is the "shared base RNG stream" used to
 * separate watermark effects from sampling randomness.
 */
export class RandomStream {
  private state: bigint;

  constructor(seed: bigint | number) {
    this.state = mix64(BigInt(seed) & MASK64);
  }

  /** Next uniform float in [0, 1). */
  next(): number {
    const { state, out } = splitmix64(this.state);
    this.state = state;
    return toUnitFloat(out);
  }

  /** Next 64-bit integer. */
  nextU64(): bigint {
    const { state, out } = splitmix64(this.state);
    this.state = state;
    return out;
  }
}

/**
 * Seeded Fisher-Yates: choose `count` distinct integers from [0, n).
 * Used for k-SemStamp valid-cluster selection (small n = K clusters).
 */
export function seededSample(seed: bigint, n: number, count: number): number[] {
  const stream = new RandomStream(seed);
  const arr = Array.from({ length: n }, (_, i) => i);
  for (let i = n - 1; i > 0; i--) {
    const j = Math.floor(stream.next() * (i + 1));
    [arr[i], arr[j]] = [arr[j], arr[i]];
  }
  return arr.slice(0, count);
}