Support logical assignment expressions in compiler HIR

compiler/packages/babel-plugin-react-compiler/src/HIR/BuildHIR.ts

@@ -2108,6 +2108,18 @@ function lowerExpression(
         }
       }
 
+      const logicalOperators: {
+        [key: string]: t.LogicalExpression['operator'];
+      } = {
+        '&&=': '&&',
+        '||=': '||',
+        '??=': '??',
+      };
+      const logicalOperator = logicalOperators[operator];
+      if (logicalOperator !== undefined) {
+        return lowerLogicalAssignment(builder, expr, logicalOperator);
+      }
+
       const operators: {
         [key: string]: Exclude<t.BinaryExpression['operator'], '|>'>;
       } = {
@@ -3323,6 +3335,163 @@ function lowerArguments(
   return args;
 }
 
+type LogicalAssignmentTarget =
+  | {kind: 'Identifier'; path: NodePath<t.Identifier>}
+  | {
+      kind: 'MemberExpression';
+      object: Place;
+      property: Place | string | number;
+      loc: SourceLocation;
+    };
+
+function lowerLogicalAssignment(
+  builder: HIRBuilder,
+  expr: NodePath<t.AssignmentExpression>,
+  operator: t.LogicalExpression['operator'],
+): InstructionValue {
+  const left = expr.get('left');
+  const exprLoc = expr.node.loc ?? GeneratedSource;
+  const leftLoc = left.node.loc ?? GeneratedSource;
+
+  if (!left.isIdentifier() && !left.isMemberExpression()) {
+    builder.recordError(
+      new CompilerErrorDetail({
+        reason: `(BuildHIR::lowerExpression) Expected Identifier or MemberExpression, got ${left.type} lval in AssignmentExpression`,
+        category: ErrorCategory.Todo,
+        loc: expr.node.loc ?? null,
+        suggestions: null,
+      }),
+    );
+    return {kind: 'UnsupportedNode', node: expr.node, loc: exprLoc};
+  }
+
+  const continuationBlock = builder.reserve(builder.currentBlockKind());
+  const testBlock = builder.reserve('value');
+  const consequentBlock = builder.reserve('value');
+  const alternateBlock = builder.reserve('value');
+  const place = buildTemporaryPlace(builder, exprLoc);
+  const previousValuePlace = buildTemporaryPlace(builder, leftLoc);
+  let assignmentTarget: LogicalAssignmentTarget;
+  let previousValue: Place;
+  if (left.isIdentifier()) {
+    previousValue = lowerExpressionToTemporary(builder, left);
+    assignmentTarget = {kind: 'Identifier', path: left};
+  } else {
+    const {object, property, value} = lowerMemberExpression(builder, left);
+    previousValue = lowerValueToTemporary(builder, value);
+    assignmentTarget = {
+      kind: 'MemberExpression',
+      object,
+      property,
+      loc: leftLoc,
+    };
+  }
+
+  builder.enterReserved(testBlock, () => {
+    builder.push({
+      id: makeInstructionId(0),
+      lvalue: {...previousValuePlace},
+      value: {
+        kind: 'LoadLocal',
+        place: previousValue,
+        loc: leftLoc,
+      },
+      effects: null,
+      loc: leftLoc,
+    });
+    return {
+      kind: 'branch',
+      test: {...previousValuePlace},
+      consequent: consequentBlock.id,
+      alternate: alternateBlock.id,
+      fallthrough: continuationBlock.id,
+      id: makeInstructionId(0),
+      loc: exprLoc,
+    };
+  });
+
+  builder.enterReserved(consequentBlock, () => {
+    lowerValueToTemporary(builder, {
+      kind: 'StoreLocal',
+      lvalue: {kind: InstructionKind.Const, place: {...place}},
+      value: {...previousValuePlace},
+      type: null,
+      loc: leftLoc,
+    });
+    return {
+      kind: 'goto',
+      block: continuationBlock.id,
+      variant: GotoVariant.Break,
+      id: makeInstructionId(0),
+      loc: leftLoc,
+    };
+  });
+
+  builder.enterReserved(alternateBlock, () => {
+    const right = lowerExpressionToTemporary(builder, expr.get('right'));
+    let assignedValue: Place;
+    if (assignmentTarget.kind === 'Identifier') {
+      assignedValue = lowerValueToTemporary(
+        builder,
+        lowerAssignment(
+          builder,
+          leftLoc,
+          InstructionKind.Reassign,
+          assignmentTarget.path,
+          right,
+          'Assignment',
+        ),
+      );
+    } else {
+      const storeValue: InstructionValue =
+        typeof assignmentTarget.property === 'string' ||
+        typeof assignmentTarget.property === 'number'
+          ? {
+              kind: 'PropertyStore',
+              object: {...assignmentTarget.object},
+              property: makePropertyLiteral(assignmentTarget.property),
+              value: {...right},
+              loc: assignmentTarget.loc,
+            }
+          : {
+              kind: 'ComputedStore',
+              object: {...assignmentTarget.object},
+              property: {...assignmentTarget.property},
+              value: {...right},
+              loc: assignmentTarget.loc,
+            };
+      assignedValue = lowerValueToTemporary(builder, storeValue);
+    }
+    lowerValueToTemporary(builder, {
+      kind: 'StoreLocal',
+      lvalue: {kind: InstructionKind.Const, place: {...place}},
+      value: {...assignedValue},
+      type: null,
+      loc: exprLoc,
+    });
+    return {
+      kind: 'goto',
+      block: continuationBlock.id,
+      variant: GotoVariant.Break,
+      id: makeInstructionId(0),
+      loc: exprLoc,
+    };
+  });
+
+  builder.terminateWithContinuation(
+    {
+      kind: 'logical',
+      fallthrough: continuationBlock.id,
+      id: makeInstructionId(0),
+      test: testBlock.id,
+      operator,
+      loc: exprLoc,
+    },
+    continuationBlock,
+  );
+  return {kind: 'LoadLocal', place, loc: place.loc};
+}
+
 type LoweredMemberExpression = {
   object: Place;
   property: Place | string | number;

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/nullish-assignment-expression.expect.md

@@ -0,0 +1,81 @@
+
+## Input
+
+```javascript
+function Component(props) {
+  const value = {count: null, other: 2, enabled: true, fallback: 0};
+  let writes = 0;
+
+  value[props.key] ??= (writes += 1, props.count);
+  value.other ??= (writes += 1, props.count);
+  value.enabled &&= (writes += 1, props.enabled);
+  value.fallback ||= (writes += 1, props.count);
+
+  return [value.count, value.other, value.enabled, value.fallback, writes];
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{key: 'count', count: 1, enabled: false}],
+  isComponent: false,
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime";
+function Component(props) {
+  const $ = _c(11);
+  let value;
+  let writes;
+  if ($[0] !== props.count || $[1] !== props.enabled || $[2] !== props.key) {
+    value = { count: null, other: 2, enabled: true, fallback: 0 };
+    writes = 0;
+
+    const t0 = props.key;
+    value[t0] ?? (value[t0] = ((writes = 1), props.count));
+    value.other ?? (value.other = ((writes = writes + 1), props.count));
+    value.enabled && (value.enabled = ((writes = writes + 1), props.enabled));
+    value.fallback || (value.fallback = ((writes = writes + 1), props.count));
+    $[0] = props.count;
+    $[1] = props.enabled;
+    $[2] = props.key;
+    $[3] = value;
+    $[4] = writes;
+  } else {
+    value = $[3];
+    writes = $[4];
+  }
+  let t0;
+  if (
+    $[5] !== value.count ||
+    $[6] !== value.enabled ||
+    $[7] !== value.fallback ||
+    $[8] !== value.other ||
+    $[9] !== writes
+  ) {
+    t0 = [value.count, value.other, value.enabled, value.fallback, writes];
+    $[5] = value.count;
+    $[6] = value.enabled;
+    $[7] = value.fallback;
+    $[8] = value.other;
+    $[9] = writes;
+    $[10] = t0;
+  } else {
+    t0 = $[10];
+  }
+  return t0;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ key: "count", count: 1, enabled: false }],
+  isComponent: false,
+};
+
+```
+
+### Eval output
+(kind: ok) [1,2,false,1,3]

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/nullish-assignment-expression.js

@@ -0,0 +1,17 @@
+function Component(props) {
+  const value = {count: null, other: 2, enabled: true, fallback: 0};
+  let writes = 0;
+
+  value[props.key] ??= (writes += 1, props.count);
+  value.other ??= (writes += 1, props.count);
+  value.enabled &&= (writes += 1, props.enabled);
+  value.fallback ||= (writes += 1, props.count);
+
+  return [value.count, value.other, value.enabled, value.fallback, writes];
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{key: 'count', count: 1, enabled: false}],
+  isComponent: false,
+};