Skip to main content

objectiveai_sdk/
arbitrary_util.rs

1//! Arbitrary helpers for types that don't natively implement `Arbitrary`
2//! or that need bounded generation to avoid extreme values.
3
4fn arbitrary_option<T>(
5    u: &mut arbitrary::Unstructured,
6    arbitrary_value: impl Fn(&mut arbitrary::Unstructured) -> arbitrary::Result<T>,
7) -> arbitrary::Result<Option<T>> {
8    if u.arbitrary().unwrap_or(false) {
9        Ok(Some(arbitrary_value(u)?))
10    } else {
11        Ok(None)
12    }
13}
14
15fn arbitrary_hash_map<K: std::hash::Hash + Eq, V>(
16    u: &mut arbitrary::Unstructured,
17    arbitrary_key: impl Fn(&mut arbitrary::Unstructured) -> arbitrary::Result<K>,
18    arbitrary_value: impl Fn(&mut arbitrary::Unstructured) -> arbitrary::Result<V>,
19) -> arbitrary::Result<std::collections::HashMap<K, V>> {
20    let mut map = std::collections::HashMap::new();
21    while u.arbitrary().unwrap_or(false) {
22        map.insert(arbitrary_key(u)?, arbitrary_value(u)?);
23    }
24    Ok(map)
25}
26
27fn arbitrary_index_map<K: std::hash::Hash + Eq, V>(
28    u: &mut arbitrary::Unstructured,
29    arbitrary_key: impl Fn(&mut arbitrary::Unstructured) -> arbitrary::Result<K>,
30    arbitrary_value: impl Fn(&mut arbitrary::Unstructured) -> arbitrary::Result<V>,
31) -> arbitrary::Result<indexmap::IndexMap<K, V>> {
32    let mut map = indexmap::IndexMap::new();
33    while u.arbitrary().unwrap_or(false) {
34        map.insert(arbitrary_key(u)?, arbitrary_value(u)?);
35    }
36    Ok(map)
37}
38
39fn arbitrary_vec<T>(
40    u: &mut arbitrary::Unstructured,
41    arbitrary_value: impl Fn(&mut arbitrary::Unstructured) -> arbitrary::Result<T>,
42) -> arbitrary::Result<Vec<T>> {
43    let mut v = Vec::new();
44    while u.arbitrary().unwrap_or(false) {
45        v.push(arbitrary_value(u)?);
46    }
47    Ok(v)
48}
49
50/// Generates an arbitrary `rust_decimal::Decimal` from an f32,
51/// scaled down by 10^13. f32::MAX is ~3.4e38, so every finite
52/// sample lands at ≤ ~3.4e25 — a hard ceiling ~2,300× below
53/// `Decimal::MAX` (~7.9e28), leaving ample headroom for the
54/// chunk-push suites that sum ~20 of these through `Usage::push`
55/// (`Decimal::AddAssign` panics on overflow). `Inf`/`NaN` samples
56/// stay non-finite after division and fall back to zero via
57/// `unwrap_or_default`.
58pub fn arbitrary_rust_decimal(
59    u: &mut arbitrary::Unstructured,
60) -> arbitrary::Result<rust_decimal::Decimal> {
61    let f: f32 = u.arbitrary()?;
62    Ok(rust_decimal::Decimal::try_from(f / 1e13).unwrap_or_default())
63}
64
65/// Generates an arbitrary `Option<rust_decimal::Decimal>`.
66pub fn arbitrary_option_rust_decimal(
67    u: &mut arbitrary::Unstructured,
68) -> arbitrary::Result<Option<rust_decimal::Decimal>> {
69    arbitrary_option(u, arbitrary_rust_decimal)
70}
71
72/// Generates an arbitrary `Vec<rust_decimal::Decimal>`.
73pub fn arbitrary_vec_rust_decimal(
74    u: &mut arbitrary::Unstructured,
75) -> arbitrary::Result<Vec<rust_decimal::Decimal>> {
76    arbitrary_vec(u, arbitrary_rust_decimal)
77}
78
79/// Generates an arbitrary `Vec<Vec<rust_decimal::Decimal>>`.
80pub fn arbitrary_vec_vec_rust_decimal(
81    u: &mut arbitrary::Unstructured,
82) -> arbitrary::Result<Vec<Vec<rust_decimal::Decimal>>> {
83    arbitrary_vec(u, arbitrary_vec_rust_decimal)
84}
85
86/// Generates an arbitrary `f64` from an f32 (bounded range).
87pub fn arbitrary_f64(
88    u: &mut arbitrary::Unstructured,
89) -> arbitrary::Result<f64> {
90    Ok(u.arbitrary::<f32>()? as f64)
91}
92
93/// Generates an arbitrary `Option<f64>` (bounded range).
94pub fn arbitrary_option_f64(
95    u: &mut arbitrary::Unstructured,
96) -> arbitrary::Result<Option<f64>> {
97    arbitrary_option(u, arbitrary_f64)
98}
99
100/// Generates an arbitrary `u64` from a u32 (bounded range).
101pub fn arbitrary_u64(
102    u: &mut arbitrary::Unstructured,
103) -> arbitrary::Result<u64> {
104    Ok(u.arbitrary::<u32>()? as u64)
105}
106
107/// Generates an arbitrary `usize` from a u16 (bounded range).
108pub fn arbitrary_usize(
109    u: &mut arbitrary::Unstructured,
110) -> arbitrary::Result<usize> {
111    Ok(u.arbitrary::<u16>()? as usize)
112}
113
114/// Generates an arbitrary `isize` from an i16 (bounded range).
115pub fn arbitrary_isize(
116    u: &mut arbitrary::Unstructured,
117) -> arbitrary::Result<isize> {
118    Ok(u.arbitrary::<i16>()? as isize)
119}
120
121/// Generates an arbitrary `Option<usize>` (bounded range).
122pub fn arbitrary_option_usize(
123    u: &mut arbitrary::Unstructured,
124) -> arbitrary::Result<Option<usize>> {
125    arbitrary_option(u, arbitrary_usize)
126}
127
128/// Generates an arbitrary `Option<isize>` (bounded range).
129pub fn arbitrary_option_isize(
130    u: &mut arbitrary::Unstructured,
131) -> arbitrary::Result<Option<isize>> {
132    arbitrary_option(u, arbitrary_isize)
133}
134
135/// Generates an arbitrary `Vec<u64>` with bounded elements.
136pub fn arbitrary_vec_u64(
137    u: &mut arbitrary::Unstructured,
138) -> arbitrary::Result<Vec<u64>> {
139    arbitrary_vec(u, arbitrary_u64)
140}
141
142/// Generates an arbitrary `IndexMap<K, V>` where both K and V implement Arbitrary.
143pub fn arbitrary_indexmap<'a, K, V>(
144    u: &mut arbitrary::Unstructured<'a>,
145) -> arbitrary::Result<indexmap::IndexMap<K, V>>
146where
147    K: arbitrary::Arbitrary<'a> + std::hash::Hash + Eq,
148    V: arbitrary::Arbitrary<'a>,
149{
150    let mut map = indexmap::IndexMap::new();
151    while u.arbitrary().unwrap_or(false) {
152        map.insert(K::arbitrary(u)?, V::arbitrary(u)?);
153    }
154    Ok(map)
155}
156
157/// Generates an arbitrary `Option<IndexMap<String, i64>>`.
158pub fn arbitrary_option_indexmap_string_i64(
159    u: &mut arbitrary::Unstructured,
160) -> arbitrary::Result<Option<indexmap::IndexMap<String, i64>>> {
161    arbitrary_option(u, |u| {
162        arbitrary_index_map(u, |u| u.arbitrary(), arbitrary_i64)
163    })
164}
165
166/// Generates an arbitrary `Option<IndexMap<String, Option<String>>>`.
167pub fn arbitrary_option_indexmap_string_option_string(
168    u: &mut arbitrary::Unstructured,
169) -> arbitrary::Result<Option<indexmap::IndexMap<String, Option<String>>>> {
170    arbitrary_option(u, |u| {
171        arbitrary_index_map(u, |u| u.arbitrary(), |u| u.arbitrary())
172    })
173}
174
175/// Generates an arbitrary `Option<u64>` (bounded range).
176pub fn arbitrary_option_u64(
177    u: &mut arbitrary::Unstructured,
178) -> arbitrary::Result<Option<u64>> {
179    arbitrary_option(u, arbitrary_u64)
180}
181
182/// Generates an arbitrary `i64` from an i32 (bounded range).
183pub fn arbitrary_i64(
184    u: &mut arbitrary::Unstructured,
185) -> arbitrary::Result<i64> {
186    Ok(u.arbitrary::<i32>()? as i64)
187}
188
189/// Generates an arbitrary `Option<i64>` (bounded range).
190pub fn arbitrary_option_i64(
191    u: &mut arbitrary::Unstructured,
192) -> arbitrary::Result<Option<i64>> {
193    arbitrary_option(u, arbitrary_i64)
194}
195
196/// Generates an arbitrary `Vec<i64>` (bounded range).
197pub fn arbitrary_vec_i64(
198    u: &mut arbitrary::Unstructured,
199) -> arbitrary::Result<Vec<i64>> {
200    arbitrary_vec(u, arbitrary_i64)
201}
202
203/// Generates an arbitrary `Option<Vec<i64>>` (bounded range).
204pub fn arbitrary_option_vec_i64(
205    u: &mut arbitrary::Unstructured,
206) -> arbitrary::Result<Option<Vec<i64>>> {
207    arbitrary_option(u, arbitrary_vec_i64)
208}
209
210/// Generates an arbitrary `serde_json::Value`. Nesting requires a "go deep"
211/// bool to be true (geometric distribution), so compound structures are
212/// exponentially rare. Collection sizes also use bool-per-element.
213pub fn arbitrary_json_value(
214    u: &mut arbitrary::Unstructured,
215) -> arbitrary::Result<serde_json::Value> {
216    if u.arbitrary().unwrap_or(false) {
217        // Compound value
218        if u.arbitrary()? {
219            let mut arr = Vec::new();
220            while u.arbitrary().unwrap_or(false) {
221                arr.push(arbitrary_json_value(u)?);
222            }
223            Ok(serde_json::Value::Array(arr))
224        } else {
225            let mut map = serde_json::Map::new();
226            while u.arbitrary().unwrap_or(false) {
227                let key: String = u.arbitrary()?;
228                map.insert(key, arbitrary_json_value(u)?);
229            }
230            Ok(serde_json::Value::Object(map))
231        }
232    } else {
233        // Leaf value
234        match u.int_in_range(0..=3)? {
235            0 => Ok(serde_json::Value::Null),
236            1 => Ok(serde_json::Value::Bool(u.arbitrary()?)),
237            2 => Ok(serde_json::Value::Number(serde_json::Number::from(
238                u.arbitrary::<i32>()?,
239            ))),
240            _ => Ok(serde_json::Value::String(u.arbitrary()?)),
241        }
242    }
243}