From e3d23766ae067c470f9ea3c8cb62e5f2a5a457dc Mon Sep 17 00:00:00 2001 From: Anand Krishnamoorthi Date: Wed, 3 Dec 2025 10:45:29 -0600 Subject: [PATCH 1/3] feat: Ensure RVM caches deterministic builtins MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Mirror interpreter implementation: - use builtins::must_cache to determine whether builtin must be cached. - reuse cached value when applicable - clear the VM’s builtin cache whenever execution state resets to avoid leaking values across runs - add a YAML regression for rand.intn set comprehensions and re-enable the rand cases in the OPA test suite Signed-off-by: Anand Krishnamoorthi --- src/rvm/vm/functions.rs | 33 ++++++++++++++++++++++-- src/rvm/vm/machine.rs | 4 +++ src/rvm/vm/state.rs | 3 +++ tests/opa.rs | 1 - tests/rvm/rego/cases/builtins_cache.yaml | 15 +++++++++++ 5 files changed, 53 insertions(+), 3 deletions(-) create mode 100644 tests/rvm/rego/cases/builtins_cache.yaml 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 From 8269968c4a9830aa8de75c87bc353dae53f3b187 Mon Sep 17 00:00:00 2001 From: Anand Krishnamoorthi Date: Wed, 3 Dec 2025 11:00:57 -0600 Subject: [PATCH 2/3] feat: Handle computed reference roots in RVM compiler MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Allow compile_chained_ref to fall back to “evaluate root expression → chain access” so literal arrays, comprehensions, and other computed roots no longer raise NotSimpleReferenceChain. Signed-off-by: Anand Krishnamoorthi --- src/languages/rego/compiler/references.rs | 62 +++++++++++++++++------ tests/opa.rs | 4 -- tests/rvm/rego/cases/chained_access.yaml | 24 +++++++++ 3 files changed, 70 insertions(+), 20 deletions(-) diff --git a/src/languages/rego/compiler/references.rs b/src/languages/rego/compiler/references.rs index 0d739b3..327db58 100644 --- a/src/languages/rego/compiler/references.rs +++ b/src/languages/rego/compiler/references.rs @@ -24,26 +24,37 @@ pub(super) enum AccessComponent { Expression(ExprRef), } +/// Root of a reference chain - either a named variable or another arbitrary expression +#[derive(Debug, Clone)] +pub(super) enum ReferenceRoot { + Variable(String), + Expression(ExprRef), +} + /// Represents a chained reference like data.a.b[expr].c[expr] #[derive(Debug, Clone)] pub(super) struct ReferenceChain { - /// The root variable (e.g., "data", "input", "local_var") - pub(super) root: String, + /// The root of the chain (variable or arbitrary expression) + pub(super) root: ReferenceRoot, /// Chain of field accesses - either literal field names or dynamic expressions pub(super) components: Vec, } impl ReferenceChain { /// Get the static prefix path (all literal components from the start) - pub(super) fn get_static_prefix(&self) -> Vec<&str> { - let mut prefix = vec![self.root.as_str()]; + pub(super) fn get_static_prefix(&self) -> Option> { + let ReferenceRoot::Variable(root) = &self.root else { + return None; + }; + + let mut prefix = vec![root.as_str()]; for component in &self.components { match component { AccessComponent::Field(field) => prefix.push(field.as_str()), AccessComponent::Expression(_) => break, } } - prefix + Some(prefix) } } @@ -57,7 +68,7 @@ pub(super) fn parse_reference_chain(expr: &ExprRef) -> Result { match current_expr.as_ref() { Expr::Var { span, .. } => { // Found the root variable - let root = span.text().to_string(); + let root = ReferenceRoot::Variable(span.text().to_string()); components.reverse(); // We built backwards, so reverse return Ok(ReferenceChain { root, components }); } @@ -81,7 +92,12 @@ pub(super) fn parse_reference_chain(expr: &ExprRef) -> Result { current_expr = refr; } _ => { - return Err(CompilerError::NotSimpleReferenceChain.at(current_expr.span())); + // Fallback root expression (e.g., array literal, function call) + components.reverse(); + return Ok(ReferenceChain { + root: ReferenceRoot::Expression(current_expr.clone()), + components, + }); } } } @@ -94,10 +110,17 @@ impl<'a> Compiler<'a> { // Parse the expression into a reference chain let chain = parse_reference_chain(expr)?; - match chain.root.as_str() { - "input" => self.compile_input_chain(&chain, span), - "data" => self.compile_data_chain(&chain, span), - _ => self.compile_local_var_chain(&chain, span), + match chain.root.clone() { + ReferenceRoot::Variable(name) => match name.as_str() { + "input" => self.compile_input_chain(&chain, span), + "data" => self.compile_data_chain(&chain, span), + _ => self.compile_local_var_chain(&name, &chain, span), + }, + ReferenceRoot::Expression(root_expr) => { + let root_reg = + self.compile_rego_expr_with_span(&root_expr, root_expr.span(), false)?; + self.compile_chain_access(root_reg, &chain.components, span) + } } } @@ -121,7 +144,9 @@ impl<'a> Compiler<'a> { } // Build the static prefix path components for rule matching - let static_prefix = chain.get_static_prefix(); + let static_prefix = chain + .get_static_prefix() + .expect("data references must have variable roots"); // Try to find the longest matching rule prefix // Start from the full path and work backwards @@ -252,9 +277,14 @@ impl<'a> Compiler<'a> { } /// Compile local variable access chain - fn compile_local_var_chain(&mut self, chain: &ReferenceChain, span: &Span) -> Result { + fn compile_local_var_chain( + &mut self, + root: &str, + chain: &ReferenceChain, + span: &Span, + ) -> Result { // Check if it's a local variable first (precedence over rules) - if let Some(var_reg) = self.lookup_variable(&chain.root) { + if let Some(var_reg) = self.lookup_variable(root) { if chain.components.is_empty() { return Ok(var_reg); } @@ -262,7 +292,7 @@ impl<'a> Compiler<'a> { } // Check if there's a rule in the current package that matches - let current_pkg_prefix = format!("{}.{}", &self.current_package, &chain.root); + let current_pkg_prefix = format!("{}.{}", &self.current_package, root); // Build static path for rule matching let mut rule_path_parts = vec![current_pkg_prefix.as_str()]; @@ -300,7 +330,7 @@ impl<'a> Compiler<'a> { // No rule found - undefined variable Err(CompilerError::UndefinedVariable { - name: chain.root.clone(), + name: root.to_string(), } .at(span)) } diff --git a/tests/opa.rs b/tests/opa.rs index d94e444..ac64f0e 100644 --- a/tests/opa.rs +++ b/tests/opa.rs @@ -32,15 +32,11 @@ const OPA_TODO_FOLDERS: &[&str] = &[ "example", "fix1863", "functions", - "jsonschema", "partialdocconstants", "partialobjectdoc", "planner-ir", "refheads", - "replacen", - "semverisvalid", "sets", - "time", "type", "varreferences", "virtualdocs", diff --git a/tests/rvm/rego/cases/chained_access.yaml b/tests/rvm/rego/cases/chained_access.yaml index d615f85..7b734be 100644 --- a/tests/rvm/rego/cases/chained_access.yaml +++ b/tests/rvm/rego/cases/chained_access.yaml @@ -196,6 +196,30 @@ cases: query: data.test.main want_result: "web1" + - note: literal_array_root_access + data: {} + modules: + - | + package test + + y := ["x", "y"][1] + + main := y + query: data.test.main + want_result: "y" + + - note: computed_object_root_access + data: {} + modules: + - | + package test + + x := strings.replace_n({k: v | k := ["f", "foo"][i]; v := ["x", "xxx"][i]}, "foo") + + main := x + query: data.test.main + want_result: "xoo" + - note: string_literal_bracket_access data: {} modules: From bedf667adc31a2b693d5fe6616c1ac0e2e65ddbc Mon Sep 17 00:00:00 2001 From: Anand Krishnamoorthi Date: Wed, 3 Dec 2025 12:59:58 -0600 Subject: [PATCH 3/3] feat: Handle literal comparisons that use = and comprehensions without loops - emit AssertCondition for equality-only assignment plans (outside soft-assert mode) so rules like `0 = 1` fail under the VM just like the interpreter - let comprehension bodies consume assertion failures by advancing or exiting their iteration context, both in run-to-completion and suspendable execution Signed-off-by: Anand Krishnamoorthi --- src/languages/rego/compiler/destructuring.rs | 3 + src/rvm/vm/comprehension.rs | 77 ++++++++++++++++++++ src/rvm/vm/loops.rs | 46 ++++++------ tests/opa.rs | 5 -- tests/rvm/rego/cases/comparisons.yaml | 12 +++ tests/rvm/rego/cases/comprehensions.yaml | 10 +++ 6 files changed, 127 insertions(+), 26 deletions(-) diff --git a/src/languages/rego/compiler/destructuring.rs b/src/languages/rego/compiler/destructuring.rs index 3cb8b4e..83e7162 100644 --- a/src/languages/rego/compiler/destructuring.rs +++ b/src/languages/rego/compiler/destructuring.rs @@ -99,6 +99,9 @@ impl<'a> Compiler<'a> { }, span, ); + if !self.soft_assert_mode { + self.emit_instruction(Instruction::AssertCondition { condition: dest }, span); + } Ok(dest) } AssignmentPlan::WildcardMatch { diff --git a/src/rvm/vm/comprehension.rs b/src/rvm/vm/comprehension.rs index 6358c0c..acaf6b7 100644 --- a/src/rvm/vm/comprehension.rs +++ b/src/rvm/vm/comprehension.rs @@ -457,6 +457,83 @@ impl RegoVM { } } + pub(super) fn handle_comprehension_condition_failure_run_to_completion( + &mut self, + ) -> Result { + if let Some(mut context) = self.comprehension_stack.pop() { + self.advance_comprehension_after_failure(&mut context)?; + self.comprehension_stack.push(context); + Ok(true) + } else { + Ok(false) + } + } + + pub(super) fn handle_comprehension_condition_failure_suspendable(&mut self) -> Result { + if let Some(mut frame) = self.execution_stack.pop() { + let handled = if let FrameKind::Comprehension { context, .. } = &mut frame.kind { + self.advance_comprehension_after_failure(context)?; + true + } else { + false + }; + + self.execution_stack.push(frame); + if handled { + return Ok(true); + } + } + Ok(false) + } + + fn advance_comprehension_after_failure( + &mut self, + context: &mut ComprehensionContext, + ) -> Result<()> { + if let Some(iter_state) = context.iteration_state.as_mut() { + self.capture_comprehension_iteration_position( + iter_state, + context.key_reg, + context.value_reg, + ); + iter_state.advance(); + let has_next = + self.setup_next_iteration(iter_state, context.key_reg, context.value_reg)?; + if has_next { + self.pc = context.body_start.saturating_sub(1) as usize; + } else { + context.iteration_state = None; + self.pc = context.comprehension_end.saturating_sub(1) as usize; + } + } else { + self.pc = context.comprehension_end.saturating_sub(1) as usize; + } + + Ok(()) + } + + fn capture_comprehension_iteration_position( + &mut self, + iter_state: &mut IterationState, + key_reg: u8, + value_reg: u8, + ) { + match iter_state { + IterationState::Object { current_key, .. } => { + let tracked_key = if key_reg != value_reg { + self.registers[key_reg as usize].clone() + } else { + self.registers[value_reg as usize].clone() + }; + *current_key = Some(tracked_key); + } + IterationState::Set { current_item, .. } => { + *current_item = Some(self.registers[value_reg as usize].clone()); + } + IterationState::Array { .. } => {} + } + } + fn execute_comprehension_end_run_to_completion(&mut self) -> Result<()> { if let Some(_context) = self.comprehension_stack.pop() { Ok(()) diff --git a/src/rvm/vm/loops.rs b/src/rvm/vm/loops.rs index 928037f..422ab86 100644 --- a/src/rvm/vm/loops.rs +++ b/src/rvm/vm/loops.rs @@ -621,6 +621,8 @@ impl RegoVM { self.pc = loop_next_pc as usize - 1; } } + } else if self.handle_comprehension_condition_failure_run_to_completion()? { + // handled by comprehension context } else { return Err(VmError::AssertionFailed); } @@ -633,29 +635,31 @@ impl RegoVM { return Ok(()); } - let (resume_pc, loop_ctx) = match self.execution_stack.last_mut() { - Some(ExecutionFrame { - kind: FrameKind::Loop { return_pc, context }, - .. - }) => (*return_pc, context), - _ => return Err(VmError::AssertionFailed), - }; - - match loop_ctx.mode { - LoopMode::Any | LoopMode::ForEach => { - loop_ctx.current_iteration_failed = true; - self.pc = loop_ctx.loop_next_pc as usize - 1; - } - LoopMode::Every => { - self.registers[loop_ctx.result_reg as usize] = Value::Bool(false); - let completed_frame = self.execution_stack.pop().expect("loop frame exists"); - if let Some(parent) = self.execution_stack.last_mut() { - parent.pc = resume_pc; + if let Some(ExecutionFrame { + kind: FrameKind::Loop { return_pc, context }, + .. + }) = self.execution_stack.last_mut() + { + let resume_pc = *return_pc; + match context.mode { + LoopMode::Any | LoopMode::ForEach => { + context.current_iteration_failed = true; + self.pc = context.loop_next_pc as usize - 1; + } + LoopMode::Every => { + self.registers[context.result_reg as usize] = Value::Bool(false); + let completed_frame = self.execution_stack.pop().expect("loop frame exists"); + if let Some(parent) = self.execution_stack.last_mut() { + parent.pc = resume_pc; + } + drop(completed_frame); } - drop(completed_frame); } + Ok(()) + } else if self.handle_comprehension_condition_failure_suspendable()? { + Ok(()) + } else { + Err(VmError::AssertionFailed) } - - Ok(()) } } diff --git a/tests/opa.rs b/tests/opa.rs index ac64f0e..262afad 100644 --- a/tests/opa.rs +++ b/tests/opa.rs @@ -22,14 +22,10 @@ const OPA_BRANCH: &str = "v1.2.0"; const OPA_TODO_FOLDERS: &[&str] = &[ "aggregates", "baseandvirtualdocs", - "comparisonexpr", "dataderef", "defaultkeyword", - "disjunction", "elsekeyword", - "eqexpr", "every", - "example", "fix1863", "functions", "partialdocconstants", @@ -38,7 +34,6 @@ const OPA_TODO_FOLDERS: &[&str] = &[ "refheads", "sets", "type", - "varreferences", "virtualdocs", "walkbuiltin", "withkeyword", diff --git a/tests/rvm/rego/cases/comparisons.yaml b/tests/rvm/rego/cases/comparisons.yaml index 99b6e92..3f88bbb 100644 --- a/tests/rvm/rego/cases/comparisons.yaml +++ b/tests/rvm/rego/cases/comparisons.yaml @@ -48,3 +48,15 @@ cases: } query: data.test.main want_result: true + + - note: equality_literal_failure + data: {} + modules: + - | + package test + + x if { + 0 = 1 + } + query: data.test.x + want_result: "#undefined" diff --git a/tests/rvm/rego/cases/comprehensions.yaml b/tests/rvm/rego/cases/comprehensions.yaml index a3851c1..4152c37 100644 --- a/tests/rvm/rego/cases/comprehensions.yaml +++ b/tests/rvm/rego/cases/comprehensions.yaml @@ -13,3 +13,13 @@ cases: main := [(x * 2) | some x in [1, 2, 3]] query: data.test.main want_result: [2, 4, 6] + + - note: comprehension_equality_failure_returns_empty + data: {} + modules: + - | + package test + + z := [x | x := 1; x == 2] + query: data.test.z + want_result: []