Fix column by name. (#2602)

Author: chriseth

ast/src/analyzed/mod.rs

@@ -136,6 +136,19 @@ impl<T> Analyzed<T> {
             .collect()
     }
 
+    /// Tries to resolve a symbol by name. Does not support individual array elements,
+    /// just plain symbols. Also supports intermediate columns.
+    pub fn try_symbol_by_name(&self, name: &str) -> Option<&Symbol> {
+        self.definitions
+            .get(name)
+            .map(|(symbol, _)| symbol)
+            .or_else(|| {
+                self.intermediate_columns
+                    .get(name)
+                    .map(|(symbol, _)| symbol)
+            })
+    }
+
     pub fn constant_polys_in_source_order(
         &self,
     ) -> impl Iterator<Item = &(Symbol, Option<FunctionValueDefinition>)> {

executor/src/witgen/jit/prover_function_heuristics.rs

@@ -15,6 +15,7 @@ use powdr_number::FieldElement;
 use crate::witgen::{machines::MachineParts, FixedData};
 
 pub trait TrySymbolByName: Copy {
+    /// Tries to resolve a symbol by name.
     fn try_symbol_by_name<'a>(&'a self, name: &str) -> Option<&'a Symbol>;
 }
 
@@ -540,13 +541,13 @@ fn unpack(mut e: &Expression) -> &Expression {
 
 impl<T: FieldElement> TrySymbolByName for &FixedData<'_, T> {
     fn try_symbol_by_name<'a>(&'a self, name: &str) -> Option<&'a Symbol> {
-        FixedData::try_symbol_by_name(self, name)
+        self.analyzed.try_symbol_by_name(name)
     }
 }
 
 impl<T> TrySymbolByName for &Analyzed<T> {
     fn try_symbol_by_name<'a>(&'a self, name: &str) -> Option<&'a Symbol> {
-        self.definitions.get(name).map(|(s, _)| s)
+        Analyzed::try_symbol_by_name(self, name)
     }
 }
 

executor/src/witgen/mod.rs

@@ -322,6 +322,7 @@ pub struct FixedData<'a, T: FieldElement> {
     bus_receives: BTreeMap<T, BusReceive<T>>,
     fixed_cols: FixedColumnMap<FixedColumn<'a, T>>,
     witness_cols: WitnessColumnMap<WitnessColumn<'a, T>>,
+    /// A map from column name to PolyID, supports array elements.
     column_by_name: HashMap<String, PolyID>,
     challenges: BTreeMap<u64, T>,
     global_range_constraints: GlobalConstraints<T>,
@@ -408,12 +409,7 @@ impl<'a, T: FieldElement> FixedData<'a, T> {
             bus_receives,
             fixed_cols,
             witness_cols,
-            column_by_name: analyzed
-                .definitions
-                .iter()
-                .filter(|(_, (symbol, _))| matches!(symbol.kind, SymbolKind::Poly(_)))
-                .flat_map(|(_, (symbol, _))| symbol.array_elements())
-                .collect(),
+            column_by_name: analyzed.name_to_poly_id().into_iter().collect(),
             challenges,
             global_range_constraints,
             intermediate_definitions,
@@ -533,17 +529,12 @@ impl<'a, T: FieldElement> FixedData<'a, T> {
         }
     }
 
+    /// Returns a PolyID given the name of it. Supports array elements
+    /// in that for `"X[7]"` it returns the PolyID of element 7 of the array `X`.
     pub fn try_column_by_name(&self, name: &str) -> Option<PolyID> {
         self.column_by_name.get(name).cloned()
     }
 
-    pub fn try_symbol_by_name(&self, name: &str) -> Option<&Symbol> {
-        self.analyzed
-            .definitions
-            .get(name)
-            .map(|(symbol, _)| symbol)
-    }
-
     fn external_witness(&self, row: DegreeType, column: &PolyID) -> Option<T> {
         self.witness_cols[column]
             .external_values