Implement hoist constant JSX

packages/forgetti/package.json

@@ -16,7 +16,7 @@
   "devDependencies": {
     "@babel/core": "^7.21.3",
     "@types/babel__core": "^7.20.0",
-    "@types/babel__traverse": "^7.18.3",
+    "@types/babel__traverse": "^7.20.1",
     "@types/node": "^18.15.3",
     "eslint": "^8.38.0",
     "eslint-config-lxsmnsyc": "^0.6.3",

packages/forgetti/src/core/checks.ts

@@ -56,6 +56,20 @@ export function shouldSkipJSX(node: t.Node): boolean {
   return false;
 }
 
+const FORGETTI_JSX_HOISTED = /^\s*@forgetti hoisted_jsx\s*$/;
+
+export function shouldSkipJSXExtraction(node: t.Node): boolean {
+  // Node without leading comments shouldn't be skipped
+  if (node.leadingComments) {
+    for (let i = 0, len = node.leadingComments.length; i < len; i++) {
+      if (FORGETTI_JSX_HOISTED.test(node.leadingComments[i].value)) {
+        return true;
+      }
+    }
+  }
+  return false;
+}
+
 export function isComponentValid(
   ctx: StateContext,
   node: ComponentNode,

packages/forgetti/src/core/hoist-jsx.ts

@@ -0,0 +1,207 @@
+import * as t from '@babel/types';
+import type { Scope, Visitor } from '@babel/traverse';
+import type { StateContext } from './types';
+
+interface VisitorState {
+  isImmutable: boolean;
+  jsxScope: Scope;
+  targetScope: Scope;
+}
+
+function declares(node: t.Identifier | t.JSXIdentifier, scope: Scope): boolean {
+  if (
+    t.isJSXIdentifier(node, { name: 'this' })
+    || t.isJSXIdentifier(node, { name: 'arguments' })
+    || t.isJSXIdentifier(node, { name: 'super' })
+    || t.isJSXIdentifier(node, { name: 'new' })
+  ) {
+    const { path } = scope;
+    return path.isFunctionParent() && !path.isArrowFunctionExpression();
+  }
+
+  return scope.hasOwnBinding(node.name);
+}
+
+function isHoistingScope({ path }: Scope): boolean {
+  return path.isFunctionParent() || path.isLoop() || path.isProgram();
+}
+
+function getHoistingScope(scope: Scope): Scope {
+  while (!isHoistingScope(scope)) scope = scope.parent;
+  return scope;
+}
+
+const hoistingVisitor = {
+  enter(path: babel.NodePath, state: VisitorState): void {
+    const stop = (): void => {
+      state.isImmutable = false;
+      path.stop();
+    };
+
+    const skip = (): void => {
+      path.skip();
+    };
+
+    if (path.isJSXClosingElement()) {
+      skip();
+      return;
+    }
+
+    // Elements with refs are not safe to hoist.
+    if (
+      path.isJSXIdentifier({ name: 'ref' })
+      && path.parentPath.isJSXAttribute({ name: path.node })
+    ) {
+      stop();
+      return;
+    }
+
+    // Ignore JSX expressions and immutable values.
+    if (
+      path.isJSXIdentifier()
+      || path.isJSXMemberExpression()
+      || path.isJSXNamespacedName()
+      || path.isImmutable()
+    ) {
+      return;
+    }
+
+    // Ignore constant bindings.
+    if (path.isIdentifier()) {
+      const binding = path.scope.getBinding(path.node.name);
+      if (binding && binding.constant) return;
+    }
+
+    if (!path.isPure()) {
+      stop();
+      return;
+    }
+
+    // If it's not immutable, it may still be a pure expression, such as string concatenation.
+    // It is still safe to hoist that, so long as its result is immutable.
+    // If not, it is not safe to replace as mutable values (e.g. objects), as it could be
+    // mutated after render.
+    // https://github.com/facebook/react/issues/3226
+    const expressionResult = path.evaluate();
+    if (expressionResult.confident) {
+      // We know the result; check its mutability.
+      if (
+        expressionResult.value === null
+        || (typeof expressionResult.value !== 'object' && typeof expressionResult.value !== 'function')
+      ) {
+        // It evaluated to an immutable value, so we can hoist it.
+        skip();
+        return;
+      }
+    } else if (expressionResult.deopt?.isIdentifier()) {
+      // It's safe to hoist here if the deopt reason is an identifier (e.g. func param).
+      // The hoister will take care of how high up it can be hoisted.
+      return;
+    }
+
+    stop();
+  },
+  ReferencedIdentifier(path: babel.NodePath<t.Identifier>, state: VisitorState): void {
+    const { node } = path;
+    let { scope } = path;
+
+    while (scope !== state.jsxScope) {
+      // If a binding is declared in an inner function, it doesn't affect hoisting.
+      if (declares(node, scope)) return;
+
+      scope = scope.parent;
+    }
+
+    // We are recursing up the scope chain, the parent may not be existed.
+    // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition
+    while (scope) {
+      // We cannot hoist outside of the previous hoisting target
+      // scope, so we return early and we don't update it.
+      if (scope === state.targetScope) return;
+
+      // If the scope declares this identifier (or we're at the function
+      // providing the lexical env binding), we can't hoist the var any
+      // higher.
+      if (declares(node, scope)) break;
+
+      scope = scope.parent;
+    }
+
+    state.targetScope = getHoistingScope(scope);
+  },
+};
+
+export default function hoistConstantJSX(
+  ctx: StateContext,
+  path: babel.NodePath<t.JSXElement>,
+): void {
+  if (ctx.hoist.jsxScopeMap.has(path.node)) return;
+
+  const { name } = path.node.openingElement;
+
+  // In order to avoid hoisting unnecessarily, we need to know which is
+  // the scope containing the current JSX element. If a parent of the
+  // current element has already been hoisted, we can consider its target
+  // scope as the base scope for the current element.
+  let jsxScope;
+  let current: babel.NodePath<t.JSX> = path;
+  while (!jsxScope && current.parentPath.isJSX()) {
+    current = current.parentPath;
+    jsxScope = ctx.hoist.jsxScopeMap.get(current.node);
+  }
+  jsxScope ??= path.scope;
+  // The initial HOISTED is set to jsxScope, s.t.
+  // if the element's JSX ancestor has been hoisted, it will be skipped
+  ctx.hoist.jsxScopeMap.set(path.node, jsxScope);
+
+  const visitorState: VisitorState = {
+    isImmutable: true,
+    jsxScope,
+    targetScope: path.scope.getProgramParent(),
+  };
+
+  path.traverse(hoistingVisitor as Visitor<VisitorState>, visitorState);
+  if (!visitorState.isImmutable) return;
+
+  const { targetScope } = visitorState;
+  // Only hoist if it would give us an advantage.
+  for (let currentScope = jsxScope; ;) {
+    if (targetScope === currentScope) return;
+    if (isHoistingScope(currentScope)) break;
+
+    currentScope = currentScope.parent;
+  }
+
+  const id = path.scope.generateUidIdentifierBasedOnNode(name);
+
+  targetScope.push({ id });
+  // If the element is to be hoisted, update HOISTED to be the target scope
+  ctx.hoist.jsxScopeMap.set(path.node, targetScope);
+  ctx.hoist.hoisted.add(path.node);
+
+  let replacement: t.Expression | t.JSXExpressionContainer = t.addComment(
+    t.logicalExpression(
+      '||',
+      id,
+      t.assignmentExpression(
+        '=',
+        id,
+        path.node,
+      ),
+    ),
+    'leading',
+    '@forgetti hoisted_jsx',
+    false,
+  );
+
+  if (
+    path.parentPath.isJSXElement()
+    || path.parentPath.isJSXAttribute()
+  ) {
+    replacement = t.jsxExpressionContainer(replacement);
+  }
+
+  // console.log(path.node, replacement.type, replacement);
+
+  path.replaceWith(replacement);
+}

packages/forgetti/src/core/optimize-jsx.ts

@@ -3,7 +3,7 @@ import * as t from '@babel/types';
 import type { ComponentNode, StateContext } from './types';
 import getImportIdentifier from './get-import-identifier';
 import { RUNTIME_MEMO } from './imports';
-import { shouldSkipJSX, isPathValid } from './checks';
+import { shouldSkipJSX, isPathValid, shouldSkipJSXExtraction } from './checks';
 import type { ImportDefinition } from './presets';
 
 interface JSXReplacement {
@@ -38,6 +38,9 @@ function extractJSXExpressions(
 ): void {
   // Iterate attributes
   if (isPathValid(path, t.isJSXElement)) {
+    if (shouldSkipJSXExtraction(path.node)) {
+      return;
+    }
     const openingElement = path.get('openingElement');
     const openingName = openingElement.get('name');
     const trueOpeningName = getJSXIdentifier(openingName);
@@ -128,10 +131,15 @@ function extractJSXExpressions(
   for (let i = 0, len = children.length; i < len; i++) {
     const child = children[i];
 
+    if (shouldSkipJSXExtraction(child.node)) {
+      continue;
+    }
+
     if (isPathValid(child, t.isJSXElement) || isPathValid(child, t.isJSXFragment)) {
       extractJSXExpressions(child, state, false);
     } else if (isPathValid(child, t.isJSXExpressionContainer)) {
       const expr = child.get('expression');
+
       if (isPathValid(expr, t.isJSXElement) || isPathValid(expr, t.isJSXFragment)) {
         extractJSXExpressions(expr, state, false);
       } else if (isPathValid(expr, t.isExpression)) {
@@ -172,6 +180,9 @@ function transformJSX(
   if (shouldSkipJSX(path.node)) {
     return;
   }
+  if (ctx.hoist.hoisted.has(path.node)) {
+    return;
+  }
   const state: State = {
     source: path.scope.generateUidIdentifier('values'),
     expressions: [],

packages/forgetti/src/core/optimizer.ts

@@ -1,7 +1,9 @@
 /* eslint-disable @typescript-eslint/no-unnecessary-condition */
 import type * as babel from '@babel/core';
 import * as t from '@babel/types';
-import { isNestedExpression, shouldSkipNode, isPathValid } from './checks';
+import {
+  isNestedExpression, shouldSkipNode, isPathValid, shouldSkipJSXExtraction,
+} from './checks';
 import getForeignBindings, { isForeignBinding } from './get-foreign-bindings';
 import getImportIdentifier from './get-import-identifier';
 import { RUNTIME_EQUALS } from './imports';
@@ -877,6 +879,9 @@ export default class Optimizer {
     if (shouldSkipNode(path.node)) {
       return optimizedExpr(path.node, undefined, true);
     }
+    if (shouldSkipJSXExtraction(path.node)) {
+      return optimizedExpr(path.node, undefined, true);
+    }
     if (isPathValid(path, isNestedExpression)) {
       return this.optimizeExpression(path.get('expression'));
     }

packages/forgetti/src/core/presets.ts

@@ -39,6 +39,9 @@ export interface Preset {
     hooks: HookRegistration[];
     hocs: ImportDefinition[];
   };
+  optimizations?: {
+    hoistConstantJsx?: boolean;
+  };
 }
 
 export interface Options {
@@ -125,6 +128,9 @@ export const PRESETS = {
         },
       ],
     },
+    optimizations: {
+      hoistConstantJsx: true,
+    },
   }),
   preact: createPreset({
     filters: {
@@ -219,5 +225,8 @@ export const PRESETS = {
         },
       ],
     },
+    optimizations: {
+      hoistConstantJsx: true,
+    },
   }),
 };

packages/forgetti/src/core/types.ts

@@ -1,5 +1,6 @@
 import type * as t from '@babel/types';
 import type * as babel from '@babel/core';
+import type { Scope } from '@babel/traverse';
 import type {
   HookRegistration,
   ImportDefinition,
@@ -33,6 +34,10 @@ export interface StateContext {
     component: RegExp;
     hook?: RegExp;
   };
+  hoist: {
+    jsxScopeMap: WeakMap<t.JSX, Scope>;
+    hoisted: WeakSet<t.JSX>;
+  };
 }
 
 export interface OptimizedExpression {

packages/forgetti/src/index.ts

@@ -1,5 +1,6 @@
 /* eslint-disable @typescript-eslint/no-unnecessary-condition */
 import type * as babel from '@babel/core';
+import type { Scope } from '@babel/traverse';
 import * as t from '@babel/types';
 import {
   isComponent,
@@ -23,6 +24,7 @@ import unwrapPath from './core/unwrap-path';
 import { expandExpressions } from './core/expand-expressions';
 import { inlineExpressions } from './core/inline-expressions';
 import { simplifyExpressions } from './core/simplify-expressions';
+import hoistConstantJSX from './core/hoist-jsx';
 import optimizeJSX from './core/optimize-jsx';
 
 export type { Options };
@@ -267,13 +269,22 @@ export default function forgettiPlugin(): babel.PluginObj<State> {
               ? new RegExp(preset.filters.hook.source, preset.filters.hook.flags)
               : undefined,
           },
+          hoist: {
+            jsxScopeMap: new WeakMap<t.JSX, Scope>(),
+            hoisted: new WeakSet<t.JSX>(),
+          },
         };
 
-        // Register all import specifiers
         programPath.traverse({
+          // Register all import specifiers
           ImportDeclaration(path) {
             extractImportIdentifiers(ctx, path);
           },
+          JSXElement(path) {
+            if (preset.optimizations?.hoistConstantJsx) {
+              hoistConstantJSX(ctx, path);
+            }
+          },
         });
 
         programPath.traverse({