Refactor public declaration and public reference (#2516)

Started as a small fix but expanded to a complete refactoring of
`PublicDeclaration` and `PublicReference`.

Now public reference work the same way as any reference and the syntax
is without the colon (`:out` -> `out`). `PublicDeclaration` are also put
under definitions in the analyzer. Both public reference and public
declaration use the absolute path, so it's easy to refer from one to the
other now.

This is needed for #2502 to match public reference to public declaration
in the backend. Currently, different namespaces can have public
references with the same name, and it will confuse the backend on which
public reference value to fetch.

Author: qwang98

asm-to-pil/src/vm_to_constrained.rs

@@ -741,7 +741,6 @@ impl<T: FieldElement> VMConverter<T> {
 
     fn process_assignment_value(&self, value: Expression) -> Vec<(T, AffineExpressionComponent)> {
         match value {
-            Expression::PublicReference(_, _) => panic!(),
             Expression::IndexAccess(_, _) => panic!(),
             Expression::FunctionCall(_, _) => panic!(),
             Expression::Reference(_, reference) => {

asmopt/src/lib.rs

@@ -204,7 +204,6 @@ fn expr_to_ref(expr: &Expression) -> Option<String> {
         Expression::Reference(_, NamespacedPolynomialReference { path, .. }) => {
             Some(path.to_string())
         }
-        Expression::PublicReference(_, pref) => Some(pref.clone()),
         _ => None,
     }
 }

ast/src/analyzed/display.rs

@@ -95,6 +95,15 @@ impl<T: Display> Display for Analyzed<T> {
                                 format_witness_column(&name, symbol, definition),
                             )?,
                             SymbolKind::Poly(PolynomialType::Intermediate) => unreachable!(),
+                            SymbolKind::Public() => {
+                                if let Some(FunctionValueDefinition::PublicDeclaration(decl)) =
+                                    definition
+                                {
+                                    writeln_indented(f, format_public_declaration(&name, decl))?;
+                                } else {
+                                    unreachable!() // For now, public symbol should always have a public declaration
+                                }
+                            }
                             SymbolKind::Other() => {
                                 assert!(symbol.stage.is_none());
                                 match definition {
@@ -159,11 +168,6 @@ impl<T: Display> Display for Analyzed<T> {
                         panic!()
                     }
                 }
-                StatementIdentifier::PublicDeclaration(name) => {
-                    let decl = &self.public_declarations[name];
-                    let name = update_namespace(&decl.name, f)?;
-                    writeln_indented(f, format_public_declaration(&name, decl))?;
-                }
                 StatementIdentifier::ProofItem(i) => {
                     writeln_indented(f, &self.identities[*i])?;
                 }
@@ -243,14 +247,7 @@ fn format_witness_column(
 }
 
 fn format_public_declaration(name: &str, decl: &PublicDeclaration) -> String {
-    format!(
-        "public {name} = {}{}({});",
-        decl.polynomial,
-        decl.array_index
-            .map(|i| format!("[{i}]"))
-            .unwrap_or_default(),
-        decl.index
-    )
+    format!("public {name} = {};", decl.value)
 }
 
 impl Display for FunctionValueDefinition {
@@ -277,7 +274,8 @@ impl Display for FunctionValueDefinition {
             FunctionValueDefinition::TypeDeclaration(_)
             | FunctionValueDefinition::TypeConstructor(_, _)
             | FunctionValueDefinition::TraitDeclaration(_)
-            | FunctionValueDefinition::TraitFunction(_, _) => {
+            | FunctionValueDefinition::TraitFunction(_, _)
+            | FunctionValueDefinition::PublicDeclaration(_) => {
                 panic!("Should not use this formatting function.")
             }
         }
@@ -484,7 +482,7 @@ impl<T: Display> Display for AlgebraicExpression<T> {
     fn fmt(&self, f: &mut Formatter<'_>) -> Result {
         match self {
             AlgebraicExpression::Reference(reference) => write!(f, "{reference}"),
-            AlgebraicExpression::PublicReference(name) => write!(f, ":{name}"),
+            AlgebraicExpression::PublicReference(name) => write!(f, "{name}"),
             AlgebraicExpression::Challenge(challenge) => {
                 write!(
                     f,

ast/src/analyzed/mod.rs

@@ -22,16 +22,15 @@ use crate::parsed::visitor::{Children, ExpressionVisitable};
 pub use crate::parsed::BinaryOperator;
 pub use crate::parsed::UnaryOperator;
 use crate::parsed::{
-    self, ArrayExpression, EnumDeclaration, EnumVariant, NamedType, SourceReference,
-    TraitDeclaration, TraitImplementation, TypeDeclaration,
+    self, ArrayExpression, EnumDeclaration, EnumVariant, FunctionCall, IndexAccess, NamedType,
+    SourceReference, TraitDeclaration, TraitImplementation, TypeDeclaration,
 };
 pub use contains_next_ref::ContainsNextRef;
 
 #[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, PartialEq, Eq, Hash)]
 pub enum StatementIdentifier {
     /// Either an intermediate column or a definition.
     Definition(String),
-    PublicDeclaration(String),
     /// Index into the vector of proof items / identities.
     ProofItem(usize),
     /// Index into the vector of prover functions.
@@ -44,7 +43,6 @@ pub enum StatementIdentifier {
 pub struct Analyzed<T> {
     pub definitions: HashMap<String, (Symbol, Option<FunctionValueDefinition>)>,
     pub solved_impls: SolvedTraitImpls,
-    pub public_declarations: HashMap<String, PublicDeclaration>,
     pub intermediate_columns: HashMap<String, (Symbol, Vec<AlgebraicExpression<T>>)>,
     pub identities: Vec<Identity<T>>,
     pub prover_functions: Vec<Expression>,
@@ -122,7 +120,10 @@ impl<T> Analyzed<T> {
     }
     /// @returns the number of public inputs
     pub fn publics_count(&self) -> usize {
-        self.public_declarations.len()
+        self.definitions
+            .iter()
+            .filter(|(_name, (symbol, _))| matches!(symbol.kind, SymbolKind::Public()))
+            .count()
     }
 
     pub fn name_to_poly_id(&self) -> BTreeMap<String, PolyID> {
@@ -187,8 +188,12 @@ impl<T> Analyzed<T> {
         &self,
     ) -> impl Iterator<Item = (&String, &PublicDeclaration)> {
         self.source_order.iter().filter_map(move |statement| {
-            if let StatementIdentifier::PublicDeclaration(name) = statement {
-                if let Some(public_declaration) = self.public_declarations.get(name) {
+            if let StatementIdentifier::Definition(name) = statement {
+                if let Some((
+                    _,
+                    Some(FunctionValueDefinition::PublicDeclaration(public_declaration)),
+                )) = self.definitions.get(name)
+                {
                     return Some((name, public_declaration));
                 }
             }
@@ -398,29 +403,26 @@ impl<T> Analyzed<T> {
     /// Retrieves (name, col_name, poly_id, offset, stage) of each public witness in the trace.
     pub fn get_publics(&self) -> Vec<(String, String, PolyID, usize, u8)> {
         let mut publics = self
-            .public_declarations
-            .values()
-            .map(|public_declaration| {
+            .public_declarations_in_source_order()
+            .map(|(name, public_declaration)| {
                 let column_name = public_declaration.referenced_poly_name();
                 let (poly_id, stage) = {
-                    let symbol = &self.definitions[&public_declaration.polynomial.name].0;
+                    let symbol = &self.definitions[&public_declaration.referenced_poly().name].0;
                     (
                         symbol
                             .array_elements()
-                            .nth(public_declaration.array_index.unwrap_or_default())
+                            .nth(
+                                public_declaration
+                                    .referenced_poly_array_index()
+                                    .unwrap_or_default(),
+                            )
                             .unwrap()
                             .1,
                         symbol.stage.unwrap_or_default() as u8,
                     )
                 };
-                let row_offset = public_declaration.index as usize;
-                (
-                    public_declaration.name.clone(),
-                    column_name,
-                    poly_id,
-                    row_offset,
-                    stage,
-                )
+                let row_offset = public_declaration.row() as usize;
+                (name.clone(), column_name, poly_id, row_offset, stage)
             })
             .collect::<Vec<_>>();
 
@@ -559,6 +561,7 @@ pub fn type_from_definition(
                     ty: trait_func.ty.clone(),
                 })
             }
+            FunctionValueDefinition::PublicDeclaration(_) => Some(Type::Expr.into()),
         }
     } else {
         assert!(
@@ -771,6 +774,8 @@ impl Symbol {
 pub enum SymbolKind {
     /// Fixed, witness or intermediate polynomial
     Poly(PolynomialType),
+    /// Public declaration
+    Public(),
     /// Other symbol, depends on the type.
     /// Examples include functions not of the type "int -> fe".
     Other(),
@@ -784,6 +789,7 @@ pub enum FunctionValueDefinition {
     TypeConstructor(Arc<EnumDeclaration>, EnumVariant),
     TraitDeclaration(TraitDeclaration),
     TraitFunction(Arc<TraitDeclaration>, NamedType),
+    PublicDeclaration(PublicDeclaration),
 }
 
 impl Children<Expression> for FunctionValueDefinition {
@@ -799,6 +805,7 @@ impl Children<Expression> for FunctionValueDefinition {
             FunctionValueDefinition::TypeConstructor(_, variant) => variant.children(),
             FunctionValueDefinition::TraitDeclaration(trait_decl) => trait_decl.children(),
             FunctionValueDefinition::TraitFunction(_, trait_func) => trait_func.children(),
+            FunctionValueDefinition::PublicDeclaration(pub_decl) => pub_decl.children(),
         }
     }
 
@@ -814,6 +821,7 @@ impl Children<Expression> for FunctionValueDefinition {
             FunctionValueDefinition::TypeConstructor(_, variant) => variant.children_mut(),
             FunctionValueDefinition::TraitDeclaration(trait_decl) => trait_decl.children_mut(),
             FunctionValueDefinition::TraitFunction(_, trait_func) => trait_func.children_mut(),
+            FunctionValueDefinition::PublicDeclaration(pub_decl) => pub_decl.children_mut(),
         }
     }
 }
@@ -836,26 +844,80 @@ impl Children<Expression> for NamedType {
     }
 }
 
-#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, Hash)]
+#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, Hash, PartialEq, Eq)]
 pub struct PublicDeclaration {
     pub id: u64,
     pub source: SourceRef,
     pub name: String,
-    pub polynomial: PolynomialReference,
-    pub array_index: Option<usize>,
-    /// The evaluation point of the polynomial, not the array index.
-    pub index: DegreeType,
+    /// The declaration value, in two possible forms: polynomial[array_index](row) OR polynomial(row)
+    /// where "row" is the evaluation point of the polynomial.
+    pub value: Expression,
 }
 
 impl PublicDeclaration {
+    pub fn referenced_poly(&self) -> &PolynomialReference {
+        match &self.value {
+            Expression::FunctionCall(_, FunctionCall { function, .. }) => match function.as_ref() {
+                Expression::Reference(_, Reference::Poly(poly)) => poly,
+                Expression::IndexAccess(_, IndexAccess { array, .. }) => match array.as_ref() {
+                    Expression::Reference(_, Reference::Poly(poly)) => poly,
+                    _ => panic!("Expected Reference."),
+                },
+                _ => panic!("Expected Reference or IndexAccess."),
+            },
+            _ => panic!("Expected FunctionCall."),
+        }
+    }
+
+    pub fn referenced_poly_array_index(&self) -> Option<usize> {
+        match &self.value {
+            Expression::FunctionCall(_, FunctionCall { function, .. }) => match function.as_ref() {
+                Expression::Reference(_, Reference::Poly(_)) => None,
+                Expression::IndexAccess(_, IndexAccess { index, .. }) => match index.as_ref() {
+                    Expression::Number(_, index) => Some(index.value.clone().try_into().unwrap()),
+                    _ => panic!("Expected Number."),
+                },
+                _ => panic!("Expected Reference or IndexAccess."),
+            },
+            _ => panic!("Expected FunctionCall."),
+        }
+    }
+
+    /// Returns the name of the polynomial referenced by the public declaration.
+    /// Includes the array index if present.
     pub fn referenced_poly_name(&self) -> String {
-        match self.array_index {
-            Some(index) => format!("{}[{}]", self.polynomial.name, index),
-            None => self.polynomial.name.clone(),
+        let name = self.referenced_poly().name.clone();
+        let index = self.referenced_poly_array_index();
+        if let Some(index) = index {
+            format!("{name}[{index}]")
+        } else {
+            name
+        }
+    }
+
+    pub fn row(&self) -> DegreeType {
+        match &self.value {
+            Expression::FunctionCall(_, FunctionCall { arguments, .. }) => {
+                assert!(arguments.len() == 1);
+                match &arguments[0] {
+                    Expression::Number(_, index) => index.value.clone().try_into().unwrap(),
+                    _ => panic!("Expected Number."),
+                }
+            }
+            _ => panic!("Expected FunctionCall."),
         }
     }
 }
 
+impl Children<Expression> for PublicDeclaration {
+    fn children(&self) -> Box<dyn Iterator<Item = &Expression> + '_> {
+        Box::new(self.value.children())
+    }
+    fn children_mut(&mut self) -> Box<dyn Iterator<Item = &mut Expression> + '_> {
+        Box::new(self.value.children_mut())
+    }
+}
+
 #[derive(
     Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Serialize, Deserialize, JsonSchema, Hash,
 )]

ast/src/asm_analysis/display.rs

@@ -120,7 +120,7 @@ impl Display for LinkDefinition {
         write!(
             f,
             "link {}{} {};",
-            if flag == 1.into() {
+            if flag == 1u32.into() {
                 "".to_string()
             } else {
                 format!("if {flag} ")

ast/src/parsed/display.rs

@@ -172,7 +172,7 @@ impl Display for LinkDeclaration {
         write!(
             f,
             "link {}{} {}",
-            if self.flag == 1.into() {
+            if self.flag == 1u32.into() {
                 "".to_string()
             } else {
                 format!("if {} ", self.flag)
@@ -555,7 +555,17 @@ impl Display for FunctionDefinition {
                 )
             }
             FunctionDefinition::Expression(e) => write!(f, " = {e}"),
-            FunctionDefinition::TypeDeclaration(_) | FunctionDefinition::TraitDeclaration(_) => {
+            FunctionDefinition::PublicDeclaration(poly, array_index, index) => {
+                write!(
+                    f,
+                    " = {poly}{}({index});",
+                    array_index
+                        .as_ref()
+                        .map(|i| format!("[{i}]"))
+                        .unwrap_or_default()
+                )
+            }
+            FunctionDefinition::TraitDeclaration(_) | FunctionDefinition::TypeDeclaration(_) => {
                 panic!("Should not use this formatting function.")
             }
         }
@@ -717,7 +727,6 @@ impl<Ref: Display> Display for Expression<Ref> {
     fn fmt(&self, f: &mut Formatter<'_>) -> Result {
         match self {
             Expression::Reference(_, reference) => write!(f, "{reference}"),
-            Expression::PublicReference(_, name) => write!(f, ":{name}"),
             Expression::Number(_, n) => write!(f, "{n}"),
             Expression::String(_, value) => write!(f, "{}", quote(value)),
             Expression::Tuple(_, items) => write!(f, "({})", format_list(items)),