From 1c6b55a32f51641808124123f683b1c5c4bc1a47 Mon Sep 17 00:00:00 2001
From: Rute Figueiredo <rfigueiredo.dev@gmail.com>
Date: Fri, 30 Aug 2024 10:44:45 +0100
Subject: [PATCH] parser for declaring arrays of signals and variables

---
 src/parser/ast/statement.rs | 14 ++++++++--
 src/parser/chiquito.lalrpop | 16 ++++++++++--
 src/parser/mod.rs           | 52 +++++++++++++++++++++++++++++++++++++
 3 files changed, 78 insertions(+), 4 deletions(-)

diff --git a/src/parser/ast/statement.rs b/src/parser/ast/statement.rs
index 896cd50e..70cd102a 100644
--- a/src/parser/ast/statement.rs
+++ b/src/parser/ast/statement.rs
@@ -1,5 +1,7 @@
 use std::fmt::Debug;
 
+use num_bigint::BigInt;
+
 use super::{expression::Expression, DebugSymRef, Identifiable, Identifier};
 
 /// An identifier with the type it is declared with.
@@ -9,6 +11,8 @@ pub struct TypedIdDecl<V> {
     pub id: V,
     /// type
     pub ty: Option<V>,
+    /// Array dimension (if any)
+    pub dimension: Option<Expression<BigInt, V>>,
 }
 
 #[derive(Clone)]
@@ -69,7 +73,10 @@ impl<F: Debug> Debug for Statement<F, Identifier> {
                 f,
                 "signal {};",
                 id.iter()
-                    .map(|id| id.id.name())
+                    .map(|id| match &id.dimension {
+                        Some(dim) => format!("{}[{:?}]", id.id.name(), dim),
+                        None => id.id.name().to_string(),
+                    })
                     .collect::<Vec<_>>()
                     .join(", ")
             ),
@@ -77,7 +84,10 @@ impl<F: Debug> Debug for Statement<F, Identifier> {
                 f,
                 "var {};",
                 id.iter()
-                    .map(|id| id.id.name())
+                    .map(|id| match &id.dimension {
+                        Some(dim) => format!("{}[{:?}]", id.id.name(), dim),
+                        None => id.id.name().to_string(),
+                    })
                     .collect::<Vec<_>>()
                     .join(", ")
             ),
diff --git a/src/parser/chiquito.lalrpop b/src/parser/chiquito.lalrpop
index 36b50a3c..47260539 100644
--- a/src/parser/chiquito.lalrpop
+++ b/src/parser/chiquito.lalrpop
@@ -249,8 +249,20 @@ ParseParamDecl: Statement<BigInt, Identifier> = {
 }
 
 ParseTypedId: TypedIdDecl<Identifier> = {
-    <id: Identifier> ":" <ty: Identifier> => TypedIdDecl { id, ty: Some(ty) },
-    <id: Identifier> => TypedIdDecl { id, ty: None },
+    <id: Identifier> ":" <ty: Identifier> <dim: ParseArrayDimension?> => TypedIdDecl { 
+        id, 
+        ty: Some(ty), 
+        dimension: dim 
+    },
+    <id: Identifier> <dim: ParseArrayDimension?> => TypedIdDecl { 
+        id, 
+        ty: None, 
+        dimension: dim 
+    },
+}
+
+ParseArrayDimension: Expression<BigInt, Identifier> = {
+    "[" <expr: Expression> "]" => expr,
 }
 
 // Terminals
diff --git a/src/parser/mod.rs b/src/parser/mod.rs
index c4031960..56959c97 100644
--- a/src/parser/mod.rs
+++ b/src/parser/mod.rs
@@ -112,6 +112,58 @@ mod test {
         assert_eq!(stmts.len(), 2);
     }
 
+    #[test]
+    fn test_parser_declarations() {
+        let debug_sym_ref_factory = DebugSymRefFactory::new("", "");
+        let signal_decl = lang::StatementsParser::new()
+            .parse(&debug_sym_ref_factory, "signal a: field, b, c;")
+            .unwrap();
+        assert_eq!(
+            format!("{:?}", signal_decl),
+            r#"[signal a, b, c;]"#
+        );
+
+        let typed_signal_decl = lang::StatementsParser::new()
+            .parse(&debug_sym_ref_factory, "signal a: field;")
+            .unwrap();
+        assert_eq!(
+            format!("{:?}", typed_signal_decl),
+            r#"[signal a;]"#
+        );
+
+        let var_decl = lang::StatementsParser::new()
+            .parse(&debug_sym_ref_factory, "var a: field, b, c;")
+            .unwrap();
+        assert_eq!(
+            format!("{:?}", var_decl),
+            r#"[var a, b, c;]"#
+        );
+
+        let signal_array_decl = lang::StatementsParser::new()
+            .parse(&debug_sym_ref_factory, "signal a[7];")
+            .unwrap();
+        assert_eq!(
+            format!("{:?}", signal_array_decl),
+            r#"[signal a[7];]"#
+        );
+
+        let signal_typed_array_decl = lang::StatementsParser::new()
+            .parse(&debug_sym_ref_factory, "signal a: field[7];")
+            .unwrap();
+        assert_eq!(
+            format!("{:?}", signal_typed_array_decl),
+            r#"[signal a[7];]"#
+        );
+
+        let var_array_decl = lang::StatementsParser::new()
+            .parse(&debug_sym_ref_factory, "var a[7];")
+            .unwrap();
+        assert_eq!(
+            format!("{:?}", var_array_decl),
+            r#"[var a[7];]"#
+        );
+    }
+
     #[test]
     fn test_parser_tracer() {
         let circuit = "