objectiveai_sdk/functions/check/example_inputs/
any_of.rs1use rand::Rng;
2use rand::rngs::StdRng;
3use rand::seq::SliceRandom;
4use rand::SeedableRng;
5
6use crate::functions::expression::{AnyOfInputSchema, InputValue};
7
8fn max_inner_permutations(schema: &AnyOfInputSchema) -> usize {
9 schema
10 .any_of
11 .iter()
12 .map(|v| super::optional::inner_permutations(v))
13 .max()
14 .unwrap_or(0)
15}
16
17pub fn permutations(schema: &AnyOfInputSchema) -> usize {
18 schema
19 .any_of
20 .len()
21 .saturating_mul(max_inner_permutations(schema))
22}
23
24pub fn generate<R: Rng>(schema: &AnyOfInputSchema, mut rng: R) -> Generator<R> {
25 let variant_count = schema.any_of.len();
26 let max_inner = max_inner_permutations(schema);
27
28 let generators: Vec<super::multi::Generator> = schema
29 .any_of
30 .iter()
31 .map(|v| super::multi::generate(v, StdRng::seed_from_u64(rng.random::<u64>())))
32 .collect();
33
34 let mut order: Vec<usize> = (0..variant_count)
36 .flat_map(|i| std::iter::repeat(i).take(max_inner))
37 .collect();
38 order.shuffle(&mut rng);
39
40 Generator {
41 generators,
42 order,
43 pos: 0,
44 rng,
45 }
46}
47
48pub struct Generator<R: Rng> {
49 generators: Vec<super::multi::Generator>,
50 order: Vec<usize>,
51 pos: usize,
52 rng: R,
53}
54
55impl<R: Rng> Iterator for Generator<R> {
56 type Item = InputValue;
57 fn next(&mut self) -> Option<InputValue> {
58 if self.order.is_empty() {
59 return None;
60 }
61 if self.pos >= self.order.len() {
62 self.order.shuffle(&mut self.rng);
63 self.pos = 0;
64 }
65 let variant_idx = self.order[self.pos];
66 self.pos += 1;
67 self.generators[variant_idx].next()
68 }
69}