prng.ts

/**
 * Seeded Pseudo-Random Number Generator (PRNG)
 *
 * Uses Xoroshiro128+ algorithm for deterministic sampling.
 * Same seed produces identical sequences - critical for reproducibility.
 *
 * Based on reference implementation:
 * http://prng.di.unimi.it/xoroshiro128plus.c
 *
 * Properties:
 * - Fast: ~0.01ms per call
 * - High quality: Passes BigCrush statistical tests
 * - Deterministic: Same seed = same sequence
 * - Matches native RNG behavior
 */

/**
 * Xoroshiro128+ PRNG implementation
 *
 * State: 128 bits (two 64-bit values)
 * Period: 2^128 - 1
 * Output: Uniformly distributed [0, 1) floats
 */
export class Xoroshiro128Plus {
  private state0: bigint;
  private state1: bigint;

  /**
   * Initialize PRNG with seed
   * @param seed 32-bit seed value
   */
  constructor(seed: number) {
    // Initialize state from seed using SplitMix64
    // This ensures good initial state distribution
    let s = BigInt(seed);
    this.state0 = this.splitmix64(s);
    this.state1 = this.splitmix64(s + 1n);
  }

  /**
   * SplitMix64 initialization
   *
   * Converts 64-bit seed to well-distributed state.
   * Used to initialize Xoroshiro128+ from simple integer seed.
   */
  private splitmix64(x: bigint): bigint {
    // Magic constants from SplitMix64 paper
    x = (x + 0x9e3779b97f4a7c15n) & 0xffffffffffffffffn;
    x = ((x ^ (x >> 30n)) * 0xbf58476d1ce4e5b9n) & 0xffffffffffffffffn;
    x = ((x ^ (x >> 27n)) * 0x94d049bb133111ebn) & 0xffffffffffffffffn;
    return (x ^ (x >> 31n)) & 0xffffffffffffffffn;
  }

  /**
   * Generate next random number [0, 1)
   *
   * Xoroshiro128+ algorithm:
   * 1. result = state0 + state1
   * 2. rotate/xor state for next value
   * 3. convert to float [0, 1)
   *
   * @returns Uniformly distributed random float [0, 1)
   */
  next(): number {
    const s0 = this.state0;
    let s1 = this.state1;
    const result = (s0 + s1) & 0xffffffffffffffffn;

    // Update state for next value
    s1 ^= s0;
    this.state0 = (this.rotl(s0, 24n) ^ s1 ^ (s1 << 16n)) & 0xffffffffffffffffn;
    this.state1 = this.rotl(s1, 37n);

    // Convert 64-bit int to [0, 1) float
    // Use top 53 bits (double precision mantissa)
    return Number(result & 0x1fffffffffffffn) / 0x20000000000000;
  }

  /**
   * Rotate left (64-bit)
   */
  private rotl(x: bigint, k: bigint): bigint {
    return ((x << k) | (x >> (64n - k))) & 0xffffffffffffffffn;
  }
}

/**
 * Global PRNG instance
 *
 * Initialized once per completion via initializePRNG(seed).
 * All random() calls use this instance for determinism.
 */
let globalPRNG: Xoroshiro128Plus | null = null;

/**
 * Initialize global PRNG with seed
 *
 * @deprecated Use instance PRNG via SampleOptions instead:
 * ```typescript
 * const prng = new Xoroshiro128Plus(seed)
 * sampleWithStrategy(logits, { tokenHistory, params, prng })
 * ```
 *
 * Call once at start of each completion:
 * ```typescript
 * const seed = samplingParams?.seed ?? Date.now();
 * initializePRNG(seed);
 * ```
 *
 * @param seed 32-bit seed value
 */
export function initializePRNG(seed: number): void {
  globalPRNG = new Xoroshiro128Plus(seed);
}

/**
 * Generate next random number [0, 1)
 *
 * @deprecated Use instance PRNG from Xoroshiro128Plus class:
 * ```typescript
 * const prng = new Xoroshiro128Plus(seed)
 * const rand = prng.next()
 * ```
 *
 * Uses global PRNG instance initialized via initializePRNG().
 *
 * CRITICAL: Must call initializePRNG(seed) before first use.
 * Throws error if called before initialization.
 *
 * @returns Uniformly distributed random float [0, 1)
 * @throws Error if PRNG not initialized
 */
export function random(): number {
  if (!globalPRNG) {
    throw new Error('PRNG not initialized - call initializePRNG(seed) first');
  }
  return globalPRNG.next();
}

/**
 * Reset global PRNG (for testing)
 *
 * @deprecated For testing only. Create new PRNG instance instead.
 *
 * Resets to uninitialized state. Use with caution.
 */
export function resetPRNG(): void {
  globalPRNG = null;
}