chore(rvm): add debug-mode invariant assertions (microsoft#737)

Encode VM stack/context/register lifecycle invariants as
debug_assert!s. Zero cost in release; surfaces violations during
debug-mode tests and CI.

Invariants covered:
- reset_execution_state postcondition: all stacks empty, registers
  resized to base and Undefined, rule_cache reset, pc/executed
  counters zeroed, builtins_cache cleared, execution_state Ready.
- Per-opcode invariant check (assert_vm_invariants) invoked at the
  top of run_stackless_loop and jump_to iterations: state is
  Ready/Running, registers non-empty, rule_cache sized to program,
  execution stack bounded by a debug-only sanity ceiling
  (DEBUG_MAX_EXECUTION_STACK_DEPTH = 4096; not a production limit).
- resume() precondition: execution_state is Suspended.
- execute_suspendable_entry precondition: clean state (callers reset
  immediately before).
- Rule finalize: call_rule_stack pop matches the finalized rule_index.
- IterationState::advance: Single iterator not advanced past
  consumption, Array index not at usize::MAX before saturating_add.

All assertions are gated by #[cfg(debug_assertions)] (directly or via
debug_assert!) so release builds are unaffected.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: anakrish

src/rvm/vm/comprehension.rs

@@ -528,7 +528,6 @@ impl RegoVM {
             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 &mut FrameKind::Comprehension {
@@ -606,6 +605,9 @@ impl RegoVM {
     }
 
     fn execute_comprehension_end_run_to_completion(&mut self) -> Result<()> {
+        // `ComprehensionEnd` is reached from a loaded program; an empty stack
+        // here means malformed user-supplied bytecode, which must still surface
+        // as a typed error rather than a panic — including in debug builds.
         self.comprehension_stack.pop().map_or_else(
             || {
                 Err(VmError::InvalidIteration {

src/rvm/vm/context.rs

@@ -54,6 +54,14 @@ impl IterationState {
     pub(super) const fn advance(&mut self) {
         match *self {
             Self::Array { ref mut index, .. } => {
+                // Array iteration uses `usize` as the cursor and advances via
+                // `saturating_add(1)`. A cursor already at `usize::MAX` here
+                // means a stuck (non-progressing) iteration was emitted by
+                // malformed bytecode; assert in debug to surface it loudly.
+                debug_assert!(
+                    *index < usize::MAX,
+                    "IterationState::Array index already at usize::MAX on advance"
+                );
                 *index = index.saturating_add(1);
             }
             Self::Object {
@@ -69,6 +77,12 @@ impl IterationState {
             Self::Single {
                 ref mut consumed, ..
             } => {
+                // `Single` yields exactly once; advancing a consumed Single
+                // means the compiler emitted a redundant LoopNext.
+                debug_assert!(
+                    !*consumed,
+                    "IterationState::Single advanced after consumption"
+                );
                 *consumed = true;
             }
         }

src/rvm/vm/errors.rs

@@ -295,6 +295,13 @@ pub enum VmError {
     #[error("Call rule stack underflow during rule finalization (pc={pc})")]
     CallRuleStackUnderflow { pc: usize },
 
+    #[error("Call rule stack mismatch during rule finalization: expected rule_index {expected}, popped {actual} (pc={pc})")]
+    CallRuleStackMismatch {
+        expected: u16,
+        actual: u16,
+        pc: usize,
+    },
+
     #[error("Internal VM error: {message} (pc={pc})")]
     Internal { message: String, pc: usize },
 }

src/rvm/vm/execution.rs

@@ -117,6 +117,10 @@ impl RegoVM {
         let target = self.convert_pc(target, "jump target")?;
         self.pc = target;
         while self.pc < program.instructions.len() {
+            // Per-instruction sanity check: every iteration of the dispatch
+            // loop must re-enter with the VM in a Running/Ready state and the
+            // working data structures coherent.
+            self.assert_vm_invariants();
             self.memory_check()?;
             if self.executed_instructions >= self.max_instructions {
                 return Err(VmError::InstructionLimitExceeded {
@@ -189,6 +193,9 @@ impl RegoVM {
     }
 
     fn execute_suspendable_entry(&mut self, entry_point_pc: usize) -> Result<Value> {
+        // Precondition: callers (execute_entry_point_by_{index,name}) reset the
+        // VM before invoking this method, so the VM must be in a clean state.
+        self.debug_assert_state_is_clean();
         self.execution_state = ExecutionState::Running;
         self.reset_execution_timer_state();
         match self.run_stackless_from(entry_point_pc) {
@@ -201,6 +208,10 @@ impl RegoVM {
     }
 
     pub fn resume(&mut self, resume_value: Option<Value>) -> Result<Value> {
+        // Precondition is enforced below by returning `VmError::InvalidResumeState`
+        // for any non-`Suspended` state. A `debug_assert!` here would diverge
+        // debug vs release behavior and, when invoked via FFI, would trip the
+        // unwind guard and poison the engine on a recoverable misuse.
         let (reason, mut last_result) = match self.execution_state.clone() {
             ExecutionState::Suspended {
                 reason,
@@ -289,6 +300,9 @@ impl RegoVM {
 
     fn run_stackless_loop(&mut self, program: &Program, last_result: &mut Value) -> Result<()> {
         while !self.execution_stack.is_empty() {
+            // Per-instruction sanity check: see `assert_vm_invariants` for the
+            // exact contract. Compiled out in release.
+            self.assert_vm_invariants();
             self.memory_check()?;
             self.frame_pc_overridden = false;
             let should_finalize_rule = self.execution_stack.last().is_some_and(|frame| {

src/rvm/vm/rules.rs

@@ -277,6 +277,19 @@ impl RegoVM {
             .call_rule_stack
             .pop()
             .ok_or(VmError::CallRuleStackUnderflow { pc: self.pc })?;
+        // Stack discipline: the context we just popped must belong to the
+        // rule we are finalizing. A mismatch indicates a missing push or an
+        // extra pop somewhere in this rule's execution and would otherwise
+        // silently restore the wrong return_pc / rule_type. Surface as a
+        // typed VmError so the contract holds the same in debug and release
+        // builds (avoiding FFI poisoning via a debug-only panic).
+        if rule_index != call_context.rule_index {
+            return Err(VmError::CallRuleStackMismatch {
+                expected: rule_index,
+                actual: call_context.rule_index,
+                pc: self.pc,
+            });
+        }
         self.pc = call_context.return_pc;
 
         let result_from_rule = if !rule_failed_due_to_inconsistency {
@@ -831,6 +844,9 @@ impl RegoVM {
 
         self.registers = parent_registers;
 
+        // Underflow here means malformed/poisoned program state; surface as a
+        // typed error rather than a debug-only panic so the public load_program
+        // contract holds the same in debug and release.
         if self.call_rule_stack.pop().is_none() {
             return Err(VmError::CallRuleStackUnderflow { pc: self.pc });
         }

src/rvm/vm/state.rs

@@ -34,6 +34,139 @@ impl RegoVM {
 
         // Builtin cache entries only live for a single execution
         self.builtins_cache.clear();
+
+        // Postcondition: every stack/cache that `reset_execution_state` touches
+        // must be in its documented "clean" shape. This catches accidental
+        // omissions in future edits to this function.
+        self.debug_assert_state_is_clean();
+    }
+
+    /// Debug-only postcondition for `reset_execution_state`.
+    ///
+    /// Asserts the invariants every caller of `reset_execution_state` relies on
+    /// before starting a fresh execution. The body is fully gated by
+    /// `#[cfg(debug_assertions)]` so this is a zero-cost no-op in release.
+    #[inline]
+    pub(super) fn debug_assert_state_is_clean(&self) {
+        #[cfg(debug_assertions)]
+        {
+            // --- Stacks: every per-execution stack must be drained. ---
+            debug_assert!(
+                self.execution_stack.is_empty(),
+                "reset_execution_state postcondition: execution_stack must be empty"
+            );
+            debug_assert!(
+                self.loop_stack.is_empty(),
+                "reset_execution_state postcondition: loop_stack must be empty"
+            );
+            debug_assert!(
+                self.comprehension_stack.is_empty(),
+                "reset_execution_state postcondition: comprehension_stack must be empty"
+            );
+            debug_assert!(
+                self.call_rule_stack.is_empty(),
+                "reset_execution_state postcondition: call_rule_stack must be empty"
+            );
+            debug_assert!(
+                self.register_stack.is_empty(),
+                "reset_execution_state postcondition: register_stack must be empty"
+            );
+
+            // --- Caches: cleared so a new program/input cannot read stale entries. ---
+            debug_assert!(
+                self.builtins_cache.is_empty(),
+                "reset_execution_state postcondition: builtins_cache must be empty"
+            );
+
+            // --- Registers: window resized to the program's base count and zeroed. ---
+            debug_assert_eq!(
+                self.registers.len(),
+                self.base_register_count,
+                "reset_execution_state postcondition: registers must be sized to base_register_count"
+            );
+            debug_assert!(
+                self.registers.iter().all(|v| matches!(v, Value::Undefined)),
+                "reset_execution_state postcondition: all registers must be Undefined"
+            );
+
+            // --- Rule cache: sized to the current program and marked uncomputed. ---
+            debug_assert_eq!(
+                self.rule_cache.len(),
+                self.program.rule_infos.len(),
+                "reset_execution_state postcondition: rule_cache size must match program rule_infos"
+            );
+            debug_assert!(
+                self.rule_cache.iter().all(|entry| !entry.0),
+                "reset_execution_state postcondition: rule_cache entries must be uncomputed"
+            );
+
+            // --- Counters and execution-state machine: zeroed and back to Ready. ---
+            debug_assert_eq!(
+                self.pc, 0,
+                "reset_execution_state postcondition: pc must be 0"
+            );
+            debug_assert_eq!(
+                self.executed_instructions, 0,
+                "reset_execution_state postcondition: executed_instructions must be 0"
+            );
+            debug_assert!(
+                matches!(self.execution_state, ExecutionState::Ready),
+                "reset_execution_state postcondition: execution_state must be Ready"
+            );
+        }
+    }
+
+    /// Per-opcode VM invariants checked from the inner dispatch loop.
+    ///
+    /// These hold every time control re-enters the dispatch loop with another
+    /// instruction to execute. Only conditions that are *purely VM-internal*
+    /// (i.e. cannot be made false by any host-supplied program or out-of-order
+    /// API call) are asserted here — anything reachable from `load_program`
+    /// input must surface as a typed `VmError` instead, to avoid panicking in
+    /// debug builds and poisoning the engine across FFI.
+    ///
+    /// Fully `#[cfg(debug_assertions)]`-gated so the method body compiles out
+    /// in release.
+    #[inline]
+    pub(super) fn assert_vm_invariants(&self) {
+        #[cfg(debug_assertions)]
+        {
+            // The dispatch loop only runs while execution is live. Once the VM
+            // has transitioned to a terminal state (Suspended/Completed/Error)
+            // the loop must have exited. Note `Ready` is also valid here because
+            // some entry points (e.g. `execute_entry_point_by_index` in
+            // RunToCompletion mode) drive `jump_to` without flipping the state.
+            // `execution_state` is mutated only inside the VM and is not
+            // host-controllable.
+            debug_assert!(
+                matches!(
+                    self.execution_state,
+                    ExecutionState::Ready | ExecutionState::Running
+                ),
+                "vm invariant: execution_state must be Ready or Running inside the dispatch loop, was {:?}",
+                self.execution_state
+            );
+
+            // Rule cache is sized once at reset (against the currently loaded
+            // program) and the VM does not resize it mid-execution. Any
+            // mismatch here would indicate an internal accounting bug rather
+            // than malformed input.
+            debug_assert_eq!(
+                self.rule_cache.len(),
+                self.program.rule_infos.len(),
+                "vm invariant: rule_cache size must equal program.rule_infos size"
+            );
+
+            // NOTE: `!registers.is_empty()` and an `execution_stack` depth
+            // ceiling were intentionally *not* asserted here: both can be
+            // triggered by a host-loaded program (registers via
+            // `RuleInfo::num_registers == 0`; stack depth via deeply nested
+            // rules/loops/comprehensions) and would therefore panic in debug
+            // and poison the engine across FFI. Register access is already
+            // guarded by `VmError::RegisterIndexOutOfBounds`; runaway recursion
+            // is bounded in production by `set_max_instructions` and
+            // `memory_check`.
+        }
     }
 
     /// Return all active objects to their respective pools for reuse