temptress-bot/src/util/weightedList.ts

import { PRNG } from 'seedrandom'

// Weighted items do not support as values:
// * Individual weighted items (what does it mean for a weighted item to have another weight inside?)
// * Arrays (we don't know the difference between an array that is an item and an array that is)
// * Items with the length or wt property as a number (that's what we use to detect an array or weighted item)
export type WeightableItem =
    | {
          readonly wt?: undefined | null | string | boolean | bigint | symbol | object | never
          readonly length?: undefined | null | string | boolean | bigint | symbol | object | never
      }
    | string
    | boolean
    | number
    | bigint
    | null
    | undefined
    | symbol
export interface WeightedItem<ItemT extends WeightableItem> {
    readonly wt: number
    readonly v: ItemT | WeightedList<ItemT>
}
type WeightedListItem<ItemT extends WeightableItem> = WeightedItem<ItemT> | ItemT | WeightedList<ItemT>
export type WeightedList<ItemT extends WeightableItem> = readonly WeightedListItem<ItemT>[]

function isWeightedItem<ItemT extends WeightableItem>(x: WeightedListItem<ItemT>): x is WeightedItem<ItemT> {
    return typeof x === 'object' && x !== null && 'wt' in x && typeof x['wt'] === 'number'
}

function getItemValue<ItemT extends WeightableItem>(x: WeightedListItem<ItemT>): WeightedList<ItemT> | ItemT {
    if (isWeightedItem(x)) {
        return x.v
    } else {
        return x
    }
}

type TotalWeightedItem<ItemT extends WeightableItem> = {
    readonly minWt: number
    readonly maxWt: number
    readonly v: WeightedList<ItemT> | ItemT
}

type TotalWeightedList<ItemT extends WeightableItem> = readonly TotalWeightedItem<ItemT>[]

function isWeightedList<ItemT extends WeightableItem>(x: ItemT | WeightedList<ItemT>): x is WeightedList<ItemT> {
    return typeof x === 'object' && x !== null && 'length' in x && typeof x['length'] === 'number'
}

export function shuffleWeightedList<ItemT extends WeightableItem>(
    x: WeightedList<ItemT> | ItemT,
    emptyValue: ItemT,
    rng: PRNG,
): ItemT {
    if (!isWeightedList(x)) {
        return x
    }
    if (x.length === 0) {
        return emptyValue
    }
    if (x.length === 1) {
        return shuffleWeightedList(getItemValue(x[0]), emptyValue, rng)
    }
    const totals = x.reduce<TotalWeightedList<ItemT>>((accumulated, next) => {
        const minWeight = accumulated.length === 0 ? 0 : accumulated[accumulated.length - 1].maxWt
        const newWeight = isWeightedItem(next) ? next.wt : 1
        if (minWeight + newWeight === minWeight) {
            // weightless items get ignored
            // this includes items whose weights are so small relative to the list that they are effectively weightless
            return accumulated
        }
        return [...accumulated, { minWt: minWeight, maxWt: minWeight + newWeight, v: getItemValue(next) }]
    }, [])
    if (totals.length === 0) {
        return emptyValue
    }
    if (totals.length === 1) {
        return shuffleWeightedList(totals[0].v, emptyValue, rng)
    }
    // Here we have a weighted list with no items with any weight.
    if (totals[totals.length - 1].maxWt <= 0) {
        return emptyValue
    }
    // Now we pick a value out of the stacked weights and binary search it up
    const target = rng.double() * totals[totals.length - 1].maxWt
    for (let start = 0, end = totals.length; ; ) {
        const nextIndex = start + Math.floor((end - start) / 2)
        if (target >= totals[nextIndex].maxWt) {
            start = nextIndex + 1
        } else if (target < totals[nextIndex].minWt) {
            end = nextIndex
        } else {
            return shuffleWeightedList(totals[nextIndex].v, emptyValue, rng)
        }
    }
}