Fix PIL optimizer

pilopt/src/lib.rs

@@ -152,10 +152,23 @@ fn collect_required_symbols<'a, T: FieldElement>(
             .values()
             .map(|p| SymbolReference::from(&p.polynomial.name)),
     );
+
+    let poly_ref_to_id = pil_file
+        .definitions
+        .values()
+        .filter_map(|(symbol, _)| matches!(symbol.kind, SymbolKind::Poly(_)).then_some(symbol))
+        .flat_map(|symbol| symbol.array_elements())
+        .collect::<BTreeMap<_, _>>();
+
     for fun in &pil_file.prover_functions {
         for e in fun.all_children() {
             if let Expression::Reference(_, Reference::Poly(poly_ref)) = e {
-                required_names.insert(SymbolReference::from(poly_ref));
+                let symbol_ref = match poly_ref_to_id.get(&poly_ref.name) {
+                    Some(poly_id) => poly_id_to_definition_name[poly_id].into(),
+                    None => SymbolReference::from(poly_ref),
+                };
+
+                required_names.insert(symbol_ref);
             }
         }
     }

pilopt/tests/optimizer.rs

@@ -281,3 +281,30 @@ fn enum_ref_by_trait() {
     let optimized = optimize(analyze_string::<GoldilocksField>(input).unwrap()).to_string();
     assert_eq!(optimized, expectation);
 }
+
+
+#[test]
+fn handle_array_references_in_prover_functions() {
+    let input = r#"namespace N(8);
+        col witness x;
+        col fixed cnt(i) { inc(i) };
+        let inc = |x| x + 1;
+        // these are removed
+        col witness k;
+        col k2 = k;
+        let rec: -> int = || rec();
+        let a: int -> int = |i| b(i + 1);
+        let b: int -> int = |j| 8;
+        // identity
+        [ x ] in [ cnt ];
+
+    "#;
+    let expectation = r#"namespace N(65536);
+    col witness x;
+    col fixed cnt(i) { N::inc(i) };
+    let inc: int -> int = |x| x + 1_int;
+    [N::x] in [N::cnt];
+"#;
+    let optimized = optimize(analyze_string::<GoldilocksField>(input).unwrap()).to_string();
+    assert_eq!(optimized, expectation);
+}