1use std::rc::Rc;
2
3use rand::seq::SliceRandom;
4
5use crate::{
6 ast::{Amount, Choose, ChooseOptions, Entry, Query},
7 Pcg,
8};
9
10pub 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 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 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}