feat(react-compiler): add reactiveConditionHelper

GrinZero · Jul 1, 2024 · 3b7af2e · 3b7af2e
1 parent 1e0d652
commit 3b7af2e
Show file tree

Hide file tree

Showing 6 changed files with 152 additions and 54 deletions.
diff --git a/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Imports.ts b/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Imports.ts
@@ -100,7 +100,8 @@ function hasExistingNonNamespacedImportOfModule(
 function addMemoCacheFunctionSpecifierToExistingImport(
   program: NodePath<t.Program>,
   moduleName: string,
-  identifierName: string
+  identifierName: string,
+  reactiveFnHelperIdentifierName?: string
 ): boolean {
   let didInsertUseMemoCache = false;
   program.traverse({
@@ -109,10 +110,17 @@ function addMemoCacheFunctionSpecifierToExistingImport(
         !didInsertUseMemoCache &&
         isNonNamespacedImport(importDeclPath, moduleName)
       ) {
-        importDeclPath.pushContainer(
-          "specifiers",
-          t.importSpecifier(t.identifier(identifierName), t.identifier("c"))
-        );
+        importDeclPath.pushContainer("specifiers", [
+          t.importSpecifier(t.identifier(identifierName), t.identifier("c")),
+          ...(reactiveFnHelperIdentifierName
+            ? [
+                t.importSpecifier(
+                  t.identifier(reactiveFnHelperIdentifierName),
+                  t.identifier("u")
+                ),
+              ]
+            : []),
+        ]);
         didInsertUseMemoCache = true;
       }
     },
@@ -123,7 +131,8 @@ function addMemoCacheFunctionSpecifierToExistingImport(
 export function updateMemoCacheFunctionImport(
   program: NodePath<t.Program>,
   moduleName: string,
-  useMemoCacheIdentifier: string
+  useMemoCacheIdentifier: string,
+  reactiveFnHelperIdentifier?: string
 ): void {
   /*
    * If there isn't already an import of * as React, insert it so useMemoCache doesn't
@@ -138,7 +147,8 @@ export function updateMemoCacheFunctionImport(
     const didUpdateImport = addMemoCacheFunctionSpecifierToExistingImport(
       program,
       moduleName,
-      useMemoCacheIdentifier
+      useMemoCacheIdentifier,
+      reactiveFnHelperIdentifier
     );
     if (!didUpdateImport) {
       throw new Error(
@@ -149,20 +159,32 @@ export function updateMemoCacheFunctionImport(
     addMemoCacheFunctionImportDeclaration(
       program,
       moduleName,
-      useMemoCacheIdentifier
+      useMemoCacheIdentifier,
+      reactiveFnHelperIdentifier
     );
   }
 }
 
 function addMemoCacheFunctionImportDeclaration(
   program: NodePath<t.Program>,
   moduleName: string,
-  localName: string
+  localName: string,
+  reactiveFnHelperIdentifierName?: string
 ): void {
   program.unshiftContainer(
     "body",
     t.importDeclaration(
-      [t.importSpecifier(t.identifier(localName), t.identifier("c"))],
+      [
+        t.importSpecifier(t.identifier(localName), t.identifier("c")),
+        ...(reactiveFnHelperIdentifierName
+          ? [
+              t.importSpecifier(
+                t.identifier(reactiveFnHelperIdentifierName),
+                t.identifier("u")
+              ),
+            ]
+          : []),
+      ],
       t.stringLiteral(moduleName)
     )
   );

diff --git a/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts b/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts
@@ -111,6 +111,7 @@ export function* run(
   config: EnvironmentConfig,
   fnType: ReactFunctionType,
   useMemoCacheIdentifier: string,
+  reactiveFnHelperIdentifier: string,
   logger: Logger | null,
   filename: string | null,
   code: string | null
@@ -123,7 +124,8 @@ export function* run(
     logger,
     filename,
     code,
-    useMemoCacheIdentifier
+    useMemoCacheIdentifier,
+    reactiveFnHelperIdentifier
   );
   yield {
     kind: "debug",
@@ -499,6 +501,7 @@ export function compileFn(
   config: EnvironmentConfig,
   fnType: ReactFunctionType,
   useMemoCacheIdentifier: string,
+  reactiveFnHelperIdentifier: string,
   logger: Logger | null,
   filename: string | null,
   code: string | null
@@ -508,6 +511,7 @@ export function compileFn(
     config,
     fnType,
     useMemoCacheIdentifier,
+    reactiveFnHelperIdentifier,
     logger,
     filename,
     code

diff --git a/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Program.ts b/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Program.ts
@@ -248,6 +248,7 @@ export function compileProgram(
 
   const environment = parseEnvironmentConfig(pass.opts.environment ?? {});
   const useMemoCacheIdentifier = program.scope.generateUidIdentifier("c");
+  const reactiveFnHelperIdentifier = program.scope.generateUidIdentifier("u");
   const moduleName = pass.opts.runtimeModule ?? "react/compiler-runtime";
   if (hasMemoCacheFunctionImport(program, moduleName)) {
     return;
@@ -318,6 +319,7 @@ export function compileProgram(
         config,
         fnType,
         useMemoCacheIdentifier.name,
+        reactiveFnHelperIdentifier.name,
         pass.opts.logger,
         pass.filename,
         pass.code
@@ -463,7 +465,10 @@ export function compileProgram(
       updateMemoCacheFunctionImport(
         program,
         moduleName,
-        useMemoCacheIdentifier.name
+        useMemoCacheIdentifier.name,
+        environment.unwrap().enableReactiveScopeConditionHelper
+          ? reactiveFnHelperIdentifier.name
+          : void 0
       );
     }
     addImportsToProgram(program, externalFunctions);

diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts
@@ -197,6 +197,8 @@ const EnvironmentConfigSchema = z.object({
 
   enableReactiveScopesInHIR: z.boolean().default(true),
 
+  enableReactiveScopeConditionHelper: z.boolean().default(false),
+
   /*
    * Enable validation of hooks to partially check that the component honors the rules of hooks.
    * When disabled, the component is assumed to follow the rules (though the Babel plugin looks
@@ -510,6 +512,7 @@ export class Environment {
   config: EnvironmentConfig;
   fnType: ReactFunctionType;
   useMemoCacheIdentifier: string;
+  reactiveFnHelperIdentifier: string;
 
   #contextIdentifiers: Set<t.Identifier>;
   #hoistedIdentifiers: Set<t.Identifier>;
@@ -521,14 +524,16 @@ export class Environment {
     logger: Logger | null,
     filename: string | null,
     code: string | null,
-    useMemoCacheIdentifier: string
+    useMemoCacheIdentifier: string,
+    reactiveFnHelperIdentifier: string
   ) {
     this.fnType = fnType;
     this.config = config;
     this.filename = filename;
     this.code = code;
     this.logger = logger;
     this.useMemoCacheIdentifier = useMemoCacheIdentifier;
+    this.reactiveFnHelperIdentifier = reactiveFnHelperIdentifier;
     this.#shapes = new Map(DEFAULT_SHAPES);
     this.#globals = new Map(DEFAULT_GLOBALS);
 

diff --git a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/CodegenReactiveFunction.ts b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/CodegenReactiveFunction.ts
@@ -156,28 +156,66 @@ export function codegenFunction(
       // HMR detection is enabled, emit code to reset the memo cache on source changes
       const index = cx.synthesizeName("$i");
       preface.push(
-        t.variableDeclaration("const", [
-          t.variableDeclarator(
-            t.identifier(index),
-            t.logicalExpression(
-              "||",
-              t.callExpression(t.identifier("u"), [
-                t.identifier(cx.synthesizeName("$")),
-                t.arrayExpression([]),
-                t.arrayExpression([]),
-                t.arrayExpression([]),
+        t.ifStatement(
+          t.binaryExpression(
+            "!==",
+            t.memberExpression(
+              t.identifier(cx.synthesizeName("$")),
+              t.numericLiteral(fastRefreshState.cacheIndex),
+              true
+            ),
+            t.stringLiteral(fastRefreshState.hash)
+          ),
+          t.blockStatement([
+            t.forStatement(
+              t.variableDeclaration("let", [
+                t.variableDeclarator(t.identifier(index), t.numericLiteral(0)),
               ]),
-              t.memberExpression(
-                t.identifier(cx.synthesizeName("$")),
-                t.numericLiteral(fastRefreshState.cacheIndex),
-                true
+              t.binaryExpression(
+                "<",
+                t.identifier(index),
+                t.numericLiteral(cacheCount)
+              ),
+              t.assignmentExpression(
+                "+=",
+                t.identifier(index),
+                t.numericLiteral(1)
+              ),
+              t.blockStatement([
+                t.expressionStatement(
+                  t.assignmentExpression(
+                    "=",
+                    t.memberExpression(
+                      t.identifier(cx.synthesizeName("$")),
+                      t.identifier(index),
+                      true
+                    ),
+                    t.callExpression(
+                      t.memberExpression(
+                        t.identifier("Symbol"),
+                        t.identifier("for")
+                      ),
+                      [t.stringLiteral(MEMO_CACHE_SENTINEL)]
+                    )
+                  )
+                ),
+              ])
+            ),
+            t.expressionStatement(
+              t.assignmentExpression(
+                "=",
+                t.memberExpression(
+                  t.identifier(cx.synthesizeName("$")),
+                  t.numericLiteral(fastRefreshState.cacheIndex),
+                  true
+                ),
+                t.stringLiteral(fastRefreshState.hash)
               )
-            )
-          ),
-        ])
+            ),
+          ])
+        )
       );
     }
-    compiled.body.body.unshift(...preface);
   }
 
   const emitInstrumentForget = fn.env.config.enableEmitInstrumentForget;
@@ -394,11 +432,9 @@ function codegenBlockNoReset(
   block: ReactiveBlock
 ): t.BlockStatement {
   const statements: Array<t.Statement> = [];
-  //TODO: hack to generate
   for (const item of block) {
     switch (item.kind) {
       case "instruction": {
-        // like a = b + c
         const statement = codegenInstructionNullable(cx, item.instruction);
         if (statement !== null) {
           statements.push(statement);
@@ -411,14 +447,12 @@ function codegenBlockNoReset(
         break;
       }
       case "scope": {
-        // like if, for, while
         const temp = new Map(cx.temp);
         codegenReactiveScope(cx, statements, item.scope, item.instructions);
         cx.temp = temp;
         break;
       }
       case "terminal": {
-        // like break, continue, return
         const statement = codegenTerminal(cx, item.terminal);
         if (statement === null) {
           break;
@@ -465,7 +499,6 @@ function wrapCacheDep(cx: Context, value: t.Expression): t.Expression {
   }
 }
 
-const count: 0 | 1 = 1;
 function codegenReactiveScope(
   cx: Context,
   statements: Array<t.Statement>,
@@ -483,18 +516,17 @@ function codegenReactiveScope(
   const changeExpressionComments: Array<string> = [];
   const outputComments: Array<string> = [];
 
-  // #region taking
-
-  if (count === 1) {
-    const indices: t.NumericLiteral[] = [];
-    const newValues: t.Expression[] = [];
+  // #region enableReactiveScopeConditionHelper
+  if (cx.env.config.enableReactiveScopeConditionHelper) {
+    const conditionIndexs: t.NumericLiteral[] = [];
+    const conditionTargets: t.Expression[] = [];
     const updateIndexs: t.NumericLiteral[] = [];
     const source: (t.Expression | null | undefined)[] = [];
 
     for (const dep of scope.dependencies) {
       const index = cx.nextCacheIndex;
-      indices.push(t.numericLiteral(index));
-      newValues.push(codegenDependency(cx, dep));
+      conditionIndexs.push(t.numericLiteral(index));
+      conditionTargets.push(codegenDependency(cx, dep));
       const comparison = t.binaryExpression(
         "!==",
         t.memberExpression(
@@ -510,8 +542,8 @@ function codegenReactiveScope(
     }
 
     const arr: t.Identifier[] = [];
-    const index = cx.nextCacheIndex;
     for (const [, { identifier }] of scope.declarations) {
+      const index = cx.nextCacheIndex;
       const name = convertIdentifier(identifier);
       arr.push(name);
       cacheLoads.push({ name, index, value: wrapCacheDep(cx, name) });
@@ -596,14 +628,16 @@ function codegenReactiveScope(
         t.arrayPattern(arr),
         t.logicalExpression(
           "||",
-          t.callExpression(t.identifier("u"), [
+          t.callExpression(t.identifier(cx.env.reactiveFnHelperIdentifier), [
             t.identifier(cx.synthesizeName("$")),
             t.arrowFunctionExpression(
               [],
               t.arrayExpression(bd as t.Expression[])
             ),
-            t.arrayExpression(indices),
-            testCondition ? t.arrayExpression(newValues) : t.nullLiteral(),
+            t.arrayExpression(conditionIndexs),
+            testCondition
+              ? t.arrayExpression(conditionTargets)
+              : t.nullLiteral(),
             t.arrayExpression(updateIndexs),
             t.arrowFunctionExpression(
               [],
@@ -688,7 +722,6 @@ function codegenReactiveScope(
 
     const name = convertIdentifier(identifier);
     outputComments.push(name.name);
-    // !!! const t4
     if (!cx.hasDeclared(identifier)) {
       statements.push(
         t.variableDeclaration("let", [t.variableDeclarator(name)])
@@ -936,11 +969,6 @@ function codegenReactiveScope(
     );
     const name: ValidIdentifierName = earlyReturnValue.value.name.value;
     statements.push(
-      /**
-       * if (name !== Symbol.for("EARLY_RETURN_SENTINEL")) {
-       *  return name;
-       * }
-       */
       t.ifStatement(
         t.binaryExpression(
           "!==",
@@ -1240,7 +1268,6 @@ function codegenInstructionNullable(
   cx: Context,
   instr: ReactiveInstruction
 ): t.Statement | null {
-  // TODO: know
   if (
     instr.value.kind === "StoreLocal" ||
     instr.value.kind === "StoreContext" ||