Support array access in public declaration.

powdr-labs · Nov 1, 2023 · 119cf82 · 119cf82
1 parent ea7951d
commit 119cf82
Show file tree

Hide file tree

Showing 7 changed files with 53 additions and 11 deletions.
diff --git a/ast/src/analyzed/mod.rs b/ast/src/analyzed/mod.rs
@@ -370,6 +370,7 @@ pub struct PublicDeclaration {
     pub source: SourceRef,
     pub name: String,
     pub polynomial: PolynomialReference,
+    pub array_index: Option<DegreeType>,
     /// The evaluation point of the polynomial, not the array index.
     pub index: DegreeType,
 }

diff --git a/ast/src/parsed/display.rs b/ast/src/parsed/display.rs
@@ -337,8 +337,15 @@ impl<T: Display> Display for PilStatement<T> {
             PilStatement::PolynomialDefinition(_, name, value) => {
                 write!(f, "pol {name} = {value};")
             }
-            PilStatement::PublicDeclaration(_, name, poly, index) => {
-                write!(f, "public {name} = {poly}({index});")
+            PilStatement::PublicDeclaration(_, name, poly, array_index, index) => {
+                write!(
+                    f,
+                    "public {name} = {poly}{}({index});",
+                    array_index
+                        .as_ref()
+                        .map(|i| format!("[{i}]"))
+                        .unwrap_or_default()
+                )
             }
             PilStatement::PolynomialConstantDeclaration(_, names) => {
                 write!(f, "pol constant {};", names.iter().format(", "))

diff --git a/ast/src/parsed/mod.rs b/ast/src/parsed/mod.rs
@@ -20,7 +20,13 @@ pub enum PilStatement<T> {
     Namespace(usize, String, Expression<T>),
     LetStatement(usize, String, Option<Expression<T>>),
     PolynomialDefinition(usize, String, Expression<T>),
-    PublicDeclaration(usize, String, NamespacedPolynomialReference, Expression<T>),
+    PublicDeclaration(
+        usize,
+        String,
+        NamespacedPolynomialReference,
+        Option<Expression<T>>,
+        Expression<T>,
+    ),
     PolynomialConstantDeclaration(usize, Vec<PolynomialName<T>>),
     PolynomialConstantDefinition(usize, String, FunctionDefinition<T>),
     PolynomialCommitDeclaration(usize, Vec<PolynomialName<T>>, Option<FunctionDefinition<T>>),

diff --git a/ast/src/parsed/visitor.rs b/ast/src/parsed/visitor.rs
@@ -204,10 +204,14 @@ impl<T> ExpressionVisitable<Expression<T, NamespacedPolynomialReference>> for Pi
             PilStatement::Namespace(_, _, e)
             | PilStatement::PolynomialDefinition(_, _, e)
             | PilStatement::PolynomialIdentity(_, e)
-            | PilStatement::PublicDeclaration(_, _, _, e)
+            | PilStatement::PublicDeclaration(_, _, _, None, e)
             | PilStatement::ConstantDefinition(_, _, e)
             | PilStatement::LetStatement(_, _, Some(e)) => e.visit_expressions_mut(f, o),
 
+            PilStatement::PublicDeclaration(_, _, _, Some(i), e) => [i, e]
+                .into_iter()
+                .try_for_each(|e| e.visit_expressions_mut(f, o)),
+
             PilStatement::PolynomialConstantDefinition(_, _, fundef)
             | PilStatement::PolynomialCommitDeclaration(_, _, Some(fundef)) => {
                 fundef.visit_expressions_mut(f, o)
@@ -240,10 +244,14 @@ impl<T> ExpressionVisitable<Expression<T, NamespacedPolynomialReference>> for Pi
             PilStatement::Namespace(_, _, e)
             | PilStatement::PolynomialDefinition(_, _, e)
             | PilStatement::PolynomialIdentity(_, e)
-            | PilStatement::PublicDeclaration(_, _, _, e)
+            | PilStatement::PublicDeclaration(_, _, _, None, e)
             | PilStatement::ConstantDefinition(_, _, e)
             | PilStatement::LetStatement(_, _, Some(e)) => e.visit_expressions(f, o),
 
+            PilStatement::PublicDeclaration(_, _, _, Some(i), e) => [i, e]
+                .into_iter()
+                .try_for_each(|e| e.visit_expressions(f, o)),
+
             PilStatement::PolynomialConstantDefinition(_, _, fundef)
             | PilStatement::PolynomialCommitDeclaration(_, _, Some(fundef)) => {
                 fundef.visit_expressions(f, o)

diff --git a/parser/src/lib.rs b/parser/src/lib.rs
@@ -235,6 +235,16 @@ public out = y(%last_row);"#;
             assert_eq!(input.trim(), printed.trim());
         }
 
+        #[test]
+        fn reparse_arrays() {
+            let input = r#"pol witness y[3];\ny - 2 = 0;\ny[2] - 2 = 0;\npublic out = y[1](2);"#;
+            let printed = format!(
+                "{}",
+                parse::<GoldilocksField>(Some("input"), input).unwrap()
+            );
+            assert_eq!(input.trim(), printed.trim());
+        }
+
         #[test]
         fn reparse_strings_and_tuples() {
             let input = r#"constant %N = ("abc", 3);"#;

diff --git a/parser/src/powdr.lalrpop b/parser/src/powdr.lalrpop
@@ -86,7 +86,10 @@ PolynomialDefinition: PilStatement<T> = {
 }
 
 PublicDeclaration: PilStatement<T> = {
-    <@L> "public" <Identifier> "=" <NamespacedPolynomialReference> "(" <Expression> ")" => PilStatement::PublicDeclaration(<>)
+    <@L> "public" <Identifier> "="
+        <NamespacedPolynomialReference>
+        <("[" <Expression> "]")?>
+        "(" <Expression> ")" => PilStatement::PublicDeclaration(<>)
 }
 
 PolynomialConstantDeclaration: PilStatement<T> = {

diff --git a/pil_analyzer/src/pil_analyzer.rs b/pil_analyzer/src/pil_analyzer.rs
@@ -126,9 +126,14 @@ impl<T: FieldElement> PILAnalyzer<T> {
                     Some(FunctionDefinition::Expression(value)),
                 );
             }
-            PilStatement::PublicDeclaration(start, name, polynomial, index) => {
-                self.handle_public_declaration(self.to_source_ref(start), name, polynomial, index)
-            }
+            PilStatement::PublicDeclaration(start, name, polynomial, array_index, index) => self
+                .handle_public_declaration(
+                    self.to_source_ref(start),
+                    name,
+                    polynomial,
+                    array_index,
+                    index,
+                ),
             PilStatement::PolynomialConstantDeclaration(start, polynomials) => self
                 .handle_polynomial_declarations(
                     self.to_source_ref(start),
@@ -416,8 +421,9 @@ impl<T: FieldElement> PILAnalyzer<T> {
         &mut self,
         source: SourceRef,
         name: String,
-        poly: ::ast::parsed::NamespacedPolynomialReference,
-        index: ::ast::parsed::Expression<T>,
+        poly: parsed::NamespacedPolynomialReference,
+        array_index: Option<parsed::Expression<T>>,
+        index: parsed::Expression<T>,
     ) {
         let id = self.public_declarations.len() as u64;
         self.public_declarations.insert(
@@ -428,6 +434,7 @@ impl<T: FieldElement> PILAnalyzer<T> {
                 name: name.to_string(),
                 polynomial: ExpressionProcessor::new(self)
                     .process_namespaced_polynomial_reference(poly),
+                array_index: array_index.map(|i| self.evaluate_expression(i).unwrap().to_degree()),
                 index: self.evaluate_expression(index).unwrap().to_degree(),
             },
         );