diff --git a/ast/src/analyzed/mod.rs b/ast/src/analyzed/mod.rs index 00cd958c4a..e27b258155 100644 --- 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, /// 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 index 51844f2635..095d8a3583 100644 --- a/ast/src/parsed/display.rs +++ b/ast/src/parsed/display.rs @@ -337,8 +337,15 @@ impl Display for PilStatement { 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 index efe63ff824..9b0ad5597c 100644 --- a/ast/src/parsed/mod.rs +++ b/ast/src/parsed/mod.rs @@ -20,7 +20,13 @@ pub enum PilStatement { Namespace(usize, String, Expression), LetStatement(usize, String, Option>), PolynomialDefinition(usize, String, Expression), - PublicDeclaration(usize, String, NamespacedPolynomialReference, Expression), + PublicDeclaration( + usize, + String, + NamespacedPolynomialReference, + Option>, + Expression, + ), PolynomialConstantDeclaration(usize, Vec>), PolynomialConstantDefinition(usize, String, FunctionDefinition), PolynomialCommitDeclaration(usize, Vec>, Option>), diff --git a/ast/src/parsed/visitor.rs b/ast/src/parsed/visitor.rs index d4f2679233..43a9555f71 100644 --- a/ast/src/parsed/visitor.rs +++ b/ast/src/parsed/visitor.rs @@ -204,10 +204,14 @@ impl ExpressionVisitable> 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 ExpressionVisitable> 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 index 723a19ce59..8a23645825 100644 --- 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::(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 index e4a2c1457c..23de7922fb 100644 --- a/parser/src/powdr.lalrpop +++ b/parser/src/powdr.lalrpop @@ -86,7 +86,10 @@ PolynomialDefinition: PilStatement = { } PublicDeclaration: PilStatement = { - <@L> "public" "=" "(" ")" => PilStatement::PublicDeclaration(<>) + <@L> "public" "=" + + <("[" "]")?> + "(" ")" => PilStatement::PublicDeclaration(<>) } PolynomialConstantDeclaration: PilStatement = { diff --git a/pil_analyzer/src/pil_analyzer.rs b/pil_analyzer/src/pil_analyzer.rs index d5a4ec6a77..5a60724aaf 100644 --- a/pil_analyzer/src/pil_analyzer.rs +++ b/pil_analyzer/src/pil_analyzer.rs @@ -126,9 +126,14 @@ impl PILAnalyzer { 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 PILAnalyzer { &mut self, source: SourceRef, name: String, - poly: ::ast::parsed::NamespacedPolynomialReference, - index: ::ast::parsed::Expression, + poly: parsed::NamespacedPolynomialReference, + array_index: Option>, + index: parsed::Expression, ) { let id = self.public_declarations.len() as u64; self.public_declarations.insert( @@ -428,6 +434,7 @@ impl PILAnalyzer { 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(), }, );