feat(oneOfWeighted): Support weighted choosing

index.d.ts

@@ -9,6 +9,7 @@ export type Input = JSONSerializable
 export type Range = number | [number, number]
 
 export type Maker<V = unknown> = ((input: Input) => V) | V
+export type WeightedMaker<V = unknown> = [number, Maker<V>]
 
 export function hash(input: Input): number
 export function bool(input: Input): boolean
@@ -115,6 +116,17 @@ declare const oneOf: OneOf
 
 export { oneOf }
 
+export interface OneOfWeighted {
+  <M extends WeightedMaker>(samples: M[]): (
+    input: Input
+  ) => WeightedMakerResult<M>
+  <M extends WeightedMaker>(input: Input, samples: M[]): WeightedMakerResult<M>
+}
+
+declare const oneOfWeighted: OneOfWeighted
+
+export { oneOfWeighted }
+
 export interface SomeOf {
   <M extends Maker>(range: Range, samples: M[]): (
     input: Input
@@ -166,6 +178,7 @@ declare const tuple: Tuple
 export { tuple }
 
 type MakerResult<M> = M extends Maker<infer R> ? R : never
+type WeightedMakerResult<M> = M extends WeightedMaker<infer R> ? R : never
 
 export type TupleReturnType<Makers extends AnyMakers> = Makers extends Makers1<
   infer V1

index.js

@@ -5,6 +5,7 @@ exports.someOf = require('./someOf')
 exports.tuple = require('./tuple')
 exports.times = require('./times')
 exports.join = require('./join')
+exports.oneOfWeighted = require('./oneOfWeighted')
 
 exports.int = require('./int')
 exports.float = require('./float')

index.mjs

@@ -1,10 +1,11 @@
 export { default as hash } from './has'
 
-export { default as oneOf } from './oneOf'
 export { default as someOf } from './someOf'
 export { default as tuple } from './tuple'
 export { default as times } from './times'
 export { default as join } from './join'
+export { default as oneOf } from './oneOf'
+export { default as oneOfWeighted } from './oneOfWeighted'
 
 export { default as int } from './int'
 export { default as float } from './float'

index.test-d.ts

@@ -7,30 +7,49 @@ import {
   float,
   word,
   tuple,
-  oneOf,
   someOf,
   times,
-  join
+  join,
+  oneOf,
+  oneOfWeighted
 } from '.'
 
 // ## function items
 expectType<[number, string]>(tuple(null, [int, word]))
 expectType<string>(join(null, ' ', [int, word]))
 expectType<number | string>(oneOf(null, [int, word]))
+expectType<number | string>(
+  oneOfWeighted(null, [
+    [0.6, int],
+    [0.4, word]
+  ])
+)
 expectType<(number | string)[]>(someOf(null, [2, 3], [int, word]))
 expectType<number[]>(times(null, [2, 3], int))
 
 // ## constant items
 expectType<[number, string]>(tuple(null, [2, '!']))
 expectType<string>(join(null, ' ', [2, '!']))
 expectType<number | string>(oneOf(null, [2, '!']))
+expectType<number | string>(
+  oneOfWeighted(null, [
+    [0.6, 2],
+    [0.4, '!']
+  ])
+)
 expectType<(number | string)[]>(someOf(null, [3, 3], [2, '!']))
 expectType<number[]>(times(null, [2, 3], 23))
 
 // ## currying
 expectType<(input: Input) => [number]>(tuple([int]))
 expectType<(input: Input) => string>(join(' ', [int, word]))
 expectType<(input: Input) => number>(oneOf([2, 3]))
+expectType<(input: Input) => number>(
+  oneOfWeighted([
+    [0.6, 2],
+    [0.6, 3]
+  ])
+)
 expectType<(input: Input) => string[]>(someOf([3, 3], ['a', 'b']))
 expectType<(input: Input) => string[]>(times([2, 3], word))
 

oneOf.js

@@ -1,13 +1,13 @@
 var hash = require('./hash')
+var resolve = require('./utils/resolve')
 
 function oneOf(a, b) {
   return b != null ? oneOfMain(a, b) : oneOfCurried(a)
 }
 
 function oneOfMain(input, samples) {
   var id = hash(input)
-  var result = samples[id % samples.length]
-  return typeof result === 'function' ? result(hash([id, 'oneOf'])) : result
+  return resolve(id, samples[id % samples.length])
 }
 
 function oneOfCurried(samples) {

oneOfWeighted.js

@@ -0,0 +1,115 @@
+var hash = require('./hash')
+var flip = require('./utils/flip')
+var resolve = require('./utils/resolve')
+
+var EPS = 0.0001
+
+function oneOfWeighted(a, b) {
+  return b != null ? oneOfWeightedMain(a, b) : oneOfWeightedCurried(a)
+}
+
+function oneOfWeightedMain(input, samples) {
+  samples = parseSamples(samples)
+  var id = hash(input)
+  var i = -1
+  var n = samples.length
+  var pRemaining = 1
+  var sample
+  var p
+
+  while (++i < n) {
+    sample = samples[i]
+    p = sample[0] / pRemaining
+    pRemaining -= p
+
+    if (flip(id, p)) {
+      return resolve(id, sample[1])
+    }
+  }
+
+  throw new Error(
+    'Unexpectedly reached end of oneOfWeighted() unresolved. If you see this please file a bug.'
+  )
+}
+
+function oneOfWeightedCurried(samples) {
+  samples = parseSamples(samples)
+
+  return function oneOfWeightedCurriedFn(input) {
+    return oneOfWeighted(input, samples)
+  }
+}
+
+function parseSamples(samples) {
+  if (samples.__parsed) {
+    return samples
+  }
+
+  var samplesLen = samples.length
+  var assignedPs = getAssignedPs(samples)
+  var sumAssignedPs = assignedPs.reduce(add, 0)
+  var assignedPsLen = assignedPs.length
+
+  if (!samplesLen) {
+    throw new Error('Empty samples given to oneOfWeighted')
+  }
+
+  if (sumAssignedPs > 1 + EPS) {
+    throw new Error(
+      'Assigned probabilities add up more than 1: ' + JSON.stringify(assignedPs)
+    )
+  } else if (samplesLen === assignedPsLen && sumAssignedPs < 1 - EPS) {
+    throw new Error(
+      'All items were assigned probabilities, yet the probabilities add up to less than 1: ' +
+        JSON.stringify(assignedPs)
+    )
+  }
+
+  var defaultP = (1 - sumAssignedPs) / (samplesLen - assignedPsLen)
+  return ensurePs(samples, defaultP)
+}
+
+function ensurePs(samples, defaultP) {
+  var result = []
+  var i = -1
+  var n = samples.length
+  var sample
+  var p
+
+  while (++i < n) {
+    sample = samples[i]
+    p = sample[0]
+
+    if (typeof p !== 'number') {
+      p = defaultP
+    }
+
+    result.push([p, sample[1]])
+  }
+
+  result.__parsed = true
+  return result
+}
+
+function getAssignedPs(samples) {
+  var i = -1
+  var n = samples.length
+  var result = []
+  var p
+
+  while (++i < n) {
+    p = samples[i][0]
+
+    if (typeof p === 'number') {
+      result.push(p)
+    }
+  }
+
+  return result
+}
+
+function add(a, b) {
+  return a + b
+}
+
+module.exports = oneOfWeighted

someOf.js

@@ -21,7 +21,7 @@ function someOfMain(input, range, samples) {
     chosenIndex = id % remainingLen
     chosen = remaining[chosenIndex]
     remaining.splice(chosenIndex, 1)
-    results.push(resolve(chosen))
+    results.push(resolve(id, chosen))
   }
 
   return results