GenerateWeightedCount Class — drizzle-orm Architecture

Source Code

drizzle-seed/src/services/Generators.ts lines 132–168

export class GenerateWeightedCount extends AbstractGenerator<{}> {
	static override readonly entityKind: string = 'GenerateWeightedCount';

	private state: {
		rng: prand.RandomGenerator;
		weightedIndices: number[];
		weightedCount: { weight: number; count: number | number[] }[];
	} | undefined;

	override init({ seed, count }: { count: { weight: number; count: number | number[] }[]; seed: number }) {
		const rng = prand.xoroshiro128plus(seed);
		const weightedIndices = getWeightedIndices(count.map((val) => val.weight));
		this.state = { rng, weightedIndices, weightedCount: count };
	}

	generate() {
		if (this.state === undefined) {
			throw new Error('state is not defined.');
		}

		// logic for this generator
		let idx: number;
		const weightedCount = this.state.weightedCount;

		[idx, this.state.rng] = prand.uniformIntDistribution(0, this.state.weightedIndices.length - 1, this.state.rng);
		const objIdx = this.state.weightedIndices[idx] as number;

		if (typeof weightedCount[objIdx]!.count === 'number') {
			return weightedCount[objIdx]!.count as number;
		} else {
			// typeof weightedCount[objIdx]!.count === 'object' // number[]
			const possCounts = weightedCount[objIdx]!.count as number[];
			[idx, this.state.rng] = prand.uniformIntDistribution(0, possCounts.length - 1, this.state.rng);
			return possCounts[idx]!;
		}
	}
}