Merge pull request #225 from bufferhe4d/plusequal

Add += functionality
zksecurity · Nov 13, 2024 · b17ea20 · b17ea20
2 parents 8c0ca9d + 7b21812
commit b17ea20
Show file tree

Hide file tree

Showing 8 changed files with 157 additions and 20 deletions.
diff --git a/examples/assignment.no b/examples/assignment.no
@@ -15,7 +15,7 @@ fn main(pub xx: Field) {
     };
 
     // ideally: thing.xx += 1;
-    thing.xx = thing.xx + 1;
+    thing.xx += 1;
 
     try_to_mutate(thing);
 

diff --git a/examples/augmented_assign.no b/examples/augmented_assign.no
@@ -0,0 +1,16 @@
+fn main(pub xx: Field, yy: Field) -> Field {
+    let mut zz = xx + yy;
+    let mut zz2 = xx + yy;
+
+    zz += 1;
+    zz2 = zz2 + 1;
+
+    zz *= zz;
+    zz2 = zz2 * zz2;
+
+    zz -= yy;
+    zz2 = zz2 - yy;
+    assert_eq(zz, zz2);
+
+    return zz;
+}
diff --git a/examples/fixture/asm/kimchi/augmented_assign.asm b/examples/fixture/asm/kimchi/augmented_assign.asm
@@ -0,0 +1,30 @@
+@ noname.0.7.0
+@ public inputs: 2
+
+DoubleGeneric<1>
+DoubleGeneric<1>
+DoubleGeneric<1,1,-1>
+DoubleGeneric<1,1,-1>
+DoubleGeneric<1,0,-1,0,1>
+DoubleGeneric<1,0,-1,0,1>
+DoubleGeneric<0,0,-1,1>
+DoubleGeneric<0,0,-1,1>
+DoubleGeneric<1,1>
+DoubleGeneric<1,1,-1>
+DoubleGeneric<1,1>
+DoubleGeneric<1,1,-1>
+DoubleGeneric<1,-1>
+DoubleGeneric<1,-1>
+(0,0) -> (13,0)
+(1,0) -> (2,0) -> (3,0)
+(2,1) -> (3,1) -> (8,0) -> (10,0)
+(2,2) -> (4,0)
+(3,2) -> (5,0)
+(4,2) -> (6,0) -> (6,1)
+(5,2) -> (7,0) -> (7,1)
+(6,2) -> (9,0)
+(7,2) -> (11,0)
+(8,1) -> (9,1)
+(9,2) -> (12,0) -> (13,1)
+(10,1) -> (11,1)
+(11,2) -> (12,1)
diff --git a/examples/fixture/asm/r1cs/augmented_assign.asm b/examples/fixture/asm/r1cs/augmented_assign.asm
@@ -0,0 +1,7 @@
+@ noname.0.7.0
+@ public inputs: 2
+
+v_4 == (v_2 + v_3 + 1) * (v_2 + v_3 + 1)
+v_5 == (v_2 + v_3 + 1) * (v_2 + v_3 + 1)
+-1 * v_3 + v_5 == (-1 * v_3 + v_4) * (1)
+-1 * v_3 + v_4 == (v_1) * (1)
diff --git a/examples/iterate.no b/examples/iterate.no
@@ -19,7 +19,7 @@ fn House.windows(house: House) -> Field {
     let mut windows_count = 0;
     for room in house.rooms {
         // ideally: windows +=
-        windows_count = windows_count + room.windows();
+        windows_count += room.windows();
     }
 
     return windows_count;

diff --git a/src/lexer/mod.rs b/src/lexer/mod.rs
@@ -158,6 +158,9 @@ pub enum TokenKind {
     DoublePipe,         // ||
     Exclamation,        // !
     Question,           // ?
+    PlusEqual,          // +=
+    MinusEqual,         // -=
+    StarEqual,          // *=
                         //    Literal,               // "thing"
 }
 
@@ -199,6 +202,9 @@ impl Display for TokenKind {
             DoublePipe => "`||`",
             Exclamation => "`!`",
             Question => "`?`",
+            PlusEqual => "`+=`",
+            MinusEqual => "`-=`",
+            StarEqual => "`*=`",
             //            TokenType::Literal => "`\"something\"",
         };
 
@@ -387,19 +393,38 @@ impl Token {
                     }
                 }
                 '+' => {
-                    tokens.push(TokenKind::Plus.new_token(ctx, 1));
-                }
-                '-' => {
                     let next_c = chars.peek();
-                    if matches!(next_c, Some(&'>')) {
-                        tokens.push(TokenKind::RightArrow.new_token(ctx, 2));
+                    if matches!(next_c, Some(&'=')) {
+                        tokens.push(TokenKind::PlusEqual.new_token(ctx, 2));
                         chars.next();
                     } else {
-                        tokens.push(TokenKind::Minus.new_token(ctx, 1));
+                        tokens.push(TokenKind::Plus.new_token(ctx, 1));
+                    }
+                }
+                '-' => {
+                    let next_c = chars.peek();
+                    match next_c {
+                        Some(&'>') => {
+                            tokens.push(TokenKind::RightArrow.new_token(ctx, 2));
+                            chars.next();
+                        }
+                        Some(&'=') => {
+                            tokens.push(TokenKind::MinusEqual.new_token(ctx, 2));
+                            chars.next();
+                        }
+                        _ => {
+                            tokens.push(TokenKind::Minus.new_token(ctx, 1));
+                        }
                     }
                 }
                 '*' => {
-                    tokens.push(TokenKind::Star.new_token(ctx, 1));
+                    let next_c = chars.peek();
+                    if matches!(next_c, Some(&'=')) {
+                        tokens.push(TokenKind::StarEqual.new_token(ctx, 2));
+                        chars.next();
+                    } else {
+                        tokens.push(TokenKind::Star.new_token(ctx, 1));
+                    }
                 }
                 '&' => {
                     let next_c = chars.peek();

diff --git a/src/parser/expr.rs b/src/parser/expr.rs
@@ -428,12 +428,17 @@ impl Expr {
         // the expression `self` we just parsed.
         // warning: ALL of the rules here should make use of `self`.
         let lhs = match tokens.peek() {
-            // assignment
+            // assignment or augmented assignment
+            // lhs = rhs | lhs += rhs | lhs -= rhs | lhs *= rhs
             Some(Token {
-                kind: TokenKind::Equal,
+                kind:
+                    TokenKind::Equal
+                    | TokenKind::PlusEqual
+                    | TokenKind::MinusEqual
+                    | TokenKind::StarEqual,
                 span,
             }) => {
-                tokens.bump(ctx); // =
+                let token = tokens.bump(ctx).unwrap();
 
                 // sanitize
                 if !matches!(
@@ -451,14 +456,49 @@ impl Expr {
                 let rhs = Expr::parse(ctx, tokens)?;
                 let span = self.span.merge_with(rhs.span);
 
-                Expr::new(
-                    ctx,
-                    ExprKind::Assignment {
-                        lhs: Box::new(self),
-                        rhs: Box::new(rhs),
-                    },
-                    span,
-                )
+                match token.kind {
+                    // regular assignment
+                    TokenKind::Equal => Expr::new(
+                        ctx,
+                        ExprKind::Assignment {
+                            lhs: Box::new(self),
+                            rhs: Box::new(rhs),
+                        },
+                        span,
+                    ),
+                    // augmented assignments
+                    TokenKind::PlusEqual | TokenKind::MinusEqual | TokenKind::StarEqual => {
+                        let op = match token.kind {
+                            TokenKind::PlusEqual => Op2::Addition,
+                            TokenKind::MinusEqual => Op2::Subtraction,
+                            TokenKind::StarEqual => Op2::Multiplication,
+                            _ => unreachable!(),
+                        };
+
+                        // create the binary operation: lhs <op> rhs
+                        let binary_op = Expr::new(
+                            ctx,
+                            ExprKind::BinaryOp {
+                                op,
+                                lhs: Box::new(self.clone()),
+                                rhs: Box::new(rhs),
+                                protected: false,
+                            },
+                            span,
+                        );
+
+                        // create the assignment with rhs obtained above: lhs = (binary_op)
+                        Expr::new(
+                            ctx,
+                            ExprKind::Assignment {
+                                lhs: Box::new(self),
+                                rhs: Box::new(binary_op),
+                            },
+                            span,
+                        )
+                    }
+                    _ => unreachable!(),
+                }
             }
 
             // binary operation

diff --git a/src/tests/examples.rs b/src/tests/examples.rs
@@ -498,6 +498,25 @@ fn test_assignment(#[case] backend: BackendKind) -> miette::Result<()> {
     Ok(())
 }
 
+#[rstest]
+#[case::kimchi_vesta(BackendKind::KimchiVesta(KimchiVesta::new(false)))]
+#[case::r1cs(BackendKind::R1csBls12_381(R1CS::new()))]
+fn test_augmented_assign(#[case] backend: BackendKind) -> miette::Result<()> {
+    let private_inputs = r#"{"yy": "2"}"#;
+    let public_inputs = r#"{"xx": "3"}"#;
+    let expected_public_output = vec!["34"];
+
+    test_file(
+        "augmented_assign",
+        public_inputs,
+        private_inputs,
+        expected_public_output,
+        backend,
+    )?;
+
+    Ok(())
+}
+
 #[rstest]
 #[case::kimchi_vesta(BackendKind::KimchiVesta(KimchiVesta::new(false)))]
 #[case::r1cs(BackendKind::R1csBls12_381(R1CS::new()))]