feat(tvix/eval): implement if/else expressions r/4483

These expressions use simple jumps to skip the correct expression
conditionally in the bytecode by advancing the instruction pointer.

Note that these expressions are already covered by a test behind the
`nix_tests` feature flag, but adding more is probably sensible.

Change-Id: Ibe0eba95d216321c883d3b6b5816e2ab6fe7eef1
Reviewed-on: https://cl.tvl.fyi/c/depot/+/6148
Tested-by: BuildkiteCI
Reviewed-by: grfn <grfn@gws.fyi>
Reviewed-by: sterni <sternenseemann@systemli.org>

4 files changed, 84 insertions, 5 deletions

diff --git a/tvix/eval/src/compiler.rs b/tvix/eval/src/compiler.rs
index 4a3a652af5..3df0ea4c98 100644
--- a/tvix/eval/src/compiler.rs
+++ b/tvix/eval/src/compiler.rs
@@ -3,7 +3,7 @@
 
 use crate::chunk::Chunk;
 use crate::errors::EvalResult;
-use crate::opcode::OpCode;
+use crate::opcode::{CodeIdx, OpCode};
 use crate::value::Value;
 
 use rnix;
@@ -68,6 +68,11 @@ impl Compiler {
                 self.compile_list(node)
             }
 
+            rnix::SyntaxKind::NODE_IF_ELSE => {
+                let node = rnix::types::IfElse::cast(node).unwrap();
+                self.compile_if_else(node)
+            }
+
             kind => {
                 println!("visiting unsupported node: {:?}", kind);
                 Ok(())
@@ -143,12 +148,16 @@ impl Compiler {
             BinOpKind::MoreOrEq => self.chunk.add_op(OpCode::OpMoreOrEq),
             BinOpKind::Concat => self.chunk.add_op(OpCode::OpConcat),
 
+            BinOpKind::And => todo!(),
+            BinOpKind::Or => todo!(),
+            BinOpKind::Implication => todo!(),
+
             BinOpKind::NotEqual => {
                 self.chunk.add_op(OpCode::OpEqual);
                 self.chunk.add_op(OpCode::OpInvert)
             }
 
-            _ => todo!(),
+            BinOpKind::IsSet => todo!("? operator"),
         };
 
         Ok(())
@@ -269,6 +278,51 @@ impl Compiler {
         self.chunk.add_op(OpCode::OpList(count));
         Ok(())
     }
+
+    // Compile conditional expressions using jumping instructions in the VM.
+    //
+    //                        ┌────────────────────┐
+    //                        │ 0  [ conditional ] │
+    //                        │ 1   JUMP_IF_FALSE →┼─┐
+    //                        │ 2  [  main body  ] │ │ Jump to else body if
+    //                       ┌┼─3─←     JUMP       │ │ condition is false.
+    //  Jump over else body  ││ 4  [  else body  ]←┼─┘
+    //  if condition is true.└┼─5─→     ...        │
+    //                        └────────────────────┘
+    fn compile_if_else(&mut self, node: rnix::types::IfElse) -> EvalResult<()> {
+        self.compile(node.condition().unwrap())?;
+
+        let then_idx = self.chunk.add_op(OpCode::OpJumpIfFalse(0));
+
+        self.chunk.add_op(OpCode::OpPop); // discard condition value
+        self.compile(node.body().unwrap())?;
+
+        let else_idx = self.chunk.add_op(OpCode::OpJump(0));
+
+        self.patch_jump(then_idx); // patch jump *to* else_body
+        self.chunk.add_op(OpCode::OpPop); // discard condition value
+        self.compile(node.else_body().unwrap())?;
+
+        self.patch_jump(else_idx); // patch jump *over* else body
+
+        Ok(())
+    }
+
+    fn patch_jump(&mut self, idx: CodeIdx) {
+        let offset = self.chunk.code.len() - 1 - idx.0;
+
+        match &mut self.chunk.code[idx.0] {
+            OpCode::OpJump(n) => {
+                *n = offset;
+            }
+
+            OpCode::OpJumpIfFalse(n) => {
+                *n = offset;
+            }
+
+            op => panic!("attempted to patch unsupported op: {:?}", op),
+        }
+    }
 }
 
 pub fn compile(ast: rnix::AST) -> EvalResult<Chunk> {
diff --git a/tvix/eval/src/opcode.rs b/tvix/eval/src/opcode.rs
index b1e6d32593..8bb8fe9a1f 100644
--- a/tvix/eval/src/opcode.rs
+++ b/tvix/eval/src/opcode.rs
@@ -12,6 +12,9 @@ pub enum OpCode {
     // Push a constant onto the stack.
     OpConstant(ConstantIdx),
 
+    // Discard a value from the stack.
+    OpPop,
+
     // Push a literal value.
     OpNull,
     OpTrue,
@@ -27,13 +30,17 @@ pub enum OpCode {
     OpMul,
     OpDiv,
 
-    // Logical binary operators
+    // Comparison operators
     OpEqual,
     OpLess,
     OpLessOrEq,
     OpMore,
     OpMoreOrEq,
 
+    // Logical operators & generic jumps
+    OpJump(usize),
+    OpJumpIfFalse(usize),
+
     // Attribute sets
     OpAttrs(usize),
     OpAttrPath(usize),
diff --git a/tvix/eval/src/value/mod.rs b/tvix/eval/src/value/mod.rs
index 0c9042dbfe..4a34379748 100644
--- a/tvix/eval/src/value/mod.rs
+++ b/tvix/eval/src/value/mod.rs
@@ -52,9 +52,9 @@ impl Value {
         }
     }
 
-    pub fn as_bool(self) -> EvalResult<bool> {
+    pub fn as_bool(&self) -> EvalResult<bool> {
         match self {
-            Value::Bool(b) => Ok(b),
+            Value::Bool(b) => Ok(*b),
             other => Err(Error::TypeError {
                 expected: "bool",
                 actual: other.type_of(),
diff --git a/tvix/eval/src/vm.rs b/tvix/eval/src/vm.rs
index a64e0b337a..4ea1b3b65c 100644
--- a/tvix/eval/src/vm.rs
+++ b/tvix/eval/src/vm.rs
@@ -83,6 +83,10 @@ impl VM {
         self.stack.push(value)
     }
 
+    fn peek(&self, offset: usize) -> &Value {
+        &self.stack[self.stack.len() - 1 - offset]
+    }
+
     fn run(&mut self) -> EvalResult<Value> {
         loop {
             match self.inc_ip() {
@@ -91,6 +95,10 @@ impl VM {
                     self.push(c);
                 }
 
+                OpCode::OpPop => {
+                    self.pop();
+                }
+
                 OpCode::OpAdd => {
                     let b = self.pop();
                     let a = self.pop();
@@ -163,6 +171,16 @@ impl VM {
                 }
 
                 OpCode::OpInterpolate(count) => self.run_interpolate(count)?,
+
+                OpCode::OpJump(offset) => {
+                    self.ip += offset;
+                }
+
+                OpCode::OpJumpIfFalse(offset) => {
+                    if !self.peek(0).as_bool()? {
+                        self.ip += offset;
+                    }
+                }
             }
 
             if self.ip == self.chunk.code.len() {