Skip to main content

jpx_core/extensions/
random.rs

1//! Random value generation functions.
2
3use std::collections::HashSet;
4
5use serde_json::Value;
6
7use crate::functions::{Function, custom_error, number_value};
8use crate::interpreter::SearchResult;
9use crate::registry::register_if_enabled;
10use crate::{Context, Runtime, defn};
11
12/// Register random functions filtered by the enabled set.
13pub fn register_filtered(runtime: &mut Runtime, enabled: &HashSet<&str>) {
14    register_if_enabled(runtime, "random", enabled, Box::new(RandomFn::new()));
15    register_if_enabled(runtime, "shuffle", enabled, Box::new(ShuffleFn::new()));
16    register_if_enabled(runtime, "sample", enabled, Box::new(SampleFn::new()));
17    register_if_enabled(runtime, "uuid", enabled, Box::new(UuidFn::new()));
18}
19
20// =============================================================================
21// random() -> number (0.0 to 1.0)
22// random(min, max) -> number in range [min, max)
23// =============================================================================
24
25pub struct RandomFn;
26
27impl Default for RandomFn {
28    fn default() -> Self {
29        Self::new()
30    }
31}
32
33impl RandomFn {
34    pub fn new() -> RandomFn {
35        RandomFn
36    }
37}
38
39impl Function for RandomFn {
40    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
41        use rand::Rng;
42
43        // Manual validation: accept 0 or 2 arguments
44        if !args.is_empty() && args.len() != 2 {
45            return Err(custom_error(ctx, "random() takes 0 or 2 arguments"));
46        }
47
48        let mut rng = rand::thread_rng();
49
50        let value: f64 = if args.is_empty() {
51            // random() - return 0.0 to 1.0
52            rng.gen_range(0.0..1.0)
53        } else {
54            // random(min, max) - return min to max
55            let min = args[0]
56                .as_f64()
57                .ok_or_else(|| custom_error(ctx, "Expected number for min"))?;
58            let max = args[1]
59                .as_f64()
60                .ok_or_else(|| custom_error(ctx, "Expected number for max"))?;
61            rng.gen_range(min..max)
62        };
63
64        Ok(number_value(value))
65    }
66}
67
68// =============================================================================
69// shuffle(array) -> array (randomly shuffled)
70// shuffle(array, seed) -> array (deterministically shuffled)
71// =============================================================================
72
73pub struct ShuffleFn;
74
75impl Default for ShuffleFn {
76    fn default() -> Self {
77        Self::new()
78    }
79}
80
81impl ShuffleFn {
82    pub fn new() -> ShuffleFn {
83        ShuffleFn
84    }
85}
86
87impl Function for ShuffleFn {
88    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
89        // Manual validation: 1 or 2 arguments
90        if args.is_empty() || args.len() > 2 {
91            return Err(custom_error(ctx, "shuffle() takes 1 or 2 arguments"));
92        }
93
94        let arr = args[0]
95            .as_array()
96            .ok_or_else(|| custom_error(ctx, "Expected array argument"))?;
97
98        use rand::SeedableRng;
99        use rand::seq::SliceRandom;
100
101        let mut result: Vec<Value> = arr.clone();
102
103        if args.len() == 2 {
104            // Deterministic shuffle with seed
105            let seed = args[1]
106                .as_f64()
107                .ok_or_else(|| custom_error(ctx, "Expected number for seed"))?
108                as u64;
109            let mut rng = rand::rngs::StdRng::seed_from_u64(seed);
110            result.shuffle(&mut rng);
111        } else {
112            // Random shuffle
113            result.shuffle(&mut rand::thread_rng());
114        }
115
116        Ok(Value::Array(result))
117    }
118}
119
120// =============================================================================
121// sample(array, n) -> array (random sample of n elements)
122// sample(array, n, seed) -> array (deterministic sample)
123// =============================================================================
124
125pub struct SampleFn;
126
127impl Default for SampleFn {
128    fn default() -> Self {
129        Self::new()
130    }
131}
132
133impl SampleFn {
134    pub fn new() -> SampleFn {
135        SampleFn
136    }
137}
138
139impl Function for SampleFn {
140    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
141        // Manual validation: 2 or 3 arguments
142        if args.len() < 2 || args.len() > 3 {
143            return Err(custom_error(ctx, "sample() takes 2 or 3 arguments"));
144        }
145
146        let arr = args[0]
147            .as_array()
148            .ok_or_else(|| custom_error(ctx, "Expected array argument"))?;
149
150        let n = args[1]
151            .as_f64()
152            .ok_or_else(|| custom_error(ctx, "Expected number argument"))? as usize;
153
154        use rand::SeedableRng;
155        use rand::seq::SliceRandom;
156
157        let sample: Vec<Value> = if args.len() == 3 {
158            // Deterministic sample with seed
159            let seed = args[2]
160                .as_f64()
161                .ok_or_else(|| custom_error(ctx, "Expected number for seed"))?
162                as u64;
163            let mut rng = rand::rngs::StdRng::seed_from_u64(seed);
164            arr.choose_multiple(&mut rng, n.min(arr.len()))
165                .cloned()
166                .collect()
167        } else {
168            // Random sample
169            arr.choose_multiple(&mut rand::thread_rng(), n.min(arr.len()))
170                .cloned()
171                .collect()
172        };
173
174        Ok(Value::Array(sample))
175    }
176}
177
178// =============================================================================
179// uuid() -> string (UUID v4)
180// =============================================================================
181
182defn!(UuidFn, vec![], None);
183
184impl Function for UuidFn {
185    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
186        self.signature.validate(args, ctx)?;
187
188        let id = uuid::Uuid::new_v4();
189        Ok(Value::String(id.to_string()))
190    }
191}