Add generator for strings

3 files changed, 177 insertions, 101 deletions

diff --git a/pbt/ts/src/gen.test.ts b/pbt/ts/src/gen.test.ts
index 75c41c8..6e891e4 100644
--- a/pbt/ts/src/gen.test.ts
+++ b/pbt/ts/src/gen.test.ts
@@ -2,6 +2,7 @@ import Prando from 'prando';
 import { none } from 'fp-ts/lib/Option.js';
 import * as G from './gen.js';
 import { test, expect } from 'vitest';
+import { chiSquare, computeChiSquaredPValue, computeCramersV } from './stats.js';
 
 test('can generate and shrink bytes', async () => {
   const prng = new Prando(Math.random() * 1000);
@@ -45,10 +46,10 @@ test('can generate arbitrary integer within bounds', () => {
 test('generated integers are uniformly distributed', () => {
   const seed = Math.random() * 1000;
   const lb = 0;
-  const ub = 65536;
+  const ub = 100000;
   const gen = G.genIntWithin(lb, ub, seed);
   const range = ub - lb;
-  const buckets = 128;
+  const buckets = 100;
   const numSamples = 100000;
   // indexed by lower bound
   const counts: Record<number, number> = {};
@@ -70,102 +71,31 @@ test('generated integers are uniformly distributed', () => {
 
   // compute χ² statistic
   const expectedCount = numSamples / buckets;
-  const chiSquared = Object.values(counts).reduce((acc, count) => {
-    const diff = count - expectedCount;
-    return acc + (diff * diff) / expectedCount;
-  });
+  const chiSquared = chiSquare(Object.values(counts), new Array(buckets).fill(expectedCount));
   // compute p-value and Cramér's V.
-  // In our case the p-value will
-  // systematically be negligible and therefore will reject the null
-  // hypothesis, eg. that the frequency distribution of numbers we
-  // generate is truly uniform. However, Cramér's V is a better
-  // measure of effect size and will be small if the distribution is
-  // uniform.
   const cramerv = computeCramersV(chiSquared, numSamples, buckets);
   const pValue = computeChiSquaredPValue(chiSquared, buckets);
   console.log(`Chi-squared: ${chiSquared}, p-value: ${pValue}, Cramer V: ${cramerv}`);
-  expect(cramerv).toBeLessThan(0.01);
+  expect(cramerv).toBeLessThan(0.015);
+  expect(pValue).toBeGreaterThan(0.05);
 });
 
-
-/**
- * Computes the right-tailed p-value for a Chi-Squared distribution
- * using the Wilson-Hilferty approximation.
- *
- * @param chiSquareValue The calculated empirical chi-square statistic (e.g., 1190.372)
- * @param numBuckets The total number of bins/categories in your test (e.g., 100)
- * @returns The p-value between 0.0 and 1.0
- */
-function computeChiSquaredPValue(chiSquareValue: number, numBuckets: number): number {
-  // Degrees of Freedom = Number of buckets - 1
-  const df = numBuckets - 1;
-
-  if (df <= 0) {
-    throw new Error("Number of buckets must be greater than 1.");
-  }
-  if (chiSquareValue < 0) {
-    return 1.0;
+test('can generate uniformly distributed arbitrary strings shorter than 256 characters', () => {
+  const seed = Math.random() * (1 << 31);
+  const gen = G.genASCIIString(seed);
+  const buckets = 16;
+  const numSamples = (1 << 16);
+  const lengths = new Array(buckets).fill(0);
+  for (let i = 0; i < numSamples; i++) {
+    const value = gen.generate();
+    lengths[value.length % buckets]++;
+    expect(value.length).toBeLessThan(256);
   }
 
-  // 1. Compute the Wilson-Hilferty Normal Approximation Z-score
-  const term1 = Math.pow(chiSquareValue / df, 1 / 3);
-  const term2 = 1 - 2 / (9 * df);
-  const denominator = Math.sqrt(2 / (9 * df));
-
-  const zScore = (term1 - term2) / denominator;
-
-  // 2. Compute right-tailed probability from Z-score using standard Normal CDF
-  return standardNormalRightTail(zScore);
-}
-
-/**
- * Standard Normal Right-Tail Probability Q(z) = P(Z > z)
- * Uses a highly precise numerical approximation for the error function (erf).
- */
-function standardNormalRightTail(z: number): number {
-  // Standard Normal Cumulative Distribution Function (CDF)
-  const cdf = 0.5 * (1 + errorFunction(z / Math.sqrt(2)));
-  // Right tail is 1 - CDF
-  return 1 - cdf;
-}
-
-/**
- * Abramowitz and Stegun approximation for the Error Function (erf),
- * accurate to 1.5 x 10^-7
- *
- * See https://en.wikipedia.org/wiki/Error_function#Bounds_and_numerical_approximations
- */
-function errorFunction(x: number): number {
-  const sign = x < 0 ? -1 : 1;
-  const absX = Math.abs(x);
-
-  // Constants for the numerical approximation
-  const p = 0.3275911;
-  const a1 = 0.254829592;
-  const a2 = -0.284496736;
-  const a3 = 1.421413741;
-  const a4 = -1.453152027;
-  const a5 = 1.061405429;
-
-  const t = 1.0 / (1.0 + p * absX);
-  const y = 1.0 - (((((a5 * t + a4) * t) + a3) * t + a2) * t + a1) * t * Math.exp(-absX * absX);
-
-  return sign * y;
-}
-
-/**
- * Calculates Cramér's V to determine the practical effect size (strength of bias).
- * See https://en.wikipedia.org/wiki/Cramér%27s_V
- *
- * @param chiSquare The empirical chi-square value (e.g., 1190.372)
- * @param totalSamples The total number of data points (e.g., 100000)
- * @param numBuckets The total number of bins/buckets (e.g., 100)
- * @returns A value between 0.0 (no bias) and 1.0 (extreme bias)
- */
-function computeCramersV(chiSquare: number, totalSamples: number, numBuckets: number): number {
-  const df = numBuckets - 1;
-  if (df <= 0 || totalSamples <= 0) return 0;
-
-  // Core Formula: sqrt( chi2 / (N * df) )
-  return Math.sqrt(chiSquare / (totalSamples * df));
-}
+  console.log(lengths);
+  const chiSquared = chiSquare(lengths, new Array(buckets).fill(numSamples / buckets));
+  const cramerv = computeCramersV(chiSquared, numSamples, buckets);
+  const pValue = computeChiSquaredPValue(chiSquared, buckets);
+  console.log(`Chi-squared: ${chiSquared}, p-value: ${pValue}, Cramer V: ${cramerv}`);
+  expect(cramerv).toBeLessThan(0.015);
+});
diff --git a/pbt/ts/src/gen.ts b/pbt/ts/src/gen.ts
index c27cfec..7d8c6f1 100644
--- a/pbt/ts/src/gen.ts
+++ b/pbt/ts/src/gen.ts
@@ -27,7 +27,7 @@ export class BytesGen implements GenBytes {
   }
 
   generate(n: number): number[] {
-    if (n <= 0) {
+    if (n < 0) {
       throw new Error('n must be greater than 0');
     }
     return this.generateBytes(n);
@@ -42,6 +42,13 @@ export class BytesGen implements GenBytes {
     return gens;
   }
 
+  reject() {
+    if (this.bytes.length === 0) {
+      throw new Error('No more bytes to shrink');
+    }
+    this.bytes.pop();
+  }
+
   replay(): GenBytes {
     return new ReplayBytes(this.bytes);
   }
@@ -106,12 +113,27 @@ class GenIntWithin implements Gen<number> {
   generate(): number {
     const range = this.max - this.min;
     const numBytes = Math.ceil(Math.log2(range) / 8);
-    const bytes = this.bytes.generate(numBytes);
-    let value = 0;
-    for (let i = 0; i < bytes.length; i++) {
-      value = (value << 8) | bytes[i];
+    const maxBytes = Math.pow(2, numBytes * 8);
+    const upperBound = maxBytes - (maxBytes % range);
+    let bytes = this.bytes.generate(numBytes);
+    let rejection = 0;
+    while (true) {
+      let value = 0;
+      for (let i = 0; i < bytes.length; i++) {
+        value = (value << 8) | bytes[i];
+      }
+
+      if (value >= upperBound) {
+        this.bytes.reject();
+        bytes = this.bytes.generate(numBytes);
+        continue;
+      }
+
+      if (rejection > 0) {
+        console.log(`Rejection sampling: rejected ${rejection} values before accepting ${value}`);
+      }
+      return this.min + (value % range);
     }
-    return this.min + (value % range);
   }
 
   generateN(n: number): number[] {
@@ -125,6 +147,30 @@ class GenIntWithin implements Gen<number> {
 
 }
 
-export const genInt = (seed : number | undefined = undefined) => new GenInt(seed);
+class GenASCIIString implements Gen<string> {
+  private bytes: BytesGen;
+
+  constructor(seed: number | undefined = undefined) {
+    this.bytes = new BytesGen(new Prando(seed ? seed : Math.random() * 1000));
+  }
+
+  generate(): string {
+    const length = this.bytes.generate(1)[0]; // limit to 20 characters
+    const bytes = this.bytes.generate(length);
+    return String.fromCharCode(...bytes.map(b => b % 128)); // limit to ASCII
+  }
+
+  generateN(n: number): string[] {
+    const result = [];
+    for (let i = 0; i < n; i++) {
+      result.push(this.generate());
+    }
+    return result;
+  }
+}
+
+export const genInt = (seed: number | undefined = undefined) => new GenInt(seed);
+
+export const genIntWithin = (min: number, max: number, seed: number | undefined = undefined) => new GenIntWithin(min, max, seed);
 
-export const genIntWithin = (min: number, max: number, seed : number | undefined = undefined) => new GenIntWithin(min, max, seed);
+export const genASCIIString = (seed: number | undefined = undefined) => new GenASCIIString(seed);
diff --git a/pbt/ts/src/stats.ts b/pbt/ts/src/stats.ts
new file mode 100644
index 0000000..8c50856
--- /dev/null
+++ b/pbt/ts/src/stats.ts
@@ -0,0 +1,100 @@
+/**
+ * Computes the right-tailed p-value for a Chi-Squared distribution
+ * using the Wilson-Hilferty approximation.
+ *
+ * @param chiSquareValue The calculated empirical chi-square statistic (e.g., 1190.372)
+ * @param numBuckets The total number of bins/categories in your test (e.g., 100)
+ * @returns The p-value between 0.0 and 1.0
+ */
+export function computeChiSquaredPValue(chiSquareValue: number, numBuckets: number): number {
+  // Degrees of Freedom = Number of buckets - 1
+  const df = numBuckets - 1;
+
+  if (df <= 0) {
+    throw new Error("Number of buckets must be greater than 1.");
+  }
+  if (chiSquareValue < 0) {
+    return 1.0;
+  }
+
+  // 1. Compute the Wilson-Hilferty Normal Approximation Z-score
+  const term1 = Math.pow(chiSquareValue / df, 1 / 3);
+  const term2 = 1 - 2 / (9 * df);
+  const denominator = Math.sqrt(2 / (9 * df));
+
+  const zScore = (term1 - term2) / denominator;
+
+  // 2. Compute right-tailed probability from Z-score using standard Normal CDF
+  return standardNormalRightTail(zScore);
+}
+
+/**
+ * Standard Normal Right-Tail Probability Q(z) = P(Z > z)
+ * Uses a highly precise numerical approximation for the error function (erf).
+ */
+function standardNormalRightTail(z: number): number {
+  // Standard Normal Cumulative Distribution Function (CDF)
+  const cdf = 0.5 * (1 + errorFunction(z / Math.sqrt(2)));
+  // Right tail is 1 - CDF
+  return 1 - cdf;
+}
+
+/**
+ * Abramowitz and Stegun approximation for the Error Function (erf),
+ * accurate to 1.5 x 10^-7
+ *
+ * See https://en.wikipedia.org/wiki/Error_function#Bounds_and_numerical_approximations
+ */
+function errorFunction(x: number): number {
+  const sign = x < 0 ? -1 : 1;
+  const absX = Math.abs(x);
+
+  // Constants for the numerical approximation
+  const p = 0.3275911;
+  const a1 = 0.254829592;
+  const a2 = -0.284496736;
+  const a3 = 1.421413741;
+  const a4 = -1.453152027;
+  const a5 = 1.061405429;
+
+  const t = 1.0 / (1.0 + p * absX);
+  const y = 1.0 - (((((a5 * t + a4) * t) + a3) * t + a2) * t + a1) * t * Math.exp(-absX * absX);
+
+  return sign * y;
+}
+
+/**
+ * Calculates Cramér's V to determine the practical effect size (strength of bias).
+ * See https://en.wikipedia.org/wiki/Cramér%27s_V
+ *
+ * @param chiSquare The empirical chi-square value (e.g., 1190.372)
+ * @param totalSamples The total number of data points (e.g., 100000)
+ * @param numBuckets The total number of bins/buckets (e.g., 100)
+ * @returns A value between 0.0 (no bias) and 1.0 (extreme bias)
+ */
+export function computeCramersV(chiSquare: number, totalSamples: number, numBuckets: number): number {
+  const df = numBuckets - 1;
+  if (df <= 0 || totalSamples <= 0) return 0;
+
+  // Core Formula: sqrt( chi2 / (N * df) )
+  return Math.sqrt(chiSquare / (totalSamples * df));
+}
+
+
+export function chiSquare (samples : number [], expected : number []) : number {
+  if (samples.length !== expected.length) {
+    throw new Error("Samples and expected arrays must have the same length.");
+  }
+
+  let chiSquareValue = 0;
+  for (let i = 0; i < samples.length; i++) {
+    const observed = samples[i];
+    const exp = expected[i];
+    if (exp === 0) {
+      throw new Error("Expected frequency cannot be zero.");
+    }
+    chiSquareValue += Math.pow(observed - exp, 2) / exp;
+  }
+
+  return chiSquareValue;
+}