Feat: Assignment to private instance properties with += / &&= / etc

crates/oxc_transformer/src/es2022/class_properties.rs

@@ -76,7 +76,7 @@ use indexmap::IndexMap;
 use rustc_hash::FxHasher;
 use serde::Deserialize;
 
-use oxc_allocator::{Address, Box as ArenaBox, GetAddress};
+use oxc_allocator::{Address, Box as ArenaBox, CloneIn, GetAddress};
 use oxc_ast::{ast::*, NONE};
 use oxc_data_structures::stack::{NonEmptyStack, SparseStack};
 use oxc_span::SPAN;
@@ -1056,16 +1056,27 @@ impl<'a, 'ctx> ClassProperties<'a, 'ctx> {
             ))
         } else {
             // `_classPrivateFieldGet(_prop, this)`
-            let arguments = [Argument::from(prop_ident), Argument::from(field_expr.object)];
-            let arguments = ctx.ast.vec_from_array(arguments);
-            self.ctx.helper_call_expr(Helper::ClassPrivateFieldGet, field_expr.span, arguments, ctx)
+            self.create_private_field_get(prop_ident, field_expr.object, field_expr.span, ctx)
         }
     }
 
     /// Transform assignment to private field.
     ///
-    /// Instance prop: `this.#prop = value` -> `_classPrivateFieldSet(_prop, this, value)`
-    /// Static prop: `this.#prop = value` -> `_prop._ = _assertClassBrand(Class, this, _prop)`
+    /// Instance prop:
+    /// * `this.#prop = value`
+    ///   -> `_classPrivateFieldSet(_prop, this, value)`
+    /// * `this.#prop += value`
+    ///   -> `_classPrivateFieldSet(_prop, this, _classPrivateFieldGet(_prop, this) + value)`
+    /// * `this.#prop &&= value`
+    ///   -> `_classPrivateFieldGet(_prop, this) && _classPrivateFieldSet(_prop, this, value)`
+    ///
+    /// Static prop:
+    /// * `this.#prop = value`
+    ///   -> `_prop._ = _assertClassBrand(Class, this, _prop)`
+    /// * `this.#prop += value`
+    ///   -> `_prop._ = _assertClassBrand(Class, this, _assertClassBrand(Class, this, _prop)._ + value)`
+    /// * `this.#prop &&= value`
+    ///   -> `_assertClassBrand(C, this, _prop)._ && (_prop._ = _assertClassBrand(C, this, value))`
     //
     // `#[inline]` so that compiler sees that `expr` is an `Expression::AssignmentExpression`
     // with `AssignmentTarget::PrivateFieldExpression` on its left.
@@ -1086,9 +1097,6 @@ impl<'a, 'ctx> ClassProperties<'a, 'ctx> {
         expr: &mut Expression<'a>,
         ctx: &mut TraverseCtx<'a>,
     ) {
-        // TODO: Handle operators other than `=`
-        // if assign_expr.operator != AssignmentOperator::Assign {}
-
         let Expression::AssignmentExpression(assign_expr) = expr else { unreachable!() };
         let AssignmentTarget::PrivateFieldExpression(field_expr) = &mut assign_expr.left else {
             unreachable!()
@@ -1100,6 +1108,9 @@ impl<'a, 'ctx> ClassProperties<'a, 'ctx> {
         let prop_ident = prop.binding.create_read_expression(ctx);
 
         if prop.is_static {
+            // Static private property
+            // TODO: Handle operators other than `=`
+
             // `_prop._ = _assertClassBrand(Class, this, _prop)`
             // TODO: Ensure there are tests for nested classes with references to private static props
             // of outer class inside inner class, to make sure we're getting the right `class_name_binding`.
@@ -1119,7 +1130,7 @@ impl<'a, 'ctx> ClassProperties<'a, 'ctx> {
                 false,
             ));
         } else {
-            // `_classPrivateFieldSet(_prop, this, value)`
+            // Instance private property
             let assign_expr = match ctx.ast.move_expression(expr) {
                 Expression::AssignmentExpression(assign_expr) => assign_expr.unbox(),
                 _ => unreachable!(),
@@ -1129,20 +1140,92 @@ impl<'a, 'ctx> ClassProperties<'a, 'ctx> {
                 _ => unreachable!(),
             };
 
-            let arguments = ctx.ast.vec_from_array([
-                Argument::from(prop_ident),
-                Argument::from(field_expr.object),
-                Argument::from(assign_expr.right),
-            ]);
-            *expr = self.ctx.helper_call_expr(
-                Helper::ClassPrivateFieldSet,
-                assign_expr.span,
-                arguments,
-                ctx,
-            );
+            if assign_expr.operator == AssignmentOperator::Assign {
+                // `_classPrivateFieldSet(_prop, this, value)`
+                *expr = self.create_private_field_set(
+                    prop_ident,
+                    field_expr.object,
+                    assign_expr.right,
+                    assign_expr.span,
+                    ctx,
+                );
+            } else {
+                let prop_ident2 = prop.binding.create_read_expression(ctx);
+                // TODO: Don't use `clone_in` - need a temp var if `field_expr.object` is not a simple value
+                let get_call = self.create_private_field_get(
+                    prop_ident,
+                    field_expr.object.clone_in(ctx.ast.allocator),
+                    SPAN,
+                    ctx,
+                );
+
+                if let Some(operator) = assign_expr.operator.to_binary_operator() {
+                    // `_classPrivateFieldSet(_prop, this, _classPrivateFieldGet(_prop, this) + value)`
+                    let value =
+                        ctx.ast.expression_binary(SPAN, get_call, operator, assign_expr.right);
+                    *expr = self.create_private_field_set(
+                        prop_ident2,
+                        field_expr.object,
+                        value,
+                        assign_expr.span,
+                        ctx,
+                    );
+                } else if let Some(operator) = assign_expr.operator.to_logical_operator() {
+                    // `_classPrivateFieldGet(_prop, this) && _classPrivateFieldSet(_prop, this, value)`
+                    let set_call = self.create_private_field_set(
+                        prop_ident2,
+                        field_expr.object,
+                        assign_expr.right,
+                        SPAN,
+                        ctx,
+                    );
+                    *expr =
+                        ctx.ast.expression_logical(assign_expr.span, get_call, operator, set_call);
+                } else {
+                    // The above covers all types of `AssignmentOperator`
+                    unreachable!();
+                }
+            }
         }
     }
 
+    /// `_classPrivateFieldGet(_prop, this)`
+    fn create_private_field_get(
+        &mut self,
+        prop_ident: Expression<'a>,
+        object: Expression<'a>,
+        span: Span,
+        ctx: &mut TraverseCtx<'a>,
+    ) -> Expression<'a> {
+        self.ctx.helper_call_expr(
+            Helper::ClassPrivateFieldGet,
+            span,
+            ctx.ast.vec_from_array([Argument::from(prop_ident), Argument::from(object)]),
+            ctx,
+        )
+    }
+
+    /// `_classPrivateFieldSet(_prop, this, value)`
+    fn create_private_field_set(
+        &mut self,
+        prop_ident: Expression<'a>,
+        object: Expression<'a>,
+        value: Expression<'a>,
+        span: Span,
+        ctx: &mut TraverseCtx<'a>,
+    ) -> Expression<'a> {
+        self.ctx.helper_call_expr(
+            Helper::ClassPrivateFieldSet,
+            span,
+            ctx.ast.vec_from_array([
+                Argument::from(prop_ident),
+                Argument::from(object),
+                Argument::from(value),
+            ]),
+            ctx,
+        )
+    }
+
     /// Lookup details of private property referred to by `ident`.
     fn lookup_private_property(
         &self,