Allow type declaration after ":" in let / let mut

Author: fkettelhoit

src/ast.rs

@@ -269,9 +269,9 @@ impl<T> Stmt<T> {
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 pub enum StmtEnum<T> {
     /// Let expression, binds variables to exprs.
-    Let(Pattern<T>, Expr<T>),
+    Let(Pattern<T>, Option<Type>, Expr<T>),
     /// Mutable let expression, bind a single variable to an expr.
-    LetMut(String, Expr<T>),
+    LetMut(String, Option<Type>, Expr<T>),
     /// Assignment of a (previously as mutable declared) variable.
     VarAssign(String, Expr<T>),
     /// Assignment of an index in a (mutable) array.

src/check.rs

@@ -627,12 +627,15 @@ impl UntypedStmt {
     ) -> Result<TypedStmt, TypeErrors> {
         let meta = self.meta;
         match &self.inner {
-            ast::StmtEnum::Let(pattern, binding) => {
+            ast::StmtEnum::Let(pattern, ty, binding) => {
                 match binding.type_check(top_level_defs, env, fns, defs) {
-                    Ok(binding) => {
+                    Ok(mut binding) => {
+                        if let Some(ty) = ty {
+                            check_type(&mut binding, ty)?;
+                        }
                         let pattern =
                             pattern.type_check(env, fns, defs, Some(binding.ty.clone()))?;
-                        Ok(Stmt::new(StmtEnum::Let(pattern, binding), meta))
+                        Ok(Stmt::new(StmtEnum::Let(pattern, ty.clone(), binding), meta))
                     }
                     Err(mut errors) => {
                         if let Err(e) = pattern.type_check(env, fns, defs, None) {
@@ -642,9 +645,12 @@ impl UntypedStmt {
                     }
                 }
             }
-            ast::StmtEnum::LetMut(identifier, binding) => {
+            ast::StmtEnum::LetMut(identifier, ty, binding) => {
                 match binding.type_check(top_level_defs, env, fns, defs) {
                     Ok(mut binding) => {
+                        if let Some(ty) = ty {
+                            check_type(&mut binding, ty)?;
+                        }
                         if binding.ty == Type::Unsigned(UnsignedNumType::Unspecified)
                             || binding.ty == Type::Signed(SignedNumType::Unspecified)
                         {
@@ -662,7 +668,7 @@ impl UntypedStmt {
                             (Some(binding.ty.clone()), Mutability::Mutable),
                         );
                         Ok(Stmt::new(
-                            StmtEnum::LetMut(identifier.clone(), binding),
+                            StmtEnum::LetMut(identifier.clone(), ty.clone(), binding),
                             meta,
                         ))
                     }

src/compile.rs

@@ -374,13 +374,13 @@ impl TypedStmt {
         circuit: &mut CircuitBuilder,
     ) -> Vec<GateIndex> {
         match &self.inner {
-            StmtEnum::Let(pattern, binding) => {
+            StmtEnum::Let(pattern, _, binding) => {
                 let binding = binding.compile(prg, env, circuit);
                 pattern.compile(&binding, prg, env, circuit);
                 vec![]
             }
             StmtEnum::Expr(expr) => expr.compile(prg, env, circuit),
-            StmtEnum::LetMut(identifier, binding) => {
+            StmtEnum::LetMut(identifier, _, binding) => {
                 let binding = binding.compile(prg, env, circuit);
                 env.let_in_current_scope(identifier.clone(), binding);
                 vec![]

src/parse.rs

@@ -390,11 +390,17 @@ impl Parser {
             if self.next_matches(&TokenEnum::KeywordMut).is_some() {
                 // let mut <identifier> = <binding>;
                 let (identifier, _) = self.expect_identifier()?;
+                let ty = if self.next_matches(&TokenEnum::Colon).is_some() {
+                    let (ty, _) = self.parse_type()?;
+                    Some(ty)
+                } else {
+                    None
+                };
                 self.expect(&TokenEnum::Eq)?;
                 if let Ok(binding) = self.parse_expr() {
                     let meta = join_meta(meta, binding.meta);
                     self.expect(&TokenEnum::Semicolon)?;
-                    return Ok(Stmt::new(StmtEnum::LetMut(identifier, binding), meta));
+                    return Ok(Stmt::new(StmtEnum::LetMut(identifier, ty, binding), meta));
                 } else {
                     self.push_error_for_next(ParseErrorEnum::ExpectedStmt);
                     self.consume_until_one_of(&[TokenEnum::Semicolon]);
@@ -405,11 +411,17 @@ impl Parser {
             } else {
                 // let <pattern> = <binding>;
                 let pattern = self.parse_pattern()?;
+                let ty = if self.next_matches(&TokenEnum::Colon).is_some() {
+                    let (ty, _) = self.parse_type()?;
+                    Some(ty)
+                } else {
+                    None
+                };
                 self.expect(&TokenEnum::Eq)?;
                 if let Ok(binding) = self.parse_expr() {
                     let meta = join_meta(meta, binding.meta);
                     self.expect(&TokenEnum::Semicolon)?;
-                    return Ok(Stmt::new(StmtEnum::Let(pattern, binding), meta));
+                    return Ok(Stmt::new(StmtEnum::Let(pattern, ty, binding), meta));
                 } else {
                     self.consume_until_one_of(&[TokenEnum::Semicolon]);
                     self.expect(&TokenEnum::Semicolon)?;

tests/compile.rs

@@ -2423,3 +2423,66 @@ pub fn main(array: [u16; PARTIES]) -> u16 {
     );
     Ok(())
 }
+
+#[test]
+fn type_declaration_in_let() -> Result<(), Error> {
+    let prg = "
+pub fn main(a: u16, b: u16) -> u16 {
+    let c: u16 = 6;
+    a + b + c
+}
+";
+    let compiled = compile(prg).map_err(|e| pretty_print(e, prg))?;
+    let mut eval = compiled.evaluator();
+    eval.set_u16(2);
+    eval.set_u16(4);
+    let output = eval.run().map_err(|e| pretty_print(e, prg))?;
+    assert_eq!(
+        u16::try_from(output).map_err(|e| pretty_print(e, prg))?,
+        2 + 4 + 6
+    );
+    Ok(())
+}
+
+#[test]
+fn type_declaration_in_let_mut() -> Result<(), Error> {
+    let prg = "
+pub fn main(a: u16, b: u16) -> u16 {
+    let mut result: u16 = 6;
+    result += a + b;
+    result
+}
+";
+    let compiled = compile(prg).map_err(|e| pretty_print(e, prg))?;
+    let mut eval = compiled.evaluator();
+    eval.set_u16(2);
+    eval.set_u16(4);
+    let output = eval.run().map_err(|e| pretty_print(e, prg))?;
+    assert_eq!(
+        u16::try_from(output).map_err(|e| pretty_print(e, prg))?,
+        2 + 4 + 6
+    );
+    Ok(())
+}
+
+#[test]
+fn type_declaration_of_array_in_let_mut() -> Result<(), Error> {
+    let prg = "
+pub fn main(a: u16, b: u16) -> u16 {
+    let mut result: [u16; 2] = [3, 3];
+    result[0] = result[0] + a;
+    result[1] = result[1] + b;
+    result[0] + result[1]
+}
+";
+    let compiled = compile(prg).map_err(|e| pretty_print(e, prg))?;
+    let mut eval = compiled.evaluator();
+    eval.set_u16(2);
+    eval.set_u16(4);
+    let output = eval.run().map_err(|e| pretty_print(e, prg))?;
+    assert_eq!(
+        u16::try_from(output).map_err(|e| pretty_print(e, prg))?,
+        2 + 4 + 3 + 3
+    );
+    Ok(())
+}