rng_query/
eval.rs

1use std::rc::Rc;
2
3use rand::seq::SliceRandom;
4
5use crate::{
6    ast::{Amount, Choose, ChooseOptions, Entry, Query},
7    Pcg,
8};
9
10/// A sample from a selected entry
11///
12/// This is an opaque type, hidden intentionally. It only expose the [`Display`]
13/// implementation to access it. The
14/// [`Display`] [alternate modifier](std::fmt#sign0) will only print the sampled
15/// value and not the whole representation.
16///
17/// [`Display`]: std::fmt::Display
18pub struct Sample(SampleData);
19
20enum SampleData {
21    Text(Rc<str>),
22    Expr(Box<dyn std::fmt::Display>),
23}
24
25impl Sample {
26    pub(crate) fn text(data: Rc<str>) -> Self {
27        Self(SampleData::Text(data))
28    }
29    pub(crate) fn expr(data: Box<dyn std::fmt::Display>) -> Self {
30        Self(SampleData::Expr(data))
31    }
32}
33
34impl std::fmt::Display for Sample {
35    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
36        match &self.0 {
37            SampleData::Text(t) => t.fmt(f),
38            SampleData::Expr(e) => e.fmt(f),
39        }
40    }
41}
42
43pub(crate) enum EvalRes {
44    Emtpy,
45    Single(Sample),
46    Many(Vec<Sample>),
47}
48
49impl From<Sample> for EvalRes {
50    fn from(value: Sample) -> Self {
51        Self::Single(value)
52    }
53}
54
55impl From<Vec<Sample>> for EvalRes {
56    fn from(value: Vec<Sample>) -> Self {
57        Self::Many(value)
58    }
59}
60
61pub(crate) trait Eval {
62    fn eval(&self, rng: &mut Pcg) -> EvalRes;
63}
64
65impl<T, R> Eval for T
66where
67    T: Fn(&mut Pcg) -> R,
68    R: Into<EvalRes>,
69{
70    fn eval(&self, rng: &mut Pcg) -> EvalRes {
71        (self)(rng).into()
72    }
73}
74
75impl Eval for Query {
76    fn eval(&self, rng: &mut Pcg) -> EvalRes {
77        self.root.eval(rng)
78    }
79}
80
81impl Eval for Choose {
82    fn eval(&self, rng: &mut Pcg) -> EvalRes {
83        let Self { entries, options } = self;
84
85        let selected = select(rng, entries, options);
86
87        if selected.is_empty() {
88            return EvalRes::Emtpy;
89        }
90
91        let mut v = Vec::with_capacity(selected.len());
92        for (_, entry) in selected {
93            match entry.eval(rng) {
94                EvalRes::Emtpy => {}
95                EvalRes::Single(s) => v.push(s),
96                EvalRes::Many(mut vv) => v.append(&mut vv),
97            }
98        }
99        EvalRes::Many(v)
100    }
101}
102
103impl Eval for Entry {
104    fn eval(&self, rng: &mut Pcg) -> EvalRes {
105        match self {
106            Entry::Text(t) => Sample::text(t.clone()).into(),
107            Entry::Expr(e) => e.eval(rng),
108        }
109    }
110}
111
112fn select(
113    rng: &mut Pcg,
114    entries: &[(usize, Entry)],
115    options: &ChooseOptions,
116) -> Vec<(usize, Entry)> {
117    if entries.is_empty() {
118        return vec![];
119    }
120
121    let n = match options.amount {
122        Amount::All => entries.len(),
123        Amount::N(n) => n as usize,
124    };
125
126    // optimization for all
127    if !options.repeating && n >= entries.len() {
128        let mut entries = entries.to_vec();
129        if !options.keep_order {
130            entries.shuffle(rng);
131        }
132        return entries;
133    }
134
135    // general case
136    let mut selected = if options.repeating {
137        let mut selected = Vec::with_capacity(n);
138        for _ in 0..n {
139            let entry = entries.choose(rng).unwrap();
140            selected.push(entry.clone());
141        }
142        selected
143    } else {
144        entries.choose_multiple(rng, n).cloned().collect()
145    };
146
147    if options.keep_order {
148        selected.sort_unstable_by_key(|e| e.0);
149    }
150    selected
151}