refactor(tvix/eval): add opcode::Count type for less ambiguity r/4622

Change-Id: Ibde0b2baa1128a74c1364ee9a6330b62db3da699
Reviewed-on: https://cl.tvl.fyi/c/depot/+/6288
Autosubmit: tazjin <tazjin@tvl.su>
Tested-by: BuildkiteCI
Reviewed-by: sterni <sternenseemann@systemli.org>

3 files changed, 23 insertions, 17 deletions

diff --git a/tvix/eval/src/compiler/mod.rs b/tvix/eval/src/compiler/mod.rs
index 37ec3915b4be..2d888a6e3bed 100644
--- a/tvix/eval/src/compiler/mod.rs
+++ b/tvix/eval/src/compiler/mod.rs
@@ -22,7 +22,7 @@ use std::rc::Rc;
 
 use crate::chunk::Chunk;
 use crate::errors::{Error, ErrorKind, EvalResult};
-use crate::opcode::{CodeIdx, JumpOffset, OpCode, StackIdx};
+use crate::opcode::{CodeIdx, Count, JumpOffset, OpCode, StackIdx};
 use crate::value::{Closure, Lambda, Value};
 use crate::warnings::{EvalWarning, WarningKind};
 
@@ -298,7 +298,7 @@ impl Compiler {
         }
 
         if count != 1 {
-            self.chunk().push_op(OpCode::OpInterpolate(count));
+            self.chunk().push_op(OpCode::OpInterpolate(Count(count)));
         }
     }
 
@@ -462,7 +462,7 @@ impl Compiler {
             self.compile(item);
         }
 
-        self.chunk().push_op(OpCode::OpList(count));
+        self.chunk().push_op(OpCode::OpList(Count(count)));
     }
 
     // Compile attribute set literals into equivalent bytecode.
@@ -547,7 +547,7 @@ impl Compiler {
             // otherwise we need to emit an instruction to construct
             // the attribute path.
             if key_count > 1 {
-                self.chunk().push_op(OpCode::OpAttrPath(key_count));
+                self.chunk().push_op(OpCode::OpAttrPath(Count(key_count)));
             }
 
             // The value is just compiled as normal so that its
@@ -556,7 +556,7 @@ impl Compiler {
             self.compile(kv.value().unwrap());
         }
 
-        self.chunk().push_op(OpCode::OpAttrs(count));
+        self.chunk().push_op(OpCode::OpAttrs(Count(count)));
     }
 
     fn compile_select(&mut self, node: ast::Select) {
@@ -925,7 +925,7 @@ impl Compiler {
         }
 
         if pops > 0 {
-            self.chunk().push_op(OpCode::OpCloseScope(pops));
+            self.chunk().push_op(OpCode::OpCloseScope(Count(pops)));
         }
 
         while !self.scope().with_stack.is_empty()
diff --git a/tvix/eval/src/opcode.rs b/tvix/eval/src/opcode.rs
index 82b72952d698..f2daf9bd8e52 100644
--- a/tvix/eval/src/opcode.rs
+++ b/tvix/eval/src/opcode.rs
@@ -21,6 +21,12 @@ pub struct StackIdx(pub usize);
 #[derive(Clone, Copy, Debug)]
 pub struct JumpOffset(pub usize);
 
+/// Provided count for an instruction (could represent e.g. a number
+/// of elements).
+#[repr(transparent)]
+#[derive(Clone, Copy, Debug)]
+pub struct Count(pub usize);
+
 #[allow(clippy::enum_variant_names)]
 #[warn(variant_size_differences)]
 #[derive(Clone, Copy, Debug)]
@@ -60,8 +66,8 @@ pub enum OpCode {
     OpJumpIfNotFound(JumpOffset),
 
     // Attribute sets
-    OpAttrs(usize),
-    OpAttrPath(usize),
+    OpAttrs(Count),
+    OpAttrPath(Count),
     OpAttrsUpdate,
     OpAttrsSelect,
     OpAttrsTrySelect,
@@ -73,11 +79,11 @@ pub enum OpCode {
     OpResolveWith,
 
     // Lists
-    OpList(usize),
+    OpList(Count),
     OpConcat,
 
     // Strings
-    OpInterpolate(usize),
+    OpInterpolate(Count),
 
     // Type assertion operators
     OpAssertBool,
@@ -86,7 +92,7 @@ pub enum OpCode {
     OpGetLocal(StackIdx),
 
     // Close scopes while leaving their expression value around.
-    OpCloseScope(usize), // number of locals to pop
+    OpCloseScope(Count), // number of locals to pop
 
     // Asserts stack top is a boolean, and true.
     OpAssert,
diff --git a/tvix/eval/src/vm.rs b/tvix/eval/src/vm.rs
index 28b1636488af..85f81bac6efa 100644
--- a/tvix/eval/src/vm.rs
+++ b/tvix/eval/src/vm.rs
@@ -6,7 +6,7 @@ use std::rc::Rc;
 use crate::{
     chunk::Chunk,
     errors::{ErrorKind, EvalResult},
-    opcode::{JumpOffset, OpCode, StackIdx},
+    opcode::{Count, JumpOffset, OpCode, StackIdx},
     value::{Closure, Lambda, NixAttrs, NixList, Value},
 };
 
@@ -199,8 +199,8 @@ impl VM {
                 OpCode::OpTrue => self.push(Value::Bool(true)),
                 OpCode::OpFalse => self.push(Value::Bool(false)),
 
-                OpCode::OpAttrs(count) => self.run_attrset(count)?,
-                OpCode::OpAttrPath(count) => self.run_attr_path(count)?,
+                OpCode::OpAttrs(Count(count)) => self.run_attrset(count)?,
+                OpCode::OpAttrPath(Count(count)) => self.run_attr_path(count)?,
 
                 OpCode::OpAttrsUpdate => {
                     let rhs = unwrap_or_clone_rc(self.pop().to_attrs()?);
@@ -252,7 +252,7 @@ impl VM {
                     self.push(Value::Bool(result));
                 }
 
-                OpCode::OpList(count) => {
+                OpCode::OpList(Count(count)) => {
                     let list =
                         NixList::construct(count, self.stack.split_off(self.stack.len() - count));
                     self.push(Value::List(list));
@@ -264,7 +264,7 @@ impl VM {
                     self.push(Value::List(lhs.concat(&rhs)))
                 }
 
-                OpCode::OpInterpolate(count) => self.run_interpolate(count)?,
+                OpCode::OpInterpolate(Count(count)) => self.run_interpolate(count)?,
 
                 OpCode::OpJump(JumpOffset(offset)) => {
                     self.frame_mut().ip += offset;
@@ -306,7 +306,7 @@ impl VM {
 
                 // Remove the given number of elements from the stack,
                 // but retain the top value.
-                OpCode::OpCloseScope(count) => {
+                OpCode::OpCloseScope(Count(count)) => {
                     // Immediately move the top value into the right
                     // position.
                     let target_idx = self.stack.len() - 1 - count;