refactor(transformer): exponentiation transform: split into 2 paths (#6316)

Exponentiation transform has 2 paths for when left-hand side of `**=` is an identifier, and when it's a member expression. But these two paths are mixed up together in the `convert_assignment_expression`, `explode`, and `get_obj_ref` functions, with each branching on the same condition (is it an identifier or a member expression?).

Refactor to separate out these 2 code paths, and make the logic easier to follow.

crates/oxc_transformer/src/es2016/exponentiation_operator.rs

@@ -58,7 +58,7 @@ impl<'a, 'ctx> ExponentiationOperator<'a, 'ctx> {
 }
 
 impl<'a, 'ctx> Traverse<'a> for ExponentiationOperator<'a, 'ctx> {
-    // NOTE: Bail bigint arguments to `Math.pow`, which are runtime errors.
+    // Note: Do not transform to `Math.pow` with BigInt arguments - that's a runtime error
     fn enter_expression(&mut self, expr: &mut Expression<'a>, ctx: &mut TraverseCtx<'a>) {
         match expr {
             // `left ** right`
@@ -80,7 +80,21 @@ impl<'a, 'ctx> Traverse<'a> for ExponentiationOperator<'a, 'ctx> {
                     return;
                 }
 
-                self.convert_assignment_expression(expr, ctx);
+                match &assign_expr.left {
+                    AssignmentTarget::AssignmentTargetIdentifier(_) => {
+                        self.convert_assignment_to_identifier(expr, ctx);
+                    }
+                    // Note: We do not match `AssignmentTarget::PrivateFieldExpression` here.
+                    // From Babel: "We can't generate property ref for private name, please install
+                    // `@babel/plugin-transform-class-properties`".
+                    // TODO: Ensure this plugin interacts correctly with class private properties
+                    // transform, so the property is transformed before this transform.
+                    AssignmentTarget::StaticMemberExpression(_)
+                    | AssignmentTarget::ComputedMemberExpression(_) => {
+                        self.convert_assignment_to_member_expression(expr, ctx);
+                    }
+                    _ => {}
+                }
             }
             _ => {}
         }
@@ -98,9 +112,75 @@ impl<'a, 'ctx> ExponentiationOperator<'a, 'ctx> {
         *expr = Self::math_pow(binary_expr.left, binary_expr.right, ctx);
     }
 
-    /// Convert `AssignmentExpression`.
+    /// Convert `AssignmentExpression` where assignee is an identifier.
-    // `left **= right` -> `left = Math.pow(left, right)`
+    ///
-    fn convert_assignment_expression(
+    /// `left **= right` transformed to:
+    /// * If `left` is a bound symbol:
+    ///   -> `left = Math.pow(left, right)`
+    /// * If `left` is unbound:
+    ///   -> `var _left; _left = left, left = Math.pow(_left, right);`
+    ///
+    /// Temporary variable `_left` is to avoid side-effects of getting `left` from running twice.
+    fn convert_assignment_to_identifier(
+        &mut self,
+        expr: &mut Expression<'a>,
+        ctx: &mut TraverseCtx<'a>,
+    ) {
+        let Expression::AssignmentExpression(assign_expr) = expr else { unreachable!() };
+        let assign_target = &mut assign_expr.left;
+        let AssignmentTarget::AssignmentTargetIdentifier(ident) = assign_target else {
+            unreachable!()
+        };
+
+        let mut nodes = ctx.ast.vec();
+
+        let symbol_id = ctx.symbols().get_reference(ident.reference_id().unwrap()).symbol_id();
+        // Make sure side-effects of evaluating `left` only happen once
+        let uid = if let Some(symbol_id) = symbol_id {
+            // This variable is declared in scope so evaluating it multiple times can't trigger a getter.
+            // No need for a temp var.
+            ctx.ast.expression_from_identifier_reference(ctx.create_bound_reference_id(
+                SPAN,
+                ident.name.clone(),
+                symbol_id,
+                ReferenceFlags::Write,
+            ))
+        } else {
+            // Unbound reference. Could possibly trigger a getter so we need to only evaluate it once.
+            // Assign to a temp var.
+            let reference = ctx.ast.expression_from_identifier_reference(
+                ctx.create_unbound_reference_id(SPAN, ident.name.clone(), ReferenceFlags::Read),
+            );
+            self.add_new_reference(reference, &mut nodes, ctx)
+        };
+
+        let reference = ctx.ast.move_assignment_target(assign_target);
+
+        *expr = Self::create_replacement(assign_expr, reference, uid, nodes, ctx);
+    }
+
+    /// Convert `AssignmentExpression` where assignee is a member expression.
+    ///
+    /// `obj.prop **= right`
+    /// * If `obj` is a bound symbol:
+    ///   -> `obj["prop"] = Math.pow(obj["prop"], right)`
+    /// * If `obj` is unbound:
+    ///   -> `var _obj; _obj = obj, _obj["prop"] = Math.pow(_obj["prop"], right)`
+    ///
+    /// `obj[name] **= right`
+    /// * If `obj` is a bound symbol:
+    ///   -> `var _name; _name = name, obj[_name] = Math.pow(obj[_name], 2)`
+    /// * If `obj` is unbound:
+    ///   -> `var _obj, _name; _obj = obj, _name = name, _obj[_name] = Math.pow(_obj[_name], 2)`
+    ///
+    /// Temporary variables are to avoid side-effects of getting `obj` or `name` being run twice.
+    ///
+    /// TODO(improve-on-babel):
+    /// 1. If `name` is bound, it doesn't need a temp variable `_name`.
+    /// 2. `obj.prop` does not need to be transformed to `obj["prop"]`.
+    /// We currently aim to produce output that exactly matches Babel, but we can improve this in future
+    /// when we no longer need to match exactly.
+    fn convert_assignment_to_member_expression(
         &mut self,
         expr: &mut Expression<'a>,
         ctx: &mut TraverseCtx<'a>,
@@ -108,18 +188,86 @@ impl<'a, 'ctx> ExponentiationOperator<'a, 'ctx> {
         let Expression::AssignmentExpression(assign_expr) = expr else { unreachable!() };
 
         let mut nodes = ctx.ast.vec();
-        let Some(Exploded { reference, uid }) =
+        let Exploded { reference, uid } =
-            self.explode(&mut assign_expr.left, &mut nodes, ctx)
+            self.explode_member_expression(&mut assign_expr.left, &mut nodes, ctx);
-        else {
+
-            return;
+        *expr = Self::create_replacement(assign_expr, reference, uid, nodes, ctx);
-        };
+    }
+
+    fn create_replacement(
+        assign_expr: &mut AssignmentExpression<'a>,
+        reference: AssignmentTarget<'a>,
+        uid: Expression<'a>,
+        mut nodes: Vec<'a, Expression<'a>>,
+        ctx: &mut TraverseCtx<'a>,
+    ) -> Expression<'a> {
         let right = ctx.ast.move_expression(&mut assign_expr.right);
         let right = Self::math_pow(uid, right, ctx);
         let assign_expr =
             ctx.ast.expression_assignment(SPAN, AssignmentOperator::Assign, reference, right);
         nodes.push(assign_expr);
 
-        *expr = ctx.ast.expression_sequence(SPAN, nodes);
+        ctx.ast.expression_sequence(SPAN, nodes)
+    }
+
+    fn explode_member_expression(
+        &mut self,
+        node: &mut AssignmentTarget<'a>,
+        nodes: &mut Vec<'a, Expression<'a>>,
+        ctx: &mut TraverseCtx<'a>,
+    ) -> Exploded<'a> {
+        let member_expr = node.to_member_expression_mut();
+
+        // Make sure side-effects of evaluating `obj` of `obj.ref` and `obj[ref]` only happen once
+        let obj = match member_expr {
+            MemberExpression::ComputedMemberExpression(e) => &mut e.object,
+            MemberExpression::StaticMemberExpression(e) => &mut e.object,
+            // This possibility is ruled out in `enter_expression`
+            MemberExpression::PrivateFieldExpression(_) => unreachable!(),
+        };
+        let mut obj = ctx.ast.move_expression(obj);
+        // If the object reference that we need to save is locally declared, evaluating it multiple times
+        // will not trigger getters or setters. `super` cannot be directly assigned, so use it directly too.
+        let needs_temp_var = match &obj {
+            Expression::Super(_) => false,
+            Expression::Identifier(ident) => {
+                !ctx.symbols().has_binding(ident.reference_id().unwrap())
+            }
+            _ => true,
+        };
+        if needs_temp_var {
+            obj = self.add_new_reference(obj, nodes, ctx);
+        }
+
+        let computed = member_expr.is_computed();
+        let prop = self.get_prop_ref(member_expr, nodes, ctx);
+        let optional = false;
+        let obj_clone = Self::clone_expression(&obj, ctx);
+        let (reference, uid) = match &prop {
+            Expression::Identifier(ident) if !computed => {
+                let ident = IdentifierName::new(SPAN, ident.name.clone());
+                (
+                    // TODO:
+                    // Both of these are the same, but it's in order to avoid after cloning without reference_id.
+                    // Related: https://github.com/oxc-project/oxc/issues/4804
+                    ctx.ast.member_expression_static(SPAN, obj_clone, ident.clone(), optional),
+                    ctx.ast.member_expression_static(SPAN, obj, ident, optional),
+                )
+            }
+            _ => {
+                let prop_clone = Self::clone_expression(&prop, ctx);
+                (
+                    ctx.ast.member_expression_computed(SPAN, obj_clone, prop_clone, optional),
+                    ctx.ast.member_expression_computed(SPAN, obj, prop, optional),
+                )
+            }
+        };
+        Exploded {
+            reference: AssignmentTarget::from(
+                ctx.ast.simple_assignment_target_member_expression(reference),
+            ),
+            uid: Expression::from(uid),
+        }
     }
 
     fn clone_expression(expr: &Expression<'a>, ctx: &mut TraverseCtx<'a>) -> Expression<'a> {
@@ -150,127 +298,6 @@ impl<'a, 'ctx> ExponentiationOperator<'a, 'ctx> {
         ctx.ast.expression_call(SPAN, callee, NONE, arguments, false)
     }
 
-    /// Change `lhs **= 2` to `var temp; temp = lhs, lhs = Math.pow(temp, 2);`.
-    /// If the lhs is a member expression `obj.ref` or `obj[ref]`, assign them to a temporary variable so side-effects are not computed twice.
-    /// For `obj.ref`, change it to `var _obj; _obj = obj, _obj["ref"] = Math.pow(_obj["ref"], 2)`.
-    /// For `obj[ref]`, change it to `var _obj, _ref; _obj = obj, _ref = ref, _obj[_ref] = Math.pow(_obj[_ref], 2);`.
-    fn explode(
-        &mut self,
-        node: &mut AssignmentTarget<'a>,
-        nodes: &mut Vec<'a, Expression<'a>>,
-        ctx: &mut TraverseCtx<'a>,
-    ) -> Option<Exploded<'a>> {
-        let (reference, uid) = match node {
-            AssignmentTarget::AssignmentTargetIdentifier(_) => {
-                let obj = self.get_obj_ref(node, nodes, ctx).unwrap();
-                let ident = ctx.ast.move_assignment_target(node);
-                (ident, obj)
-            }
-            match_member_expression!(AssignmentTarget) => {
-                let obj = self.get_obj_ref(node, nodes, ctx)?;
-                let member_expr = node.to_member_expression_mut();
-                let computed = member_expr.is_computed();
-                let prop = self.get_prop_ref(member_expr, nodes, ctx);
-                let optional = false;
-                let obj_clone = Self::clone_expression(&obj, ctx);
-                let (reference, uid) = match &prop {
-                    Expression::Identifier(ident) if !computed => {
-                        let ident = IdentifierName::new(SPAN, ident.name.clone());
-                        (
-                            // TODO:
-                            // Both of these are the same, but it's in order to avoid after cloning without reference_id.
-                            // Related: https://github.com/oxc-project/oxc/issues/4804
-                            ctx.ast.member_expression_static(
-                                SPAN,
-                                obj_clone,
-                                ident.clone(),
-                                optional,
-                            ),
-                            ctx.ast.member_expression_static(SPAN, obj, ident, optional),
-                        )
-                    }
-                    _ => {
-                        let prop_clone = Self::clone_expression(&prop, ctx);
-                        (
-                            ctx.ast
-                                .member_expression_computed(SPAN, obj_clone, prop_clone, optional),
-                            ctx.ast.member_expression_computed(SPAN, obj, prop, optional),
-                        )
-                    }
-                };
-                (
-                    AssignmentTarget::from(
-                        ctx.ast.simple_assignment_target_member_expression(reference),
-                    ),
-                    Expression::from(uid),
-                )
-            }
-            _ => return None,
-        };
-        Some(Exploded { reference, uid })
-    }
-
-    /// Make sure side-effects of evaluating `obj` of `obj.ref` and `obj[ref]` only happen once.
-    fn get_obj_ref(
-        &mut self,
-        node: &mut AssignmentTarget<'a>,
-        nodes: &mut Vec<'a, Expression<'a>>,
-        ctx: &mut TraverseCtx<'a>,
-    ) -> Option<Expression<'a>> {
-        let reference = match node {
-            AssignmentTarget::AssignmentTargetIdentifier(ident) => {
-                let reference = ctx.symbols().get_reference(ident.reference_id().unwrap());
-                if let Some(symbol_id) = reference.symbol_id() {
-                    // this variable is declared in scope so we can be 100% sure
-                    // that evaluating it multiple times won't trigger a getter
-                    // or something else
-                    return Some(ctx.ast.expression_from_identifier_reference(
-                        ctx.create_bound_reference_id(
-                            SPAN,
-                            ident.name.clone(),
-                            symbol_id,
-                            ReferenceFlags::Write,
-                        ),
-                    ));
-                }
-                // could possibly trigger a getter so we need to only evaluate it once
-                ctx.ast.expression_from_identifier_reference(ctx.create_unbound_reference_id(
-                    SPAN,
-                    ident.name.clone(),
-                    ReferenceFlags::Read,
-                ))
-            }
-            match_member_expression!(AssignmentTarget) => {
-                let expr = match node.to_member_expression_mut() {
-                    MemberExpression::ComputedMemberExpression(e) => &mut e.object,
-                    MemberExpression::StaticMemberExpression(e) => &mut e.object,
-                    // From Babel: "We can't generate property ref for private name, please install
-                    // `@babel/plugin-transform-class-properties`".
-                    // TODO: Ensure this plugin interacts correctly with class private properties
-                    // transform, so the property is transformed before this transform.
-                    MemberExpression::PrivateFieldExpression(_) => return None,
-                };
-                let expr = ctx.ast.move_expression(expr);
-                // the object reference that we need to save is locally declared
-                // so as per the previous comment we can be 100% sure evaluating
-                // it multiple times will be safe
-                // Super cannot be directly assigned so lets return it also
-                if matches!(expr, Expression::Super(_))
-                    || matches!(&expr, Expression::Identifier(ident) if ident
-                        .reference_id
-                        .get()
-                        .is_some_and(|reference_id| ctx.symbols().has_binding(reference_id)))
-                {
-                    return Some(expr);
-                }
-
-                expr
-            }
-            _ => return None,
-        };
-        Some(self.add_new_reference(reference, nodes, ctx))
-    }
-
     /// Make sure side-effects of evaluating `ref` of `obj.ref` and `obj[ref]` only happen once.
     fn get_prop_ref(
         &mut self,
@@ -289,7 +316,7 @@ impl<'a, 'ctx> ExponentiationOperator<'a, 'ctx> {
             MemberExpression::StaticMemberExpression(expr) => {
                 ctx.ast.expression_string_literal(SPAN, expr.property.name.clone())
             }
-            // This possibility is ruled out in earlier call to `get_obj_ref`
+            // This possibility is ruled out in `enter_expression`
             MemberExpression::PrivateFieldExpression(_) => unreachable!(),
         }
     }