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: []