vrl/stdlib/
random_int.rs

1use crate::compiler::prelude::*;
2use rand::{Rng, thread_rng};
3use std::ops::Range;
4
5const INVALID_RANGE_ERR: &str = "max must be greater than min";
6
7fn random_int(min: Value, max: Value) -> Resolved {
8    let range = get_range(min, max)?;
9
10    let i: i64 = thread_rng().gen_range(range);
11
12    Ok(Value::Integer(i))
13}
14
15fn get_range(min: Value, max: Value) -> std::result::Result<Range<i64>, &'static str> {
16    let min = min.try_integer().expect("min must be an integer");
17    let max = max.try_integer().expect("max must be an integer");
18
19    if max <= min {
20        return Err(INVALID_RANGE_ERR);
21    }
22
23    Ok(min..max)
24}
25
26#[derive(Clone, Copy, Debug)]
27pub struct RandomInt;
28
29impl Function for RandomInt {
30    fn identifier(&self) -> &'static str {
31        "random_int"
32    }
33
34    fn usage(&self) -> &'static str {
35        "Returns a random integer between [min, max)."
36    }
37
38    fn category(&self) -> &'static str {
39        Category::Random.as_ref()
40    }
41
42    fn internal_failure_reasons(&self) -> &'static [&'static str] {
43        &["`max` is not greater than `min`."]
44    }
45
46    fn return_kind(&self) -> u16 {
47        kind::INTEGER
48    }
49
50    fn parameters(&self) -> &'static [Parameter] {
51        const PARAMETERS: &[Parameter] = &[
52            Parameter::required("min", kind::INTEGER, "Minimum value (inclusive)."),
53            Parameter::required("max", kind::INTEGER, "Maximum value (exclusive)."),
54        ];
55        PARAMETERS
56    }
57
58    fn examples(&self) -> &'static [Example] {
59        &[example! {
60            title: "Random integer from 0 to 10, not including 10",
61            source: indoc! {"
62                i = random_int(0, 10)
63                i >= 0 && i < 10
64            "},
65            result: Ok("true"),
66        }]
67    }
68
69    fn compile(
70        &self,
71        state: &state::TypeState,
72        _ctx: &mut FunctionCompileContext,
73        arguments: ArgumentList,
74    ) -> Compiled {
75        let min = arguments.required("min");
76        let max = arguments.required("max");
77
78        if let (Some(min), Some(max)) = (min.resolve_constant(state), max.resolve_constant(state)) {
79            // check if range is valid
80            let _: Range<i64> =
81                get_range(min, max.clone()).map_err(|err| function::Error::InvalidArgument {
82                    keyword: "max",
83                    value: max,
84                    error: err,
85                })?;
86        }
87
88        Ok(RandomIntFn { min, max }.as_expr())
89    }
90}
91
92#[derive(Debug, Clone)]
93struct RandomIntFn {
94    min: Box<dyn Expression>,
95    max: Box<dyn Expression>,
96}
97
98impl FunctionExpression for RandomIntFn {
99    fn resolve(&self, ctx: &mut Context) -> Resolved {
100        let min = self.min.resolve(ctx)?;
101        let max = self.max.resolve(ctx)?;
102
103        random_int(min, max)
104    }
105
106    fn type_def(&self, state: &state::TypeState) -> TypeDef {
107        match (
108            self.min.resolve_constant(state),
109            self.max.resolve_constant(state),
110        ) {
111            (Some(min), Some(max)) => {
112                if get_range(min, max).is_ok() {
113                    TypeDef::integer()
114                } else {
115                    TypeDef::integer().fallible()
116                }
117            }
118            _ => TypeDef::integer().fallible(),
119        }
120    }
121}
122
123#[cfg(test)]
124mod tests {
125    use super::*;
126    use crate::value;
127    // positive tests are handled by examples
128
129    test_function![
130        random_int => RandomInt;
131
132        bad_range {
133            args: func_args![min: value!(1), max: value!(1)],
134            want: Err("invalid argument"),
135            tdef: TypeDef::integer().fallible(),
136        }
137    ];
138}