turbopack_bench/util/
module_picker.rs

1use std::path::PathBuf;
2
3use rand::{SeedableRng, rngs::StdRng, seq::IndexedRandom};
4use rustc_hash::FxHashMap;
5
6/// Picks modules at random, but with a fixed seed so runs are somewhat
7/// reproducible.
8///
9/// This must be initialized outside of `bench_with_input` so we don't repeat
10/// the same sequence in different samples.
11pub struct ModulePicker {
12    depths: Vec<usize>,
13    modules_by_depth: FxHashMap<usize, Vec<PathBuf>>,
14    rng: parking_lot::Mutex<StdRng>,
15}
16
17impl ModulePicker {
18    /// Creates a new module picker.
19    pub fn new(mut modules: Vec<(PathBuf, usize)>) -> Self {
20        let rng = StdRng::seed_from_u64(42);
21
22        // Ensure the module order is deterministic.
23        modules.sort();
24
25        let mut modules_by_depth: FxHashMap<_, Vec<_>> = FxHashMap::default();
26        for (module, depth) in modules {
27            modules_by_depth.entry(depth).or_default().push(module);
28        }
29        let mut depths: Vec<_> = modules_by_depth.keys().copied().collect();
30        // Ensure the depth order is deterministic.
31        depths.sort();
32
33        Self {
34            depths,
35            modules_by_depth,
36            rng: parking_lot::Mutex::new(rng),
37        }
38    }
39
40    /// Picks a random module with a uniform distribution over all depths.
41    pub fn pick(&self) -> &PathBuf {
42        let mut rng = self.rng.lock();
43        // Sample from all depths uniformly.
44        let depth = self.depths.choose(&mut *rng).unwrap();
45        self.modules_by_depth[depth].choose(&mut *rng).unwrap()
46    }
47}