Skip to main content

objectiveai_sdk/functions/check/example_inputs/
any_of.rs

1use 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    // Each variant appears max_inner times per cycle
35    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}