Assignment operators (&=, |=, +=, -=, *=, /=, %=, >>=, <<=) (#2301)

<!-- ELLIPSIS_HIDDEN -->


> [!IMPORTANT]
> Add support for assignment operators (`+=`, `-=`, etc.) in BAML,
updating compiler, parser, and VM with tests.
> 
>   - **Behavior**:
> - Adds support for assignment operators (`+=`, `-=`, `*=`, `/=`, `%=`)
in `codegen.rs` by introducing `AssignOp` handling in
`compile_statement()`.
> - Updates `typecheck_statement()` in `typecheck.rs` to handle
`AssignOp` for type checking.
>     - Extends `Stmt` enum in `stmt.rs` to include `AssignOpStmt`.
>   - **Parser**:
>     - Updates `parse_expr.rs` to parse `AssignOpStmt`.
> - Modifies `datamodel.pest` to include grammar rules for assignment
operators.
>   - **HIR**:
>     - Adds `AssignOp` to `Statement` enum in `mod.rs`.
> - Implements `AssignOp` handling in `lowering.rs` for AST to HIR
conversion.
>   - **Tests**:
> - Adds tests in `vm.rs` to verify execution of assignment operators.
> - Includes tests for each operator (`+=`, `-=`, `*=`, `/=`, `%=`) to
ensure correct VM execution.
> 
> <sup>This description was created by </sup>[<img alt="Ellipsis"
src="https://img.shields.io/badge/Ellipsis-blue?color=175173">](https://www.ellipsis.dev?ref=BoundaryML%2Fbaml&utm_source=github&utm_medium=referral)<sup>
for 08330af. You can
[customize](https://app.ellipsis.dev/BoundaryML/settings/summaries) this
summary. It will automatically update as commits are pushed.</sup>


<!-- ELLIPSIS_HIDDEN -->

Author: antoniosarosi

engine/baml-compiler/src/codegen.rs

@@ -298,11 +298,37 @@ impl<'g> HirCompiler<'g> {
                 self.track_local(name);
             }
 
-            hir::Statement::Assign { name, value } => {
+            hir::Statement::Assign { name, value, .. } => {
                 self.compile_expression(value);
                 self.emit(Instruction::StoreVar(self.locals[name]));
             }
 
+            hir::Statement::AssignOp {
+                name,
+                value,
+                assign_op,
+                ..
+            } => {
+                self.emit(Instruction::LoadVar(self.locals[name]));
+                self.compile_expression(value);
+
+                self.emit(match assign_op {
+                    hir::AssignOp::AddAssign => Instruction::BinOp(BinOp::Add),
+                    hir::AssignOp::SubAssign => Instruction::BinOp(BinOp::Sub),
+                    hir::AssignOp::MulAssign => Instruction::BinOp(BinOp::Mul),
+                    hir::AssignOp::DivAssign => Instruction::BinOp(BinOp::Div),
+                    hir::AssignOp::ModAssign => Instruction::BinOp(BinOp::Mod),
+
+                    hir::AssignOp::BitAndAssign => Instruction::BinOp(BinOp::BitAnd),
+                    hir::AssignOp::BitOrAssign => Instruction::BinOp(BinOp::BitOr),
+                    hir::AssignOp::BitXorAssign => Instruction::BinOp(BinOp::BitXor),
+                    hir::AssignOp::ShlAssign => Instruction::BinOp(BinOp::Shl),
+                    hir::AssignOp::ShrAssign => Instruction::BinOp(BinOp::Shr),
+                });
+
+                self.emit(Instruction::StoreVar(self.locals[name]));
+            }
+
             hir::Statement::DeclareAndAssign { name, value, .. } => {
                 self.compile_expression(value);
                 self.track_local(name);
@@ -1515,4 +1541,29 @@ mod tests {
             )],
         })
     }
+
+    #[test]
+    fn basic_assign_add() -> anyhow::Result<()> {
+        assert_compiles(Program {
+            source: r#"
+                fn main() -> int {
+                    let mut x = 1;
+                    x += 2;
+                    x
+                }
+            "#,
+            expected: vec![(
+                "main",
+                vec![
+                    Instruction::LoadConst(0),
+                    Instruction::LoadVar(1),
+                    Instruction::LoadConst(1),
+                    Instruction::BinOp(BinOp::Add),
+                    Instruction::StoreVar(1),
+                    Instruction::LoadVar(1),
+                    Instruction::Return,
+                ],
+            )],
+        })
+    }
 }

engine/baml-compiler/src/hir/dump.rs

@@ -3,9 +3,9 @@
 use pretty::RcDoc;
 
 use crate::hir::{
-    Arrow, BinaryOperator, Block, Class, ClassConstructorField, Enum, EnumVariant, ExprFunction,
-    Expression, Field, Hir, LlmFunction, Parameter, Statement, TypeArg, TypeM, TypeMeta,
-    UnaryOperator,
+    Arrow, AssignOp, BinaryOperator, Block, Class, ClassConstructorField, Enum, EnumVariant,
+    ExprFunction, Expression, Field, Hir, LlmFunction, Parameter, Statement, TypeArg, TypeM,
+    TypeMeta, UnaryOperator,
 };
 
 impl Hir {
@@ -117,12 +117,23 @@ impl Statement {
                 .append(RcDoc::space())
                 .append(RcDoc::text(name.clone()))
                 .append(RcDoc::text(";")),
-            Statement::Assign { name, value } => RcDoc::text(name.clone())
+            Statement::Assign { name, value, .. } => RcDoc::text(name.clone())
                 .append(RcDoc::space())
                 .append(RcDoc::text("="))
                 .append(RcDoc::space())
                 .append(value.to_doc())
                 .append(RcDoc::text(";")),
+            Statement::AssignOp {
+                name,
+                value,
+                assign_op,
+                ..
+            } => RcDoc::text(name.clone())
+                .append(RcDoc::space())
+                .append(assign_op.to_doc())
+                .append(RcDoc::space())
+                .append(value.to_doc())
+                .append(RcDoc::text(";")),
             Statement::DeclareAndAssign { name, value, .. } => RcDoc::text("var")
                 .append(RcDoc::space())
                 .append(RcDoc::text(name.clone()))
@@ -529,6 +540,23 @@ impl std::fmt::Display for UnaryOperator {
     }
 }
 
+impl std::fmt::Display for AssignOp {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.write_str(match self {
+            AssignOp::AddAssign => "+=",
+            AssignOp::SubAssign => "-=",
+            AssignOp::MulAssign => "*=",
+            AssignOp::DivAssign => "/=",
+            AssignOp::ModAssign => "%=",
+            AssignOp::BitXorAssign => "^=",
+            AssignOp::BitAndAssign => "&=",
+            AssignOp::BitOrAssign => "|=",
+            AssignOp::ShlAssign => "<<=",
+            AssignOp::ShrAssign => ">>=",
+        })
+    }
+}
+
 impl BinaryOperator {
     pub fn to_doc(&self) -> RcDoc<'static, ()> {
         RcDoc::text(self.to_string())
@@ -540,3 +568,9 @@ impl UnaryOperator {
         RcDoc::text(self.to_string())
     }
 }
+
+impl AssignOp {
+    pub fn to_doc(&self) -> RcDoc<'static, ()> {
+        RcDoc::text(self.to_string())
+    }
+}

engine/baml-compiler/src/hir/lowering.rs

@@ -298,15 +298,34 @@ impl Block {
                     expr,
                     span,
                 }) => {
-                    // Assignment statement (like let, without let) - check for if expressions in nested contexts
-                    // NOTE: Since we're not desugaring assignments, there will be no
-                    // lifted statements.
-                    let mut temp_counter = 0;
-                    let lifted_expr = Expression::from_ast(expr);
-
                     statements.push(Statement::Assign {
                         name: identifier.to_string(),
-                        value: lifted_expr,
+                        value: Expression::from_ast(expr),
+                        span: span.clone(),
+                    });
+                }
+                ast::Stmt::AssignOp(ast::AssignOpStmt {
+                    identifier,
+                    assign_op,
+                    expr,
+                    span,
+                }) => {
+                    statements.push(Statement::AssignOp {
+                        name: identifier.to_string(),
+                        assign_op: match assign_op {
+                            ast::AssignOp::AddAssign => hir::AssignOp::AddAssign,
+                            ast::AssignOp::SubAssign => hir::AssignOp::SubAssign,
+                            ast::AssignOp::MulAssign => hir::AssignOp::MulAssign,
+                            ast::AssignOp::DivAssign => hir::AssignOp::DivAssign,
+                            ast::AssignOp::ModAssign => hir::AssignOp::ModAssign,
+                            ast::AssignOp::BitXorAssign => hir::AssignOp::BitXorAssign,
+                            ast::AssignOp::BitAndAssign => hir::AssignOp::BitAndAssign,
+                            ast::AssignOp::BitOrAssign => hir::AssignOp::BitOrAssign,
+                            ast::AssignOp::ShlAssign => hir::AssignOp::ShlAssign,
+                            ast::AssignOp::ShrAssign => hir::AssignOp::ShrAssign,
+                        },
+                        value: Expression::from_ast(expr),
+                        span: span.clone(),
                     });
                 }
                 ast::Stmt::Let(ast::LetStmt {

engine/baml-compiler/src/hir/mod.rs

@@ -334,6 +334,13 @@ pub enum Statement {
     Assign {
         name: String,
         value: Expression,
+        span: Span,
+    },
+    AssignOp {
+        name: String,
+        assign_op: AssignOp,
+        value: Expression,
+        span: Span,
     },
     /// Declare and assign a mutable reference in one statement.
     DeclareAndAssign {
@@ -368,6 +375,30 @@ pub enum Statement {
     },
 }
 
+#[derive(Debug, Clone, Copy)]
+pub enum AssignOp {
+    /// The `+=` operator (addition)
+    AddAssign,
+    /// The `-=` operator (subtraction)
+    SubAssign,
+    /// The `*=` operator (multiplication)
+    MulAssign,
+    /// The `/=` operator (division)
+    DivAssign,
+    /// The `%=` operator (modulus)
+    ModAssign,
+    /// The `^=` operator (bitwise xor)
+    BitXorAssign,
+    /// The `&=` operator (bitwise and)
+    BitAndAssign,
+    /// The `|=` operator (bitwise or)
+    BitOrAssign,
+    /// The `<<=` operator (shift left)
+    ShlAssign,
+    /// The `>>=` operator (shift right)
+    ShrAssign,
+}
+
 /// Expressions
 #[derive(Clone, Debug)]
 pub enum Expression {

engine/baml-compiler/src/thir/typecheck.rs

@@ -344,7 +344,9 @@ fn typecheck_statement(
                 span: span.clone(),
             })
         }
-        hir::Statement::Assign { name, value } => {
+        // TODO: assign op needs more type checking?
+        hir::Statement::Assign { name, value, .. }
+        | hir::Statement::AssignOp { name, value, .. } => {
             let typed_value = typecheck_expression(value, context, diagnostics);
 
             // validate/update type.

engine/baml-lib/ast/src/ast.rs

@@ -36,7 +36,7 @@ pub use identifier::{Identifier, RefIdentifier};
 pub use indentation_type::IndentationType;
 pub use internal_baml_diagnostics::Span;
 pub use newline_type::NewlineType;
-pub use stmt::{AssignStmt, ForLoopStmt, LetStmt, Stmt};
+pub use stmt::{AssignOp, AssignOpStmt, AssignStmt, ForLoopStmt, LetStmt, Stmt};
 pub use template_string::TemplateString;
 pub use top::Top;
 pub use traits::{WithAttributes, WithDocumentation, WithIdentifier, WithName, WithSpan};

engine/baml-lib/ast/src/ast/stmt.rs

@@ -17,6 +17,38 @@ pub struct AssignStmt {
     pub span: Span,
 }
 
+#[derive(Debug, Clone)]
+pub struct AssignOpStmt {
+    pub identifier: Identifier,
+    pub assign_op: AssignOp,
+    pub expr: Expression,
+    pub span: Span,
+}
+
+#[derive(Debug, Clone, Copy)]
+pub enum AssignOp {
+    /// The `+=` operator (addition)
+    AddAssign,
+    /// The `-=` operator (subtraction)
+    SubAssign,
+    /// The `*=` operator (multiplication)
+    MulAssign,
+    /// The `/=` operator (division)
+    DivAssign,
+    /// The `%=` operator (modulus)
+    ModAssign,
+    /// The `^=` operator (bitwise xor)
+    BitXorAssign,
+    /// The `&=` operator (bitwise and)
+    BitAndAssign,
+    /// The `|=` operator (bitwise or)
+    BitOrAssign,
+    /// The `<<=` operator (shift left)
+    ShlAssign,
+    /// The `>>=` operator (shift right)
+    ShrAssign,
+}
+
 #[derive(Debug, Clone)]
 pub struct ForLoopStmt {
     pub identifier: Identifier,
@@ -33,6 +65,24 @@ pub enum Stmt {
     /// Expression with trailing semicolon.
     Expression(Expression),
     Assign(AssignStmt),
+    AssignOp(AssignOpStmt),
+}
+
+impl fmt::Display for AssignOp {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.write_str(match self {
+            AssignOp::AddAssign => "+=",
+            AssignOp::SubAssign => "-=",
+            AssignOp::MulAssign => "*=",
+            AssignOp::DivAssign => "/=",
+            AssignOp::ModAssign => "%=",
+            AssignOp::BitAndAssign => "&=",
+            AssignOp::BitOrAssign => "|=",
+            AssignOp::BitXorAssign => "^=",
+            AssignOp::ShlAssign => "<<=",
+            AssignOp::ShrAssign => ">>=",
+        })
+    }
 }
 
 impl fmt::Display for Stmt {
@@ -42,6 +92,9 @@ impl fmt::Display for Stmt {
             Stmt::ForLoop(stmt) => write!(f, "for {} in {}", stmt.identifier, stmt.iterator)?,
             Stmt::Expression(expr) => write!(f, "{expr}")?,
             Stmt::Assign(stmt) => write!(f, "{} = {}", stmt.identifier, stmt.expr)?,
+            Stmt::AssignOp(stmt) => {
+                write!(f, "{} {} {}", stmt.identifier, stmt.assign_op, stmt.expr)?
+            }
         }
         Ok(())
     }
@@ -79,6 +132,7 @@ impl Stmt {
             Stmt::ForLoop(stmt) => &stmt.identifier,
             Stmt::Expression(expr) => panic!("expressions don't have identifiers"),
             Stmt::Assign(stmt) => &stmt.identifier,
+            Stmt::AssignOp(stmt) => &stmt.identifier,
         }
     }
 
@@ -88,6 +142,7 @@ impl Stmt {
             Stmt::ForLoop(stmt) => &stmt.span,
             Stmt::Expression(expr) => expr.span(),
             Stmt::Assign(stmt) => &stmt.span,
+            Stmt::AssignOp(stmt) => &stmt.span,
         }
     }
 

engine/baml-lib/ast/src/parser/datamodel.pest

@@ -323,12 +323,26 @@ expr_block = { BLOCK_OPEN ~ NEWLINE? ~ (stmt | comment_block | empty_lines)* ~ e
 // With operator overloading (implicit function call) + side effects.
 //
 // Just allow any expr and fix the parser.
-stmt = { (((let_expr | assign_stmt | fn_app | generic_fn_app) ~ SEMICOLON) | for_loop | if_expression) ~ trailing_comment? ~ NEWLINE? }
+stmt = { (((let_expr | assign_op_stmt | assign_stmt | fn_app | generic_fn_app) ~ SEMICOLON) | for_loop | if_expression) ~ trailing_comment? ~ NEWLINE? }
 
 // Let-binding statement.
 let_expr = { "let" ~ MUT_KEYWORD? ~ identifier ~ "=" ~ expression }
 
 assign_stmt = { identifier ~ "=" ~ expression }
+assign_op_stmt = { identifier ~ assign_op ~ expression }
+
+assign_op = _{ BIT_SHL_ASSIGN | BIT_SHR_ASSIGN | ADD_ASSIGN | SUB_ASSIGN | MUL_ASSIGN | DIV_ASSIGN | MOD_ASSIGN | BIT_AND_ASSIGN | BIT_OR_ASSIGN | BIT_XOR_ASSIGN }
+
+ADD_ASSIGN = { "+=" }
+SUB_ASSIGN = { "-=" }
+MUL_ASSIGN = { "*=" }
+DIV_ASSIGN = { "/=" }
+MOD_ASSIGN = { "%=" }
+BIT_AND_ASSIGN = { "&=" }
+BIT_OR_ASSIGN = { "|=" }
+BIT_XOR_ASSIGN = { "^=" }
+BIT_SHL_ASSIGN = { "<<=" }
+BIT_SHR_ASSIGN = { ">>=" }
 
 fn_args = { expression? ~ ("," ~ expression)* }