diff --git a/src/rvm/vm/functions.rs b/src/rvm/vm/functions.rs index 7898acf..9bd0beb 100644 --- a/src/rvm/vm/functions.rs +++ b/src/rvm/vm/functions.rs @@ -1,5 +1,6 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT License. +use crate::builtins; use crate::value::Value; use alloc::string::String; use alloc::vec::Vec; @@ -41,6 +42,11 @@ impl RegoVM { }); } + if args.iter().any(|a| a == &Value::Undefined) { + self.registers[params.dest as usize] = Value::Undefined; + return Ok(()); + } + if let Some(builtin_fcn) = self.program.get_resolved_builtin(params.builtin_index) { let dummy_source = crate::lexer::Source::from_contents("arg".into(), String::new())?; let dummy_span = crate::lexer::Span { @@ -61,8 +67,31 @@ impl RegoVM { dummy_exprs.push(crate::ast::Ref::new(dummy_expr)); } - let result = (builtin_fcn.0)(&dummy_span, &dummy_exprs, &args, true)?; - self.registers[params.dest as usize] = result.clone(); + let cache_name = builtins::must_cache(builtin_info.name.as_str()); + if let Some(name) = cache_name { + if let Some(value) = self.builtins_cache.get(&(name, args.clone())) { + self.registers[params.dest as usize] = value.clone(); + return Ok(()); + } + } + + let result = + match (builtin_fcn.0)(&dummy_span, &dummy_exprs, &args, self.strict_builtin_errors) + { + Ok(value) => value, + Err(_) if !self.strict_builtin_errors => Value::Undefined, + Err(err) => return Err(err.into()), + }; + + if result == Value::Undefined { + self.registers[params.dest as usize] = Value::Undefined; + } else { + self.registers[params.dest as usize] = result.clone(); + } + + if let Some(name) = cache_name { + self.builtins_cache.insert((name, args), result); + } } else { return Err(VmError::BuiltinNotResolved { name: builtin_info.name.clone(), diff --git a/src/rvm/vm/machine.rs b/src/rvm/vm/machine.rs index e9d3478..b7ecbb2 100644 --- a/src/rvm/vm/machine.rs +++ b/src/rvm/vm/machine.rs @@ -98,6 +98,9 @@ pub struct RegoVM { /// Whether builtins should raise errors strictly or return undefined on failure pub(super) strict_builtin_errors: bool, + + /// Cache for builtin calls that must stay deterministic across a single evaluation + pub(super) builtins_cache: BTreeMap<(&'static str, Vec), Value>, } impl Default for RegoVM { @@ -135,6 +138,7 @@ impl RegoVM { execution_mode: ExecutionMode::RunToCompletion, frame_pc_overridden: false, strict_builtin_errors: false, + builtins_cache: BTreeMap::new(), } } diff --git a/src/rvm/vm/state.rs b/src/rvm/vm/state.rs index 86e307a..3001448 100644 --- a/src/rvm/vm/state.rs +++ b/src/rvm/vm/state.rs @@ -32,6 +32,9 @@ impl RegoVM { self.registers.clear(); self.registers .resize(self.base_register_count, Value::Undefined); + + // Builtin cache entries only live for a single execution + self.builtins_cache.clear(); } /// Return all active objects to their respective pools for reuse diff --git a/tests/opa.rs b/tests/opa.rs index 39b7a40..d94e444 100644 --- a/tests/opa.rs +++ b/tests/opa.rs @@ -36,7 +36,6 @@ const OPA_TODO_FOLDERS: &[&str] = &[ "partialdocconstants", "partialobjectdoc", "planner-ir", - "rand", "refheads", "replacen", "semverisvalid", diff --git a/tests/rvm/rego/cases/builtins_cache.yaml b/tests/rvm/rego/cases/builtins_cache.yaml new file mode 100644 index 0000000..b43c460 --- /dev/null +++ b/tests/rvm/rego/cases/builtins_cache.yaml @@ -0,0 +1,15 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +cases: + - note: builtin_rand_intn_cache_consistency + data: {} + modules: + - | + package test + + rands := { rand.intn("seed", 100) | numbers.range(1, 100)[_] } + + np := count(rands) + query: data.test.np + want_result: 1