From b90489f0af85b5cc828624bf12088f51a59f55ff Mon Sep 17 00:00:00 2001
From: Kata Choi <kata.choi@gmail.com>
Date: Sat, 21 Sep 2024 18:22:34 +0800
Subject: [PATCH] feat: support binary operations for function return type
 (#184)

* feat: support binary operations for function return type
---
 examples/generic_repeated_array.no |  6 ++---
 src/negative_tests.rs              | 30 +++++++++++++++++++++-
 src/parser/types.rs                | 41 +++++++++---------------------
 3 files changed, 44 insertions(+), 33 deletions(-)

diff --git a/examples/generic_repeated_array.no b/examples/generic_repeated_array.no
index 2e33f0022..d169c9ba1 100644
--- a/examples/generic_repeated_array.no
+++ b/examples/generic_repeated_array.no
@@ -1,10 +1,10 @@
-fn init_arr(const LEFT: Field) -> [Field; LEFT] {
-    let arr = [0; LEFT];
+fn init_arr(const LEFT: Field) -> [Field; 1 + (LEFT * 2)] {
+    let arr = [0; 1 + (LEFT * 2)];
     return arr;
 }
 
 fn main(pub public_input: Field) -> [Field; 3] {
-    let mut arr = init_arr(3);
+    let mut arr = init_arr(1);
     for ii in 0..3 {
         arr[ii] = public_input;
     }
diff --git a/src/negative_tests.rs b/src/negative_tests.rs
index 374e90b9d..9c7013d4d 100644
--- a/src/negative_tests.rs
+++ b/src/negative_tests.rs
@@ -1,13 +1,27 @@
 use crate::{
     backends::r1cs::{R1csBn254Field, R1CS},
     circuit_writer::CircuitWriter,
-    compiler::{typecheck_next_file_inner, Sources},
+    compiler::{get_nast, typecheck_next_file_inner, Sources},
     error::{ErrorKind, Result},
     mast::Mast,
+    name_resolution::NAST,
     type_checker::TypeChecker,
     witness::CompiledCircuit,
 };
 
+fn nast_pass(code: &str) -> Result<(NAST<R1CS<R1csBn254Field>>, usize)> {
+    let mut source = Sources::new();
+    let res = get_nast(
+        None,
+        &mut source,
+        "example.no".to_string(),
+        code.to_string(),
+        0,
+    );
+
+    res
+}
+
 fn tast_pass(code: &str) -> (Result<usize>, TypeChecker<R1CS<R1csBn254Field>>, Sources) {
     let mut source = Sources::new();
     let mut tast = TypeChecker::<R1CS<R1csBn254Field>>::new();
@@ -371,3 +385,17 @@ fn test_multiplication_mismatch() {
     let res = tast_pass(code).0;
     assert!(matches!(res.unwrap_err().kind, ErrorKind::MismatchType(..)));
 }
+
+#[test]
+fn test_generic_missing_parenthesis() {
+    let code = r#"
+    fn init_arr(const LEFT: Field) -> [Field; 1 + LEFT * 2] {
+        let arr = [0; 1 + (LEFT * 2)];
+        return arr;
+    }
+    "#;
+
+    let res = nast_pass(code).err();
+    println!("{:?}", res);
+    assert!(matches!(res.unwrap().kind, ErrorKind::MissingParenthesis));
+}
diff --git a/src/parser/types.rs b/src/parser/types.rs
index acbd7d825..32fe42ad1 100644
--- a/src/parser/types.rs
+++ b/src/parser/types.rs
@@ -247,6 +247,7 @@ impl Symbolic {
                 let lhs = Symbolic::parse(lhs)?;
                 let rhs = Symbolic::parse(rhs);
 
+                // no protected flags are needed, as this is based on expression nodes which already ordered the operations
                 match op {
                     Op2::Addition => Ok(Symbolic::Add(Box::new(lhs), Box::new(rhs?))),
                     Op2::Multiplication => Ok(Symbolic::Mul(Box::new(lhs), Box::new(rhs?))),
@@ -472,43 +473,25 @@ impl Ty {
 
                 // [type; size]
                 //         ^
-                let siz_first = tokens.bump_err(ctx, ErrorKind::InvalidToken)?;
-
-                // [type; size]
-                //            ^
-                let siz_second = tokens.bump_err(ctx, ErrorKind::InvalidToken)?;
-
-                // return Array(ty, siz) if size is a number and right_paren is ]
-                match (&siz_first.kind, &siz_second.kind) {
-                    (TokenKind::BigUInt(b), TokenKind::RightBracket) => {
-                        let siz: u32 = b
-                            .try_into()
-                            .map_err(|_e| ctx.error(ErrorKind::InvalidArraySize, siz_first.span))?;
-                        let span = span.merge_with(siz_second.span);
-
+                let expr = Expr::parse(ctx, tokens)?;
+                match expr.kind {
+                    ExprKind::BigUInt(n) => {
+                        tokens.bump_expected(ctx, TokenKind::RightBracket)?;
                         Ok(Ty {
-                            kind: TyKind::Array(Box::new(ty.kind), siz),
+                            kind: TyKind::Array(Box::new(ty.kind), n.try_into().unwrap()),
                             span,
                         })
                     }
-                    // [Field; nn]
-                    // [Field; NN]
-                    //         ^^^
-                    (TokenKind::Identifier(name), TokenKind::RightBracket) => {
-                        let siz = Ident::new(name.to_string(), siz_first.span);
-                        let span = span.merge_with(siz_second.span);
-                        let sym = if is_generic_parameter(name) {
-                            Symbolic::Generic(siz)
-                        } else {
-                            Symbolic::Constant(siz)
-                        };
-
+                    _ => {
+                        tokens.bump_expected(ctx, TokenKind::RightBracket)?;
                         Ok(Ty {
-                            kind: TyKind::GenericSizedArray(Box::new(ty.kind), sym),
+                            kind: TyKind::GenericSizedArray(
+                                Box::new(ty.kind),
+                                Symbolic::parse(&expr)?,
+                            ),
                             span,
                         })
                     }
-                    _ => Err(ctx.error(ErrorKind::InvalidSymbolicSize, siz_first.span)),
                 }
             }