refactor(tvix/eval): Compile OpAssert using conditional jumps r/5082

In order to behave nicely with tryEval, asserts need to leave the
instruction pointer in a reasonable place even if they fail - whereas
with the previous implementation catching a failed assert would still
end up running the op for the *body* of the assert. With this change, we
compile asserts much more like an `if` expression with conditional jumps
rather than having an OpAssert op.

Change-Id: I1b266c3be90185c84000da6b1995ac3e6fd5471b
Reviewed-on: https://cl.tvl.fyi/c/depot/+/6925
Reviewed-by: tazjin <tazjin@tvl.su>
Tested-by: BuildkiteCI

3 files changed, 26 insertions, 10 deletions

diff --git a/tvix/eval/src/compiler/mod.rs b/tvix/eval/src/compiler/mod.rs
index 1463bb100eae..941b945a8bb6 100644
--- a/tvix/eval/src/compiler/mod.rs
+++ b/tvix/eval/src/compiler/mod.rs
@@ -630,16 +630,34 @@ impl Compiler<'_> {
         self.patch_jump(final_jump);
     }
 
+    /// Compile `assert` expressions using jumping instructions in the VM.
+    ///
+    /// ```notrust
+    ///                        ┌─────────────────────┐
+    ///                        │ 0  [ conditional ]  │
+    ///                        │ 1   JUMP_IF_FALSE  →┼─┐
+    ///                        │ 2  [  main body  ]  │ │ Jump to else body if
+    ///                       ┌┼─3─←     JUMP        │ │ condition is false.
+    ///  Jump over else body  ││ 4   OP_ASSERT_FAIL ←┼─┘
+    ///  if condition is true.└┼─5─→     ...         │
+    ///                        └─────────────────────┘
+    /// ```
     fn compile_assert(&mut self, slot: LocalIdx, node: &ast::Assert) {
         // Compile the assertion condition to leave its value on the stack.
         self.compile(slot, &node.condition().unwrap());
         self.emit_force(&node.condition().unwrap());
-        self.push_op(OpCode::OpAssert, &node.condition().unwrap());
+        let then_idx = self.push_op(OpCode::OpJumpIfFalse(JumpOffset(0)), node);
 
-        // The runtime will abort evaluation at this point if the
-        // assertion failed, if not the body simply continues on like
-        // normal.
+        self.push_op(OpCode::OpPop, node);
         self.compile(slot, &node.body().unwrap());
+
+        let else_idx = self.push_op(OpCode::OpJump(JumpOffset(0)), node);
+
+        self.patch_jump(then_idx);
+        self.push_op(OpCode::OpPop, node);
+        self.push_op(OpCode::OpAssertFail, &node.condition().unwrap());
+
+        self.patch_jump(else_idx);
     }
 
     /// Compile conditional expressions using jumping instructions in the VM.
diff --git a/tvix/eval/src/opcode.rs b/tvix/eval/src/opcode.rs
index c2eabba1c28e..15f60a538c41 100644
--- a/tvix/eval/src/opcode.rs
+++ b/tvix/eval/src/opcode.rs
@@ -121,8 +121,8 @@ pub enum OpCode {
     /// Close scopes while leaving their expression value around.
     OpCloseScope(Count), // number of locals to pop
 
-    /// Asserts stack top is a boolean, and true.
-    OpAssert,
+    /// Return an error indicating that an `assert` failed
+    OpAssertFail,
 
     // Lambdas & closures
     OpCall,
diff --git a/tvix/eval/src/vm.rs b/tvix/eval/src/vm.rs
index b89b61f7626a..01299ccc17da 100644
--- a/tvix/eval/src/vm.rs
+++ b/tvix/eval/src/vm.rs
@@ -561,10 +561,8 @@ impl<'o> VM<'o> {
                     }
                 }
 
-                OpCode::OpAssert => {
-                    if !fallible!(self, self.pop().as_bool()) {
-                        return Err(self.error(ErrorKind::AssertionFailed));
-                    }
+                OpCode::OpAssertFail => {
+                    return Err(self.error(ErrorKind::AssertionFailed));
                 }
 
                 OpCode::OpCall => {