1use sqlite_loadable::api::ValueType;
2use sqlite_loadable::{api, define_scalar_function, define_scalar_function_with_aux, Result};
3use sqlite_loadable::{prelude::*, Error};
4
5use fastrand::Rng;
6use std::cell::{Ref, RefCell};
7use std::iter::repeat_with;
8use std::ops::Range;
9use std::rc::Rc;
10
11enum BtwnValues<'a> {
12 Range(&'a *mut sqlite3_value, &'a *mut sqlite3_value),
13 From(&'a *mut sqlite3_value),
14 To(&'a *mut sqlite3_value),
15 Full,
16}
17
18fn calculate_btwn_values(values: &[*mut sqlite3_value]) -> Result<BtwnValues> {
19 let a = values.get(0);
20 let b = values.get(1);
21 return match (a, b) {
22 (Some(a), Some(b)) => {
23 match (
24 api::value_type(a) == ValueType::Null,
25 api::value_type(b) == ValueType::Null,
26 ) {
27 (true, true) => Ok(BtwnValues::Full),
28 (false, true) => Ok(BtwnValues::From(a)),
29 (false, false) => Ok(BtwnValues::Range(a, b)),
30 (true, false) => Ok(BtwnValues::To(b)),
31 }
32 }
33 (None, None) => Ok(BtwnValues::Full),
34 (None, Some(_)) => Err(Error::new_message(
35 "sqlite-fastrand internal error: argv[1] provided but not argv[0]",
36 )),
37 (Some(a), None) => {
38 if api::value_type(a) == ValueType::Null {
39 Ok(BtwnValues::Full)
40 } else {
41 Ok(BtwnValues::From(a))
42 }
43 }
44 };
45}
46
47pub fn fastrand_version(
48 context: *mut sqlite3_context,
49 _values: &[*mut sqlite3_value],
50) -> Result<()> {
51 api::result_text(context, format!("v{}", env!("CARGO_PKG_VERSION")))?;
52 Ok(())
53}
54
55pub fn fastrand_debug(context: *mut sqlite3_context, _values: &[*mut sqlite3_value]) -> Result<()> {
56 api::result_text(
57 context,
58 format!(
59 "Version: v{}
60Source: {}
61",
62 env!("CARGO_PKG_VERSION"),
63 env!("GIT_HASH")
64 ),
65 )?;
66 Ok(())
67}
68
69fn rng_borrow_or_err(rng: &Rc<RefCell<Rng>>) -> Result<Ref<Rng>> {
70 match rng.try_borrow() {
71 Ok(rng) => Ok(rng),
72 Err(err) => Err(Error::new_message(
73 format!("internal sqlite-fastrand error, could not borrow shared RNG: {err}").as_str(),
74 )),
75 }
76}
77
78pub fn fastrand_seed_set(
79 context: *mut sqlite3_context,
80 values: &[*mut sqlite3_value],
81 rng: &Rc<RefCell<Rng>>,
82) -> Result<()> {
83 let seed: u64 = api::value_int64(
84 values
85 .get(0)
86 .ok_or_else(|| Error::new_message("expected seed as 1st argument"))?,
87 )
88 .try_into()
89 .map_err(|e| {
90 Error::new_message(format!("seed must be an usigned 64 bit integer: {e}").as_str())
91 })?;
92 rng.borrow_mut().seed(seed);
93 api::result_bool(context, true);
94 Ok(())
95}
96
97pub fn fastrand_seed_get(
98 context: *mut sqlite3_context,
99 _values: &[*mut sqlite3_value],
100 rng: &Rc<RefCell<Rng>>,
101) -> Result<()> {
102 api::result_text(context, rng_borrow_or_err(rng)?.get_seed().to_string())?;
103 Ok(())
104}
105
106pub fn fastrand_int(
107 context: *mut sqlite3_context,
108 values: &[*mut sqlite3_value],
109 rng: &Rc<RefCell<Rng>>,
110) -> Result<()> {
111 let i = match calculate_btwn_values(values)? {
112 BtwnValues::Range(start, end) => {
113 let range: Range<i32> = std::ops::Range {
114 start: api::value_int(start),
115 end: api::value_int(end),
116 };
117 if range.is_empty() {
118 return Err("fuk".into());
119 }
120 rng_borrow_or_err(rng)?.i32(range)
121 }
122 BtwnValues::From(start) => {
123 let range = std::ops::RangeFrom {
124 start: api::value_int(start),
125 };
126 rng_borrow_or_err(rng)?.i32(range)
127 }
128 BtwnValues::To(end) => {
129 let range = std::ops::RangeTo {
130 end: api::value_int(end),
131 };
132 rng_borrow_or_err(rng)?.i32(range)
133 }
134 BtwnValues::Full => rng_borrow_or_err(rng)?.i32(..),
135 };
136 api::result_int(context, i);
137 Ok(())
138}
139
140pub fn fastrand_int64(
141 context: *mut sqlite3_context,
142 values: &[*mut sqlite3_value],
143 rng: &Rc<RefCell<Rng>>,
144) -> Result<()> {
145 let i = match calculate_btwn_values(values)? {
146 BtwnValues::Range(start, end) => {
147 let range: Range<i64> = std::ops::Range {
148 start: api::value_int64(start),
149 end: api::value_int64(end),
150 };
151 if range.is_empty() {
152 return Err("fuk".into());
153 }
154 rng_borrow_or_err(rng)?.i64(range)
155 }
156 BtwnValues::From(start) => {
157 let range = std::ops::RangeFrom {
158 start: api::value_int64(start),
159 };
160 rng_borrow_or_err(rng)?.i64(range)
161 }
162 BtwnValues::To(end) => {
163 let range = std::ops::RangeTo {
164 end: api::value_int64(end),
165 };
166 rng_borrow_or_err(rng)?.i64(range)
167 }
168 BtwnValues::Full => rng_borrow_or_err(rng)?.i64(..),
169 };
170 api::result_int64(context, i);
171 Ok(())
172}
173
174pub fn fastrand_double(
175 context: *mut sqlite3_context,
176 _values: &[*mut sqlite3_value],
177 rng: &Rc<RefCell<Rng>>,
178) -> Result<()> {
179 api::result_double(context, rng_borrow_or_err(rng)?.f64());
180 Ok(())
181}
182
183pub fn fastrand_blob(
184 context: *mut sqlite3_context,
185 values: &[*mut sqlite3_value],
186 rng: &Rc<RefCell<Rng>>,
187) -> Result<()> {
188 let n: usize = api::value_int64(
189 values
190 .get(0)
191 .ok_or_else(|| Error::new_message("expected N as 1st argument"))?,
192 ) as usize;
193 let rng = rng_borrow_or_err(rng)?;
194 let bytes: Vec<u8> = repeat_with(|| rng.u8(..)).take(n).collect();
195 api::result_blob(context, bytes.as_slice());
196 Ok(())
197}
198
199pub fn fastrand_bool(
200 context: *mut sqlite3_context,
201 _values: &[*mut sqlite3_value],
202 rng: &Rc<RefCell<Rng>>,
203) -> Result<()> {
204 api::result_bool(context, rng_borrow_or_err(rng)?.bool());
205 Ok(())
206}
207
208pub fn fastrand_alphabetic(
209 context: *mut sqlite3_context,
210 _values: &[*mut sqlite3_value],
211 rng: &Rc<RefCell<Rng>>,
212) -> Result<()> {
213 api::result_text(context, rng_borrow_or_err(rng)?.alphabetic().to_string())?;
214 Ok(())
215}
216
217pub fn fastrand_alphanumeric(
218 context: *mut sqlite3_context,
219 _values: &[*mut sqlite3_value],
220 rng: &Rc<RefCell<Rng>>,
221) -> Result<()> {
222 api::result_text(context, rng_borrow_or_err(rng)?.alphanumeric().to_string())?;
223 Ok(())
224}
225
226pub fn fastrand_lowercase(
227 context: *mut sqlite3_context,
228 _values: &[*mut sqlite3_value],
229 rng: &Rc<RefCell<Rng>>,
230) -> Result<()> {
231 api::result_text(context, rng_borrow_or_err(rng)?.lowercase().to_string())?;
232 Ok(())
233}
234
235pub fn fastrand_uppercase(
236 context: *mut sqlite3_context,
237 _values: &[*mut sqlite3_value],
238 rng: &Rc<RefCell<Rng>>,
239) -> Result<()> {
240 api::result_text(context, rng_borrow_or_err(rng)?.uppercase().to_string())?;
241 Ok(())
242}
243
244pub fn fastrand_digit(
245 context: *mut sqlite3_context,
246 values: &[*mut sqlite3_value],
247 rng: &Rc<RefCell<Rng>>,
248) -> Result<()> {
249 let base: u32 = api::value_int64(
250 values
251 .get(0)
252 .ok_or_else(|| Error::new_message("expected base as 1st argument"))?,
253 )
254 .try_into()
255 .map_err(|e| {
256 Error::new_message(format!("base must be an unsigned 32-bit integer: {e}").as_str())
257 })?;
258 match base {
259 1..=35 => api::result_text(context, rng_borrow_or_err(rng)?.digit(base).to_string())?,
260 _ => {
261 return Err(Error::new_message(
262 "base must be greater than 0 and less than 26",
263 ))
264 }
265 }
266
267 Ok(())
268}
269
270#[sqlite_entrypoint]
271pub fn sqlite3_fastrand_init(db: *mut sqlite3) -> Result<()> {
272 let flags = FunctionFlags::empty();
273 let rng = Rc::new(RefCell::new(Rng::new()));
274
275 define_scalar_function(
276 db,
277 "fastrand_version",
278 0,
279 fastrand_version,
280 FunctionFlags::UTF8 | FunctionFlags::DETERMINISTIC,
281 )?;
282 define_scalar_function(
283 db,
284 "fastrand_debug",
285 0,
286 fastrand_debug,
287 FunctionFlags::UTF8 | FunctionFlags::DETERMINISTIC,
288 )?;
289
290 define_scalar_function_with_aux(
291 db,
292 "fastrand_seed_get",
293 0,
294 fastrand_seed_get,
295 flags,
296 Rc::clone(&rng),
297 )?;
298 define_scalar_function_with_aux(
299 db,
300 "fastrand_seed_set",
301 1,
302 fastrand_seed_set,
303 flags,
304 Rc::clone(&rng),
305 )?;
306
307 define_scalar_function_with_aux(db, "fastrand_int", 0, fastrand_int, flags, Rc::clone(&rng))?;
308 define_scalar_function_with_aux(db, "fastrand_int", 1, fastrand_int, flags, Rc::clone(&rng))?;
309 define_scalar_function_with_aux(db, "fastrand_int", 2, fastrand_int, flags, Rc::clone(&rng))?;
310 define_scalar_function_with_aux(
311 db,
312 "fastrand_int64",
313 0,
314 fastrand_int64,
315 flags,
316 Rc::clone(&rng),
317 )?;
318 define_scalar_function_with_aux(
319 db,
320 "fastrand_int64",
321 1,
322 fastrand_int64,
323 flags,
324 Rc::clone(&rng),
325 )?;
326 define_scalar_function_with_aux(
327 db,
328 "fastrand_int64",
329 2,
330 fastrand_int64,
331 flags,
332 Rc::clone(&rng),
333 )?;
334 define_scalar_function_with_aux(
335 db,
336 "fastrand_double",
337 0,
338 fastrand_double,
339 flags,
340 Rc::clone(&rng),
341 )?;
342 define_scalar_function_with_aux(
343 db,
344 "fastrand_blob",
345 1,
346 fastrand_blob,
347 flags,
348 Rc::clone(&rng),
349 )?;
350
351 define_scalar_function_with_aux(
352 db,
353 "fastrand_bool",
354 0,
355 fastrand_bool,
356 flags,
357 Rc::clone(&rng),
358 )?;
359
360 define_scalar_function_with_aux(
361 db,
362 "fastrand_alphabetic",
363 0,
364 fastrand_alphabetic,
365 flags,
366 Rc::clone(&rng),
367 )?;
368 define_scalar_function_with_aux(
369 db,
370 "fastrand_alphanumeric",
371 0,
372 fastrand_alphanumeric,
373 flags,
374 Rc::clone(&rng),
375 )?;
376 define_scalar_function_with_aux(
377 db,
378 "fastrand_uppercase",
379 0,
380 fastrand_uppercase,
381 flags,
382 Rc::clone(&rng),
383 )?;
384 define_scalar_function_with_aux(
385 db,
386 "fastrand_lowercase",
387 0,
388 fastrand_lowercase,
389 flags,
390 Rc::clone(&rng),
391 )?;
392
393 define_scalar_function_with_aux(
394 db,
395 "fastrand_digit",
396 1,
397 fastrand_digit,
398 flags,
399 Rc::clone(&rng),
400 )?;
401 Ok(())
402}