Skip to main content

objectiveai_sdk/functions/check/example_inputs/
integer.rs

1use rand::Rng;
2use rand::seq::SliceRandom;
3
4use crate::functions::expression::{InputValue, IntegerInputSchema};
5
6pub fn permutations(schema: &IntegerInputSchema) -> usize {
7    let min = schema.minimum;
8    let max = schema.maximum;
9    let mut count = 0;
10
11    // 0 if in range
12    if min.map_or(true, |m| m <= 0) && max.map_or(true, |m| m >= 0) {
13        count += 1;
14    }
15
16    // min boundary (if specified and not 0)
17    if let Some(m) = min {
18        if m != 0 {
19            count += 1;
20        }
21    }
22
23    // max boundary (if specified and not 0)
24    if let Some(m) = max {
25        if m != 0 {
26            count += 1;
27        }
28    }
29
30    // random negative — only if min is None
31    if min.is_none() {
32        count += 1;
33    }
34
35    // random positive — only if max is None
36    if max.is_none() {
37        count += 1;
38    }
39
40    count.max(1)
41}
42
43#[derive(Clone, Copy)]
44enum Variant {
45    Zero,
46    Min(i64),
47    Max(i64),
48    RandomNegative,
49    RandomPositive,
50}
51
52fn variants(schema: &IntegerInputSchema) -> Vec<Variant> {
53    let min = schema.minimum;
54    let max = schema.maximum;
55    let mut v = Vec::with_capacity(5);
56
57    if min.map_or(true, |m| m <= 0) && max.map_or(true, |m| m >= 0) {
58        v.push(Variant::Zero);
59    }
60    if let Some(m) = min {
61        if m != 0 {
62            v.push(Variant::Min(m));
63        }
64    }
65    if let Some(m) = max {
66        if m != 0 {
67            v.push(Variant::Max(m));
68        }
69    }
70    if min.is_none() {
71        v.push(Variant::RandomNegative);
72    }
73    if max.is_none() {
74        v.push(Variant::RandomPositive);
75    }
76    if v.is_empty() {
77        v.push(Variant::Zero);
78    }
79    v
80}
81
82pub fn generate<R: Rng>(
83    schema: &IntegerInputSchema,
84    mut rng: R,
85) -> Generator<R> {
86    let vars = variants(schema);
87    let mut indices: Vec<usize> = (0..vars.len()).collect();
88    indices.shuffle(&mut rng);
89    Generator {
90        variants: vars,
91        indices,
92        pos: 0,
93        rng,
94        min: schema.minimum,
95        max: schema.maximum,
96    }
97}
98
99pub struct Generator<R: Rng> {
100    variants: Vec<Variant>,
101    indices: Vec<usize>,
102    pos: usize,
103    rng: R,
104    min: Option<i64>,
105    max: Option<i64>,
106}
107
108impl<R: Rng> Iterator for Generator<R> {
109    type Item = InputValue;
110    fn next(&mut self) -> Option<InputValue> {
111        if self.pos >= self.indices.len() {
112            self.indices.shuffle(&mut self.rng);
113            self.pos = 0;
114        }
115        let index = self.indices[self.pos];
116        self.pos += 1;
117        let value = match self.variants[index] {
118            Variant::Zero => 0,
119            Variant::Min(m) => m,
120            Variant::Max(m) => m,
121            Variant::RandomNegative => {
122                let upper = self.max.unwrap_or(-1).min(-1);
123                self.rng.random_range(-100..=upper)
124            }
125            Variant::RandomPositive => {
126                let lower = self.min.unwrap_or(1).max(1);
127                self.rng.random_range(lower..=100)
128            }
129        };
130        Some(InputValue::Integer(value))
131    }
132}