feat(tvix/eval): implement simple arithmetic binary operations r/4410

Implements simple arithmetic operations (+, -, *, /).

There is some scaffolding included to pop and coerce pairs of numbers,
as the Nix language will let arithmetic operators apply to arbitrary
pairs of number types (always resulting in floats if the types are
mixed).

Change-Id: I5f62c363bdea8baa6ef812cc64c5406759d257cf
Reviewed-on: https://cl.tvl.fyi/c/depot/+/6074
Tested-by: BuildkiteCI
Reviewed-by: grfn <grfn@gws.fyi>

5 files changed, 110 insertions, 5 deletions

diff --git a/tvix/eval/src/compiler.rs b/tvix/eval/src/compiler.rs
index 4a23db3dcb2b..65b97cd56ccc 100644
--- a/tvix/eval/src/compiler.rs
+++ b/tvix/eval/src/compiler.rs
@@ -6,17 +6,18 @@ use crate::errors::EvalResult;
 use crate::opcode::OpCode;
 use crate::value::Value;
 use rnix;
+use rnix::types::TypedNode;
 
 struct Compiler {
     chunk: Chunk,
 }
 
 impl Compiler {
-    fn compile(&mut self, node: &rnix::SyntaxNode) -> EvalResult<()> {
+    fn compile(&mut self, node: rnix::SyntaxNode) -> EvalResult<()> {
         match node.kind() {
             // Root of a file contains no content, it's just a marker
             // type.
-            rnix::SyntaxKind::NODE_ROOT => self.compile(&node.first_child().expect("TODO")),
+            rnix::SyntaxKind::NODE_ROOT => self.compile(node.first_child().expect("TODO")),
 
             // Literals contain a single token comprising of the
             // literal itself.
@@ -25,6 +26,11 @@ impl Compiler {
                 self.compile_literal(token)
             }
 
+            rnix::SyntaxKind::NODE_BIN_OP => {
+                let op = rnix::types::BinOp::cast(node).expect("TODO (should not be possible)");
+                self.compile_binop(op)
+            }
+
             kind => {
                 println!("visiting unsupported node: {:?}", kind);
                 Ok(())
@@ -52,6 +58,25 @@ impl Compiler {
             rnix::NixValue::Path(_, _) => todo!(),
         }
     }
+
+    fn compile_binop(&mut self, op: rnix::types::BinOp) -> EvalResult<()> {
+        self.compile(op.lhs().unwrap())?;
+        self.compile(op.rhs().unwrap())?;
+
+        use rnix::types::BinOpKind;
+
+        let opcode = match op.operator().unwrap() {
+            BinOpKind::Add => OpCode::OpAdd,
+            BinOpKind::Sub => OpCode::OpSub,
+            BinOpKind::Mul => OpCode::OpMul,
+            BinOpKind::Div => OpCode::OpDiv,
+
+            _ => todo!(),
+        };
+
+        self.chunk.add_op(opcode);
+        Ok(())
+    }
 }
 
 pub fn compile(ast: rnix::AST) -> EvalResult<Chunk> {
@@ -59,7 +84,7 @@ pub fn compile(ast: rnix::AST) -> EvalResult<Chunk> {
         chunk: Chunk::default(),
     };
 
-    c.compile(&ast.node())?;
+    c.compile(ast.node())?;
 
     Ok(c.chunk)
 }
diff --git a/tvix/eval/src/eval.rs b/tvix/eval/src/eval.rs
index 0c4900e7e59d..efa016ca9846 100644
--- a/tvix/eval/src/eval.rs
+++ b/tvix/eval/src/eval.rs
@@ -12,7 +12,7 @@ pub fn interpret(code: String) -> EvalResult<String> {
     }
 
     let mut out = String::new();
-    writeln!(out, "{}", ast.root().dump()).ok();
+    println!("{}", ast.root().dump());
 
     let code = crate::compiler::compile(ast)?;
     writeln!(out, "code: {:?}", code).ok();
diff --git a/tvix/eval/src/opcode.rs b/tvix/eval/src/opcode.rs
index b43f75efaa8f..307b695f6d9b 100644
--- a/tvix/eval/src/opcode.rs
+++ b/tvix/eval/src/opcode.rs
@@ -16,4 +16,10 @@ pub enum OpCode {
     OpNull,
     OpTrue,
     OpFalse,
+
+    // Simple binary operators
+    OpAdd,
+    OpSub,
+    OpMul,
+    OpDiv,
 }
diff --git a/tvix/eval/src/value.rs b/tvix/eval/src/value.rs
index 037284b144b2..cbfff55b49e1 100644
--- a/tvix/eval/src/value.rs
+++ b/tvix/eval/src/value.rs
@@ -8,3 +8,28 @@ pub enum Value {
     Integer(i64),
     Float(f64),
 }
+
+impl Value {
+    pub fn is_number(&self) -> bool {
+        match self {
+            Value::Integer(_) => true,
+            Value::Float(_) => true,
+            _ => false,
+        }
+    }
+
+    pub fn type_of(&self) -> &'static str {
+        match self {
+            Value::Null => "null",
+            Value::Bool(_) => "bool",
+            Value::Integer(_) => "int",
+            Value::Float(_) => "float",
+        }
+    }
+}
+
+#[derive(Clone, Copy, Debug, PartialEq)]
+pub enum NumberPair {
+    Floats(f64, f64),
+    Integer(i64, i64),
+}
diff --git a/tvix/eval/src/vm.rs b/tvix/eval/src/vm.rs
index e35cd99b1e3d..077acfcc1b8d 100644
--- a/tvix/eval/src/vm.rs
+++ b/tvix/eval/src/vm.rs
@@ -1,7 +1,12 @@
 //! This module implements the virtual (or abstract) machine that runs
 //! Tvix bytecode.
 
-use crate::{chunk::Chunk, errors::EvalResult, opcode::OpCode, value::Value};
+use crate::{
+    chunk::Chunk,
+    errors::{Error, EvalResult},
+    opcode::OpCode,
+    value::{NumberPair, Value},
+};
 
 pub struct VM {
     ip: usize,
@@ -18,6 +23,30 @@ impl VM {
         self.stack.pop().expect("TODO")
     }
 
+    fn pop_number_pair(&mut self) -> EvalResult<NumberPair> {
+        let v2 = self.pop();
+        let v1 = self.pop();
+
+        match (v1, v2) {
+            (Value::Integer(i1), Value::Integer(i2)) => Ok(NumberPair::Integer(i1, i2)),
+
+            (Value::Float(f1), Value::Float(f2)) => Ok(NumberPair::Floats(f1, f2)),
+
+            (Value::Integer(i1), Value::Float(f2)) => Ok(NumberPair::Floats(i1 as f64, f2)),
+
+            (Value::Float(f1), Value::Integer(i2)) => Ok(NumberPair::Floats(f1, i2 as f64)),
+
+            _ => Err(Error::TypeError {
+                expected: "number (either int or float)",
+                actual: if v1.is_number() {
+                    v2.type_of()
+                } else {
+                    v1.type_of()
+                },
+            }),
+        }
+    }
+
     fn inc_ip(&mut self) -> OpCode {
         let op = self.chunk.code[self.ip];
         self.ip += 1;
@@ -32,6 +61,26 @@ impl VM {
                     self.push(c);
                 }
 
+                OpCode::OpAdd => match self.pop_number_pair()? {
+                    NumberPair::Floats(f1, f2) => self.push(Value::Float(f1 + f2)),
+                    NumberPair::Integer(i1, i2) => self.push(Value::Integer(i1 + i2)),
+                },
+
+                OpCode::OpSub => match self.pop_number_pair()? {
+                    NumberPair::Floats(f1, f2) => self.push(Value::Float(f1 - f2)),
+                    NumberPair::Integer(i1, i2) => self.push(Value::Integer(i1 - i2)),
+                },
+
+                OpCode::OpMul => match self.pop_number_pair()? {
+                    NumberPair::Floats(f1, f2) => self.push(Value::Float(f1 * f2)),
+                    NumberPair::Integer(i1, i2) => self.push(Value::Integer(i1 * i2)),
+                },
+
+                OpCode::OpDiv => match self.pop_number_pair()? {
+                    NumberPair::Floats(f1, f2) => self.push(Value::Float(f1 / f2)),
+                    NumberPair::Integer(i1, i2) => self.push(Value::Integer(i1 / i2)),
+                },
+
                 OpCode::OpNull => todo!(),
                 OpCode::OpTrue => todo!(),
                 OpCode::OpFalse => todo!(),