Fix utility function filter. (powdr-labs#1994)

gzanitti · Oct 31, 2024 · 5c37eba · 5c37eba
1 parent d4767fe
commit 5c37eba
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 29 deletions.
diff --git a/airgen/src/lib.rs b/airgen/src/lib.rs
@@ -45,7 +45,7 @@ pub fn compile(input: AnalysisASMFile) -> PILGraph {
                 main,
                 entry_points: Default::default(),
                 objects: [(main_location, Default::default())].into(),
-                statements: utility_functions(input),
+                statements: module_level_pil_statements(input),
             };
         }
         // if there is a single machine, treat it as main
@@ -188,7 +188,7 @@ pub fn compile(input: AnalysisASMFile) -> PILGraph {
         main,
         entry_points,
         objects,
-        statements: utility_functions(input),
+        statements: module_level_pil_statements(input),
     }
 }
 
@@ -216,35 +216,14 @@ fn resolve_submachine_arg(
     }
 }
 
-fn utility_functions(asm_file: AnalysisASMFile) -> BTreeMap<AbsoluteSymbolPath, Vec<PilStatement>> {
+/// Returns the pil statements at module level for each module.
+fn module_level_pil_statements(
+    asm_file: AnalysisASMFile,
+) -> BTreeMap<AbsoluteSymbolPath, Vec<PilStatement>> {
     asm_file
         .modules
         .into_iter()
-        .map(|(module_path, module)| {
-            (
-                module_path,
-                module
-                    .into_inner()
-                    .1
-                    .into_iter()
-                    .filter(|s| match s {
-                        PilStatement::EnumDeclaration(..)
-                        | PilStatement::StructDeclaration(..)
-                        | PilStatement::LetStatement(..) => true,
-                        PilStatement::Include(..) => false,
-                        PilStatement::Namespace(..) => false,
-                        PilStatement::PolynomialDefinition(..) => false,
-                        PilStatement::PublicDeclaration(..) => false,
-                        PilStatement::PolynomialConstantDeclaration(..) => false,
-                        PilStatement::PolynomialConstantDefinition(..) => false,
-                        PilStatement::PolynomialCommitDeclaration(..) => false,
-                        PilStatement::TraitImplementation(..) => false,
-                        PilStatement::TraitDeclaration(..) => false,
-                        PilStatement::Expression(..) => false,
-                    })
-                    .collect(),
-            )
-        })
+        .map(|(module_path, module)| (module_path, module.into_inner().1))
         .collect()
 }
 

diff --git a/ast/src/asm_analysis/mod.rs b/ast/src/asm_analysis/mod.rs
@@ -856,6 +856,7 @@ pub enum StatementReference {
 #[derive(Default, Clone, Debug)]
 pub struct Module {
     machines: BTreeMap<String, Machine>,
+    /// Module-level PIL statements.
     statements: Vec<PilStatement>,
     ordering: Vec<StatementReference>,
 }

diff --git a/ast/src/object/mod.rs b/ast/src/object/mod.rs
@@ -44,7 +44,8 @@ pub struct PILGraph {
     pub main: Machine,
     pub entry_points: Vec<Operation>,
     pub objects: BTreeMap<Location, Object>,
-    /// PIL utilities by module path
+    /// List of module-level PIL statements (utility functions,
+    /// data structures, etc) by module path
     pub statements: BTreeMap<AbsoluteSymbolPath, Vec<PilStatement>>,
 }