Skip to main content

objectiveai_sdk/functions/check/example_inputs/
array.rs

1use rand::Rng;
2use rand::rngs::StdRng;
3use rand::seq::SliceRandom;
4use rand::SeedableRng;
5
6use crate::functions::expression::{ArrayInputSchema, InputValue};
7
8fn length_range(schema: &ArrayInputSchema) -> (usize, usize) {
9    match (schema.min_items, schema.max_items) {
10        (Some(min), Some(max)) => (min as usize, max as usize),
11        (Some(min), None) => (min as usize, min as usize + 10),
12        (None, Some(max)) => ((max as usize).saturating_sub(10), max as usize),
13        (None, None) => (0, 10),
14    }
15}
16
17/// Compute the distinct lengths: min, mid=(min+max)/2, max — deduplicated.
18fn distinct_lengths(schema: &ArrayInputSchema) -> Vec<usize> {
19    let (min, max) = length_range(schema);
20    let mid = (min + max) / 2;
21    let mut lengths = vec![min];
22    if mid != min {
23        lengths.push(mid);
24    }
25    if max != min && max != mid {
26        lengths.push(max);
27    }
28    lengths
29}
30
31pub fn permutations(schema: &ArrayInputSchema) -> usize {
32    let item_perms = super::optional::inner_permutations(schema.items.as_ref());
33    item_perms.saturating_mul(distinct_lengths(schema).len())
34}
35
36pub fn generate<R: Rng>(schema: &ArrayInputSchema, mut rng: R) -> Generator<R> {
37    let item_schema = schema.items.as_ref().clone();
38    let lengths = distinct_lengths(schema);
39
40    // For each distinct length, create independent generators per index.
41    let per_length: Vec<Vec<Box<super::multi::Generator>>> = lengths
42        .iter()
43        .map(|&len| {
44            (0..len)
45                .map(|_| Box::new(super::multi::generate(&item_schema, StdRng::seed_from_u64(rng.random::<u64>()))))
46                .collect()
47        })
48        .collect();
49
50    let mut length_order: Vec<usize> = (0..lengths.len()).collect();
51    length_order.shuffle(&mut rng);
52
53    Generator {
54        per_length,
55        length_order,
56        pos: 0,
57        rng,
58    }
59}
60
61pub struct Generator<R: Rng> {
62    /// For each length option, a vec of generators (one per array index).
63    per_length: Vec<Vec<Box<super::multi::Generator>>>,
64    length_order: Vec<usize>,
65    pos: usize,
66    rng: R,
67}
68
69impl<R: Rng> Iterator for Generator<R> {
70    type Item = InputValue;
71    fn next(&mut self) -> Option<InputValue> {
72        // Round-robin through shuffled length order, reshuffle each cycle
73        let li = self.length_order[self.pos];
74        self.pos += 1;
75        if self.pos >= self.length_order.len() {
76            self.pos = 0;
77            self.length_order.shuffle(&mut self.rng);
78        }
79
80        // Build array: each index has its own generator for this length
81        let arr: Vec<InputValue> = self.per_length[li]
82            .iter_mut()
83            .map(|g| g.next().unwrap())
84            .collect();
85
86        Some(InputValue::Array(arr))
87    }
88}