Skip to main content

objectiveai_sdk/functions/check/example_inputs/
array.rs

1use rand::Rng;
2use rand::SeedableRng;
3use rand::rngs::StdRng;
4use rand::seq::SliceRandom;
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(|_| {
46                    Box::new(super::multi::generate(
47                        &item_schema,
48                        StdRng::seed_from_u64(rng.random::<u64>()),
49                    ))
50                })
51                .collect()
52        })
53        .collect();
54
55    let mut length_order: Vec<usize> = (0..lengths.len()).collect();
56    length_order.shuffle(&mut rng);
57
58    Generator {
59        per_length,
60        length_order,
61        pos: 0,
62        rng,
63    }
64}
65
66pub struct Generator<R: Rng> {
67    /// For each length option, a vec of generators (one per array index).
68    per_length: Vec<Vec<Box<super::multi::Generator>>>,
69    length_order: Vec<usize>,
70    pos: usize,
71    rng: R,
72}
73
74impl<R: Rng> Iterator for Generator<R> {
75    type Item = InputValue;
76    fn next(&mut self) -> Option<InputValue> {
77        // Round-robin through shuffled length order, reshuffle each cycle
78        let li = self.length_order[self.pos];
79        self.pos += 1;
80        if self.pos >= self.length_order.len() {
81            self.pos = 0;
82            self.length_order.shuffle(&mut self.rng);
83        }
84
85        // Build array: each index has its own generator for this length
86        let arr: Vec<InputValue> = self.per_length[li]
87            .iter_mut()
88            .map(|g| g.next().unwrap())
89            .collect();
90
91        Some(InputValue::Array(arr))
92    }
93}