Clean up var_count implementation

mkeeter · May 25, 2024 · fda50ba · fda50ba
1 parent a91452e
commit fda50ba
Show file tree

Hide file tree

Showing 5 changed files with 7 additions and 19 deletions.
diff --git a/fidget/src/core/compiler/alloc.rs b/fidget/src/core/compiler/alloc.rs
@@ -273,7 +273,7 @@ impl<const N: usize> RegisterAllocator<N> {
     #[inline(always)]
     pub fn op(&mut self, op: SsaOp) {
         match op {
-            SsaOp::Input(out, i) => self.op_input(out, i.try_into().unwrap()),
+            SsaOp::Input(out, i) => self.op_input(out, i),
             SsaOp::CopyImm(out, imm) => self.op_copy_imm(out, imm),
 
             SsaOp::NegReg(..)
@@ -674,8 +674,6 @@ impl<const N: usize> RegisterAllocator<N> {
     /// Pushes an [`Input`](crate::compiler::RegOp::Input) operation to the tape
     #[inline(always)]
     fn op_input(&mut self, out: u32, i: u32) {
-        // TODO: tightly pack variables (which may be sparse) into slots
-        self.out.var_count = self.out.var_count.max(i as u32 + 1);
         self.op_out_only(out, |out| RegOp::Input(out, i));
     }
 }
diff --git a/fidget/src/core/compiler/reg_tape.rs b/fidget/src/core/compiler/reg_tape.rs
@@ -10,9 +10,6 @@ pub struct RegTape {
 
     /// Total allocated slots
     pub(super) slot_count: u32,
-
-    /// Number of variables
-    pub(super) var_count: u32,
 }
 
 impl RegTape {
@@ -35,7 +32,6 @@ impl RegTape {
         Self {
             tape: vec![],
             slot_count: 1,
-            var_count: 0,
         }
     }
 
@@ -50,11 +46,6 @@ impl RegTape {
     pub fn slot_count(&self) -> usize {
         self.slot_count as usize
     }
-    /// Returns the number of variables (inputs) used in this tape
-    #[inline]
-    pub fn var_count(&self) -> usize {
-        self.var_count as usize
-    }
     /// Returns the number of elements in the tape
     #[inline]
     pub fn len(&self) -> usize {

diff --git a/fidget/src/core/eval/tracing.rs b/fidget/src/core/eval/tracing.rs
@@ -51,9 +51,7 @@ pub trait TracingEvaluator: Default {
         vars: &[Self::Data],
         var_count: usize,
     ) -> Result<(), Error> {
-        // It's fine if the caller has given us extra variables (e.g. due to
-        // tape simplification), but it must have given us enough.
-        if vars.len() < var_count {
+        if vars.len() != var_count {
             Err(Error::BadVarSlice(vars.len(), var_count))
         } else {
             Ok(())

diff --git a/fidget/src/core/vm/data.rs b/fidget/src/core/vm/data.rs
@@ -109,9 +109,12 @@ impl<const N: usize> VmData<N> {
         self.asm.slot_count()
     }
 
-    /// Returns the number of variables (inputs) in the inner VM tape
+    /// Returns the number of variables that may be used
+    ///
+    /// Note that this can sometimes be an overestimate, if the inner tape has
+    /// been simplified.
     pub fn var_count(&self) -> usize {
-        self.asm.var_count()
+        self.vars.len()
     }
 
     /// Simplifies both inner tapes, using the provided choice array

diff --git a/fidget/src/jit/mod.rs b/fidget/src/jit/mod.rs
@@ -992,7 +992,6 @@ impl JitTracingEval {
     ) -> (T, Option<&VmTrace>) {
         let mut simplify = 0;
         self.choices.resize(tape.choice_count, Choice::Unknown);
-        assert!(tape.var_count <= 3);
         self.choices.fill(Choice::Unknown);
         let out = unsafe {
             (tape.fn_trace)(
@@ -1114,7 +1113,6 @@ unsafe impl<T> Sync for JitBulkFn<T> {}
 impl<T: From<f32> + Copy + SimdSize> JitBulkEval<T> {
     /// Evaluate multiple points
     fn eval(&mut self, tape: &JitBulkFn<T>, vars: &[&[T]]) -> &[T] {
-        assert!(tape.var_count <= 3);
         let n = vars.first().map(|v| v.len()).unwrap_or(0);
         self.out.resize(n, f32::NAN.into());
         self.out.fill(f32::NAN.into());