jpx_core/extensions/
random.rs1use 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
12pub 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
20pub 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 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 rng.gen_range(0.0..1.0)
53 } else {
54 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
68pub 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 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 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 result.shuffle(&mut rand::thread_rng());
114 }
115
116 Ok(Value::Array(result))
117 }
118}
119
120pub 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 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 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 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
178defn!(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}