Skip to main content

objectiveai_sdk/functions/check/example_inputs/
optional.rs

1use rand::Rng;
2use rand::SeedableRng;
3use rand::rngs::StdRng;
4use rand::seq::SliceRandom;
5
6use crate::functions::expression::{InputSchema, InputValue};
7
8pub fn permutations(schema: &InputSchema) -> usize {
9    inner_permutations(schema) * 2
10}
11
12pub fn inner_permutations(schema: &InputSchema) -> usize {
13    match schema {
14        InputSchema::Boolean(s) => super::boolean::permutations(s),
15        InputSchema::String(s) => super::string::permutations(s),
16        InputSchema::Integer(s) => super::integer::permutations(s),
17        InputSchema::Number(s) => super::number::permutations(s),
18        InputSchema::Image(s) => super::image::permutations(s),
19        InputSchema::Audio(s) => super::audio::permutations(s),
20        InputSchema::Video(s) => super::video::permutations(s),
21        InputSchema::File(s) => super::file::permutations(s),
22        InputSchema::Object(s) => super::object::permutations(s),
23        InputSchema::Array(s) => super::array::permutations(s),
24        InputSchema::AnyOf(s) => super::any_of::permutations(s),
25    }
26}
27
28pub fn generate(schema: &InputSchema, mut rng: StdRng) -> Generator {
29    let inner_count = inner_permutations(schema);
30    // 0..inner_count = present, inner_count..inner_count*2 = absent
31    let total = inner_count * 2;
32    let mut indices: Vec<usize> = (0..total).collect();
33    indices.shuffle(&mut rng);
34
35    let inner = super::multi::generate(
36        schema,
37        StdRng::seed_from_u64(rng.random::<u64>()),
38    );
39
40    Generator {
41        inner,
42        inner_count,
43        indices,
44        pos: 0,
45        rng,
46    }
47}
48
49pub struct Generator {
50    inner: super::multi::Generator,
51    inner_count: usize,
52    indices: Vec<usize>,
53    pos: usize,
54    rng: StdRng,
55}
56
57impl Iterator for Generator {
58    type Item = Option<InputValue>;
59    fn next(&mut self) -> Option<Option<InputValue>> {
60        if self.indices.is_empty() {
61            return Some(None);
62        }
63        if self.pos >= self.indices.len() {
64            self.indices.shuffle(&mut self.rng);
65            self.pos = 0;
66        }
67        let index = self.indices[self.pos];
68        self.pos += 1;
69        if index < self.inner_count {
70            Some(self.inner.next())
71        } else {
72            Some(None)
73        }
74    }
75}