Merge pull request microsoft#516 from anakrish/rvm-opa-2

anakrish · web-flow · commit 252ae0e31249 · 2025-12-16T11:28:36.000-06:00
Handle more OPA semantics in RVM and compiler
diff --git 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/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<AccessComponent>,
 }
 
 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<Vec<&str>> {
+        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<ReferenceChain> {
         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<ReferenceChain> {
                 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,17 +277,22 @@ impl<'a> Compiler<'a> {
     }
 
     /// Compile local variable access chain
-    fn compile_local_var_chain(&mut self, chain: &ReferenceChain, span: &Span) -> Result<Register> {
+    fn compile_local_var_chain(
+        &mut self,
+        root: &str,
+        chain: &ReferenceChain,
+        span: &Span,
+    ) -> Result<Register> {
         // 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);
             }
             return self.compile_chain_access(var_reg, &chain.components, span);
         }
 
         // 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/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<bool> {
+        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<bool> {
+        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/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/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/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>), 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
@@ -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
@@ -22,28 +22,18 @@ const OPA_BRANCH: &str = "v1.2.0";
 const OPA_TODO_FOLDERS: &[&str] = &[
     "aggregates",
     "baseandvirtualdocs",
-    "comparisonexpr",
     "dataderef",
     "defaultkeyword",
-    "disjunction",
     "elsekeyword",
-    "eqexpr",
     "every",
-    "example",
     "fix1863",
     "functions",
-    "jsonschema",
     "partialdocconstants",
     "partialobjectdoc",
     "planner-ir",
-    "rand",
     "refheads",
-    "replacen",
-    "semverisvalid",
     "sets",
-    "time",
     "type",
-    "varreferences",
     "virtualdocs",
     "walkbuiltin",
     "withkeyword",
diff --git a/tests/rvm/rego/cases/builtins_cache.yaml b/tests/rvm/rego/cases/builtins_cache.yaml
diff --git a/tests/rvm/rego/cases/chained_access.yaml b/tests/rvm/rego/cases/chained_access.yaml
diff --git a/tests/rvm/rego/cases/comparisons.yaml b/tests/rvm/rego/cases/comparisons.yaml
diff --git a/tests/rvm/rego/cases/comprehensions.yaml b/tests/rvm/rego/cases/comprehensions.yaml

Original file line number	Diff line number	Diff line change
`@@ -98,6 +98,9 @@ pub struct RegoVM {`
`98`	`98`
`99`	`99`	`/// Whether builtins should raise errors strictly or return undefined on failure`
`100`	`100`	`pub(super) strict_builtin_errors: bool,`
	`101`	`+`
	`102`	`+ /// Cache for builtin calls that must stay deterministic across a single evaluation`
	`103`	`+ pub(super) builtins_cache: BTreeMap<(&'static str, Vec<Value>), Value>,`
`101`	`104`	`}`
`102`	`105`
`103`	`106`	`impl Default for RegoVM {`
`@@ -135,6 +138,7 @@ impl RegoVM {`
`135`	`138`	`execution_mode: ExecutionMode::RunToCompletion,`
`136`	`139`	`frame_pc_overridden: false,`
`137`	`140`	`strict_builtin_errors: false,`
	`141`	`+ builtins_cache: BTreeMap::new(),`
`138`	`142`	`}`
`139`	`143`	`}`
`140`	`144`
Original file line number	Diff line number	Diff line change
`@@ -32,6 +32,9 @@ impl RegoVM {`
`32`	`32`	`self.registers.clear();`
`33`	`33`	`self.registers`
`34`	`34`	`.resize(self.base_register_count, Value::Undefined);`
	`35`	`+`
	`36`	`+ // Builtin cache entries only live for a single execution`
	`37`	`+ self.builtins_cache.clear();`
`35`	`38`	`}`
`36`	`39`
`37`	`40`	`/// Return all active objects to their respective pools for reuse`