Always use first parameter as operation ID. (#2498)

Author: chriseth

executor/src/witgen/jit/function_cache.rs

@@ -81,31 +81,43 @@ impl<'a, T: FieldElement> FunctionCache<'a, T> {
         identity_id: u64,
         known_args: &BitVec,
         known_concrete: Option<(usize, T)>,
-    ) -> Option<&CacheEntry<T>> {
-        let cache_key = CacheKey {
+    ) -> &Option<CacheEntry<T>> {
+        // First try the generic version, then the specific.
+        let mut key = CacheKey {
             identity_id,
             known_args: known_args.clone(),
-            known_concrete,
+            known_concrete: None,
         };
-        self.ensure_cache(can_process, &cache_key);
-        self.witgen_functions.get(&cache_key).unwrap().as_ref()
-    }
 
-    fn ensure_cache(&mut self, can_process: impl CanProcessCall<T>, cache_key: &CacheKey<T>) {
-        if self.witgen_functions.contains_key(cache_key) {
-            return;
+        if self.ensure_cache(can_process.clone(), &key).is_none() && known_concrete.is_some() {
+            key = CacheKey {
+                identity_id,
+                known_args: known_args.clone(),
+                known_concrete,
+            };
+            self.ensure_cache(can_process.clone(), &key);
         }
+        self.ensure_cache(can_process, &key)
+    }
 
-        record_start("Auto-witgen code derivation");
-        let f = match T::known_field() {
-            // Currently, we only support the Goldilocks fields
-            Some(KnownField::GoldilocksField) => {
-                self.compile_witgen_function(can_process, cache_key)
-            }
-            _ => None,
-        };
-        assert!(self.witgen_functions.insert(cache_key.clone(), f).is_none());
-        record_end("Auto-witgen code derivation");
+    fn ensure_cache(
+        &mut self,
+        can_process: impl CanProcessCall<T>,
+        cache_key: &CacheKey<T>,
+    ) -> &Option<CacheEntry<T>> {
+        if !self.witgen_functions.contains_key(cache_key) {
+            record_start("Auto-witgen code derivation");
+            let f = match T::known_field() {
+                // Currently, we only support the Goldilocks fields
+                Some(KnownField::GoldilocksField) => {
+                    self.compile_witgen_function(can_process, cache_key)
+                }
+                _ => None,
+            };
+            assert!(self.witgen_functions.insert(cache_key.clone(), f).is_none());
+            record_end("Auto-witgen code derivation");
+        }
+        self.witgen_functions.get(cache_key).unwrap()
     }
 
     fn compile_witgen_function(

executor/src/witgen/machines/block_machine.rs

@@ -172,18 +172,16 @@ impl<'a, T: FieldElement> Machine<'a, T> for BlockMachine<'a, T> {
         known_arguments: &BitVec,
         range_constraints: Vec<RangeConstraint<T>>,
     ) -> (bool, Vec<RangeConstraint<T>>) {
-        // We use the input range constraints to see if there is a column
-        // containing the substring "operation_id" which is constrained to a
-        // single value and use that value as part of the cache key.
-        let operation_id = self.find_operation_id(identity_id).and_then(|index| {
-            let v = range_constraints[index].try_to_single_value()?;
-            Some((index, v))
-        });
+        let fixed_first_input = if !known_arguments.is_empty() && known_arguments[0] {
+            range_constraints[0].try_to_single_value().map(|v| (0, v))
+        } else {
+            None
+        };
         match self.function_cache.compile_cached(
             can_process,
             identity_id,
             known_arguments,
-            operation_id,
+            fixed_first_input,
         ) {
             Some(entry) => (true, entry.range_constraints.clone()),
             None => (false, range_constraints),
@@ -200,12 +198,10 @@ impl<'a, T: FieldElement> Machine<'a, T> for BlockMachine<'a, T> {
             return Err(EvalError::RowsExhausted(self.name.clone()));
         }
 
-        let operation_id =
-            self.find_operation_id(identity_id)
-                .and_then(|index| match &values[index] {
-                    LookupCell::Input(v) => Some((index, **v)),
-                    LookupCell::Output(_) => None,
-                });
+        let fixed_first_input = match &values.first() {
+            Some(LookupCell::Input(v)) => Some((0, **v)),
+            None | Some(LookupCell::Output(_)) => None,
+        };
 
         self.data.finalize_all();
         let data = self.data.append_new_finalized_rows(self.block_size);
@@ -215,7 +211,7 @@ impl<'a, T: FieldElement> Machine<'a, T> for BlockMachine<'a, T> {
             identity_id,
             values,
             data,
-            operation_id,
+            fixed_first_input,
         )?;
         assert!(success);
         self.block_count_jit += 1;
@@ -452,13 +448,12 @@ impl<'a, T: FieldElement> BlockMachine<'a, T> {
         }
 
         let known_inputs = arguments.iter().map(|e| e.is_constant()).collect();
-        let operation_id = self.find_operation_id(identity_id).and_then(|index| {
-            let v = arguments[index].constant_value()?;
-            Some((index, v))
-        });
+        let fixed_first_input = arguments
+            .first()
+            .and_then(|a| a.constant_value().map(|v| (0, v)));
         if self
             .function_cache
-            .compile_cached(mutable_state, identity_id, &known_inputs, operation_id)
+            .compile_cached(mutable_state, identity_id, &known_inputs, fixed_first_input)
             .is_some()
         {
             let caller_data = CallerData::new(arguments, range_constraints);
@@ -537,12 +532,10 @@ impl<'a, T: FieldElement> BlockMachine<'a, T> {
         self.data.finalize_all();
 
         let mut lookup_cells = caller_data.as_lookup_cells();
-        let operation_id =
-            self.find_operation_id(identity_id)
-                .and_then(|index| match &lookup_cells[index] {
-                    LookupCell::Input(v) => Some((index, **v)),
-                    LookupCell::Output(_) => None,
-                });
+        let operation_id = match &lookup_cells.first() {
+            Some(LookupCell::Input(v)) => Some((0, **v)),
+            None | Some(LookupCell::Output(_)) => None,
+        };
         let data = self.data.append_new_finalized_rows(self.block_size);
 
         let success = self.function_cache.process_lookup_direct(
@@ -557,13 +550,6 @@ impl<'a, T: FieldElement> BlockMachine<'a, T> {
         caller_data.into()
     }
 
-    fn find_operation_id(&self, identity_id: u64) -> Option<usize> {
-        let right = &self.parts.connections[&identity_id].right.expressions;
-        right.iter().position(|r| {
-            try_to_simple_poly(r).is_some_and(|poly| poly.name.contains("operation_id"))
-        })
-    }
-
     fn process<'b, Q: QueryCallback<T>>(
         &self,
         mutable_state: &MutableState<'a, T, Q>,