objectiveai_sdk/functions/check/example_inputs/
integer.rs1use 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 if min.map_or(true, |m| m <= 0) && max.map_or(true, |m| m >= 0) {
13 count += 1;
14 }
15
16 if let Some(m) = min {
18 if m != 0 {
19 count += 1;
20 }
21 }
22
23 if let Some(m) = max {
25 if m != 0 {
26 count += 1;
27 }
28 }
29
30 if min.is_none() {
32 count += 1;
33 }
34
35 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}