refactor(tvix/eval): use CodeIdx wrapper for instruction pointer r/4838

As suggested by sterni in cl/6453.

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

3 files changed, 34 insertions, 9 deletions

diff --git a/tvix/eval/src/observer.rs b/tvix/eval/src/observer.rs
index d662ea8a3867..e5562e3697d3 100644
--- a/tvix/eval/src/observer.rs
+++ b/tvix/eval/src/observer.rs
@@ -53,7 +53,7 @@ pub trait Observer {
 
     /// Called when the runtime *begins* executing an instruction. The
     /// provided stack is the state at the beginning of the operation.
-    fn observe_execute_op(&mut self, _ip: usize, _: &OpCode, _: &[Value]) {}
+    fn observe_execute_op(&mut self, _ip: CodeIdx, _: &OpCode, _: &[Value]) {}
 }
 
 #[derive(Default)]
@@ -162,8 +162,8 @@ impl<W: Write> Observer for TracingObserver<W> {
         );
     }
 
-    fn observe_execute_op(&mut self, ip: usize, op: &OpCode, stack: &[Value]) {
-        let _ = write!(&mut self.writer, "{:04} {:?}\t[ ", ip, op);
+    fn observe_execute_op(&mut self, ip: CodeIdx, op: &OpCode, stack: &[Value]) {
+        let _ = write!(&mut self.writer, "{:04} {:?}\t[ ", ip.0, op);
 
         for val in stack {
             let _ = write!(&mut self.writer, "{} ", val);
diff --git a/tvix/eval/src/opcode.rs b/tvix/eval/src/opcode.rs
index 27fceaf9677c..aee45d7a4452 100644
--- a/tvix/eval/src/opcode.rs
+++ b/tvix/eval/src/opcode.rs
@@ -1,6 +1,8 @@
 //! This module implements the instruction set running on the abstract
 //! machine implemented by tvix.
 
+use std::ops::{AddAssign, Sub};
+
 /// Index of a constant in the current code chunk.
 #[repr(transparent)]
 #[derive(Clone, Copy, Debug)]
@@ -11,6 +13,20 @@ pub struct ConstantIdx(pub usize);
 #[derive(Clone, Copy, Debug)]
 pub struct CodeIdx(pub usize);
 
+impl AddAssign<usize> for CodeIdx {
+    fn add_assign(&mut self, rhs: usize) {
+        *self = CodeIdx(self.0 + rhs)
+    }
+}
+
+impl Sub<usize> for CodeIdx {
+    type Output = Self;
+
+    fn sub(self, rhs: usize) -> Self::Output {
+        CodeIdx(self.0 - rhs)
+    }
+}
+
 /// Index of a value in the runtime stack.
 #[repr(transparent)]
 #[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd)]
diff --git a/tvix/eval/src/vm.rs b/tvix/eval/src/vm.rs
index 9fb51de544b3..ca024c9d65ec 100644
--- a/tvix/eval/src/vm.rs
+++ b/tvix/eval/src/vm.rs
@@ -13,9 +13,18 @@ use crate::{
 };
 
 struct CallFrame {
+    /// The lambda currently being executed.
     lambda: Rc<Lambda>,
+
+    /// Optional captured upvalues of this frame (if a thunk or
+    /// closure if being evaluated).
     upvalues: Upvalues,
-    ip: usize,
+
+    /// Instruction pointer to the instruction currently being
+    /// executed.
+    ip: CodeIdx,
+
+    /// Stack offset, i.e. the frames "view" into the VM's full stack.
     stack_offset: usize,
 }
 
@@ -134,7 +143,7 @@ impl<'o> VM<'o> {
     }
 
     fn inc_ip(&mut self) -> OpCode {
-        let op = self.chunk().code[self.frame().ip];
+        let op = self.chunk()[self.frame().ip];
         self.frame_mut().ip += 1;
         op
     }
@@ -154,7 +163,7 @@ impl<'o> VM<'o> {
     /// Returns the source span of the instruction currently being
     /// executed.
     fn current_span(&self) -> codemap::Span {
-        self.chunk().get_span(CodeIdx(self.frame().ip - 1))
+        self.chunk().get_span(self.frame().ip - 1)
     }
 
     /// Construct an error from the given ErrorKind and the source
@@ -181,7 +190,7 @@ impl<'o> VM<'o> {
         let frame = CallFrame {
             lambda,
             upvalues,
-            ip: 0,
+            ip: CodeIdx(0),
             stack_offset: self.stack.len() - arg_count,
         };
 
@@ -200,7 +209,7 @@ impl<'o> VM<'o> {
             // Break the loop if this call frame has already run to
             // completion, pop it off, and return the value to the
             // caller.
-            if self.frame().ip == self.chunk().code.len() {
+            if self.frame().ip.0 == self.chunk().code.len() {
                 self.frames.pop();
                 return Ok(self.pop());
             }
@@ -465,7 +474,7 @@ impl<'o> VM<'o> {
                             let mut frame = self.frame_mut();
                             frame.lambda = lambda;
                             frame.upvalues = closure.upvalues().clone();
-                            frame.ip = 0; // reset instruction pointer to beginning
+                            frame.ip = CodeIdx(0); // reset instruction pointer to beginning
                         }
 
                         _ => return Err(self.error(ErrorKind::NotCallable)),