allowMutablePropsOnTags: cache JSX constant elements with function props

packages/babel-plugin-transform-react-constant-elements/src/index.ts

@@ -41,17 +41,18 @@ export default declare((api, options) => {
     return scope;
   }
 
-  const analyzer = {
+  const immutabilityVisitor = {
     enter(path, state) {
       const stop = () => {
         state.isImmutable = false;
         path.stop();
       };
 
-      if (path.isJSXClosingElement()) {
+      const skip = () => {
         path.skip();
-        return;
-      }
+      };
+
+      if (path.isJSXClosingElement()) return skip();
 
       // Elements with refs are not safe to hoist.
       if (
@@ -61,54 +62,70 @@ export default declare((api, options) => {
         return stop();
       }
 
-      // Ignore identifiers & JSX expressions.
+      // Ignore JSX expressions and immutable values.
       if (
         path.isJSXIdentifier() ||
         path.isJSXMemberExpression() ||
-        path.isJSXNamespacedName()
+        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.isImmutable()) {
-        // 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 (like objects) could be mutated after render.
-        // https://github.com/facebook/react/issues/3226
-        if (path.isPure()) {
-          const expressionResult = path.evaluate();
-          if (expressionResult.confident) {
-            // We know the result; check its mutability.
-            const { value } = expressionResult;
-            const isMutable =
-              (!state.mutablePropsAllowed &&
-                value &&
-                typeof value === "object") ||
-              typeof value === "function";
-            if (!isMutable) {
-              // It evaluated to an immutable value, so we can hoist it.
-              path.skip();
-              return;
-            }
-          } else if (t.isIdentifier(expressionResult.deopt)) {
-            // 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;
-          }
+      // If we allow mutable props, tags with function expressions can be
+      // safely hoisted.
+      const { mutablePropsAllowed } = state;
+      if (mutablePropsAllowed && path.isFunction()) {
+        path.traverse(targetScopeVisitor, state);
+        return skip();
+      }
+
+      if (!path.isPure()) return stop();
+
+      // 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 (like objects) 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.
+        const { value } = expressionResult;
+        if (
+          mutablePropsAllowed ||
+          value === null ||
+          (typeof value !== "object" && typeof value !== "function")
+        ) {
+          // It evaluated to an immutable value, so we can hoist it.
+          return skip();
         }
-        stop();
+      } else if (t.isIdentifier(expressionResult.deopt)) {
+        // 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();
     },
+  };
 
+  const targetScopeVisitor = {
     ReferencedIdentifier(path, state) {
       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;
+      }
+
       while (scope) {
         // We cannot hoist outside of the previous hoisting target
         // scope, so we return early and we don't update it.
@@ -124,72 +141,13 @@ export default declare((api, options) => {
 
       state.targetScope = getHoistingScope(scope);
     },
-
-    /*
-    See the discussion at https://github.com/babel/babel/pull/12967#discussion_r587948958
-    to uncomment this code.
-
-    ReferencedIdentifier(path, state) {
-      const { node } = path;
-      let { scope } = path;
-      let targetScope;
-
-      let isNestedScope = true;
-      let needsHoisting = true;
-
-      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;
-
-        // When we hit the scope of our JSX element, we must start
-        // checking if they declare the binding of the current
-        // ReferencedIdentifier.
-        // We don't case about bindings declared in nested scopes,
-        // because the whole nested scope is hoisted alongside the
-        // JSX element so it doesn't impose any extra constraint.
-        if (scope === state.jsxScope) {
-          isNestedScope = false;
-        }
-
-        // If we are in an upper scope and hoisting to this scope has
-        // any benefit, we update the possible targetScope to the
-        // current one.
-        if (!isNestedScope && needsHoisting) {
-          targetScope = scope;
-        }
-
-        // When we start walking in upper scopes, avoid hoisting JSX
-        // elements until we hit a scope introduced by a function or
-        // loop.
-        // This is because hoisting from the inside to the outside
-        // of block or if statements doesn't give any performance
-        // benefit, and it just unnecessarily increases the code size.
-        if (scope === state.jsxScope) {
-          needsHoisting = false;
-        }
-        if (!needsHoisting && isHoistingScope(scope)) {
-          needsHoisting = true;
-        }
-
-        // If the current scope declares the ReferencedIdentifier we
-        // are checking, we break out of this loop. There are two
-        // possible scenarios:
-        //  1. We are in a nested scope, this this declaration means
-        //     that this reference doesn't affect the target scope.
-        //     The targetScope variable is still undefined.
-        //  2. We are in an upper scope, so this declaration defines
-        //     a new hoisting constraint. The targetScope variable
-        //     refers to the current scope.
-        if (declares(node, scope)) break;
-
-        scope = scope.parent;
-      }
-
-      if (targetScope) state.targetScope = targetScope;
-    },*/
   };
 
+  // We cannot use traverse.visitors.merge because it doesn't support
+  // immutabilityVisitor's bare `enter` visitor.
+  // It's safe to just use ... because the two visitors don't share any key.
+  const hoistingVisitor = { ...immutabilityVisitor, ...targetScopeVisitor };
+
   return {
     name: "transform-react-constant-elements",
 
@@ -216,21 +174,6 @@ export default declare((api, options) => {
           mutablePropsAllowed = allowMutablePropsOnTags.includes(elementName);
         }
 
-        const state = {
-          isImmutable: true,
-          mutablePropsAllowed,
-          targetScope: path.scope.getProgramParent(),
-        };
-
-        // Traverse all props passed to this element for immutability,
-        // and compute the target hoisting scope
-        path.traverse(analyzer, state);
-
-        if (!state.isImmutable) return;
-
-        const { targetScope } = state;
-        HOISTED.set(path.node, targetScope);
-
         // 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
@@ -243,6 +186,18 @@ export default declare((api, options) => {
         }
         jsxScope ??= getHoistingScope(path.scope);
 
+        const visitorState = {
+          isImmutable: true,
+          mutablePropsAllowed,
+          jsxScope,
+          targetScope: path.scope.getProgramParent(),
+        };
+        path.traverse(hoistingVisitor, visitorState);
+        if (!visitorState.isImmutable) return;
+
+        const { targetScope } = visitorState;
+        HOISTED.set(path.node, targetScope);
+
         // Only hoist if it would give us an advantage.
         if (targetScope === jsxScope) return;
 

packages/babel-plugin-transform-react-constant-elements/test/fixtures/allowMutablePropsOnTags/function-closure/input.js

@@ -0,0 +1,3 @@
+function Component({ increment }) {
+  return () => <Counter onClick={value => value + increment} />;
+}

packages/babel-plugin-transform-react-constant-elements/test/fixtures/allowMutablePropsOnTags/function-closure/options.json

@@ -0,0 +1,9 @@
+{
+  "plugins": [
+    [
+      "transform-react-constant-elements",
+      { "allowMutablePropsOnTags": ["Counter"] }
+    ],
+    "syntax-jsx"
+  ]
+}

packages/babel-plugin-transform-react-constant-elements/test/fixtures/allowMutablePropsOnTags/function-closure/output.js

@@ -0,0 +1,7 @@
+function Component({
+  increment
+}) {
+  var _Counter;
+
+  return () => _Counter || (_Counter = <Counter onClick={value => value + increment} />);
+}

packages/babel-plugin-transform-react-constant-elements/test/fixtures/allowMutablePropsOnTags/function-impure-body/input.js

@@ -0,0 +1,3 @@
+function Component() {
+  return () => <Counter onClick={value => value + prompt("Increment:")} />;
+}

packages/babel-plugin-transform-react-constant-elements/test/fixtures/allowMutablePropsOnTags/function-impure-body/options.json

@@ -0,0 +1,9 @@
+{
+  "plugins": [
+    [
+      "transform-react-constant-elements",
+      { "allowMutablePropsOnTags": ["Counter"] }
+    ],
+    "syntax-jsx"
+  ]
+}

packages/babel-plugin-transform-react-constant-elements/test/fixtures/allowMutablePropsOnTags/function-impure-body/output.js

@@ -0,0 +1,5 @@
+var _Counter;
+
+function Component() {
+  return () => _Counter || (_Counter = <Counter onClick={value => value + prompt("Increment:")} />);
+}

packages/babel-plugin-transform-react-constant-elements/test/fixtures/allowMutablePropsOnTags/function-same-binding-name/input.js

@@ -0,0 +1,3 @@
+function Component({ value }) {
+  return () => <Counter onClick={value => value + 1} />;
+}

packages/babel-plugin-transform-react-constant-elements/test/fixtures/allowMutablePropsOnTags/function-same-binding-name/options.json

@@ -0,0 +1,9 @@
+{
+  "plugins": [
+    [
+      "transform-react-constant-elements",
+      { "allowMutablePropsOnTags": ["Counter"] }
+    ],
+    "syntax-jsx"
+  ]
+}

packages/babel-plugin-transform-react-constant-elements/test/fixtures/allowMutablePropsOnTags/function-same-binding-name/output.js

@@ -0,0 +1,7 @@
+var _Counter;
+
+function Component({
+  value
+}) {
+  return () => _Counter || (_Counter = <Counter onClick={value => value + 1} />);
+}

packages/babel-plugin-transform-react-constant-elements/test/fixtures/allowMutablePropsOnTags/function/input.js

@@ -0,0 +1,3 @@
+function Component() {
+  return () => <Counter onClick={value => value + 1} />;
+}

packages/babel-plugin-transform-react-constant-elements/test/fixtures/allowMutablePropsOnTags/function/options.json

@@ -0,0 +1,9 @@
+{
+  "plugins": [
+    [
+      "transform-react-constant-elements",
+      { "allowMutablePropsOnTags": ["Counter"] }
+    ],
+    "syntax-jsx"
+  ]
+}

packages/babel-plugin-transform-react-constant-elements/test/fixtures/allowMutablePropsOnTags/function/output.js

@@ -0,0 +1,5 @@
+var _Counter;
+
+function Component() {
+  return () => _Counter || (_Counter = <Counter onClick={value => value + 1} />);
+}

packages/babel-plugin-transform-react-constant-elements/test/fixtures/allowMutablePropsOnTags/iife/input.js

@@ -0,0 +1,3 @@
+function Component() {
+  return () => <Counter init={(value => value + prompt("Increment:"))(2)} />;
+}

packages/babel-plugin-transform-react-constant-elements/test/fixtures/allowMutablePropsOnTags/iife/options.json

@@ -0,0 +1,9 @@
+{
+  "plugins": [
+    [
+      "transform-react-constant-elements",
+      { "allowMutablePropsOnTags": ["Counter"] }
+    ],
+    "syntax-jsx"
+  ]
+}

packages/babel-plugin-transform-react-constant-elements/test/fixtures/allowMutablePropsOnTags/iife/output.js

@@ -0,0 +1,3 @@
+function Component() {
+  return () => <Counter init={(value => value + prompt("Increment:"))(2)} />;
+}

packages/babel-plugin-transform-react-constant-elements/test/fixtures/constant-elements/pure-expression-whitelist-deopt/input.js → packages/babel-plugin-transform-react-constant-elements/test/fixtures/allowMutablePropsOnTags/not-allowed-element/input.js

@@ -1,5 +1,5 @@
-// This is the same as the pure-expression-whitelist test, but 'FormattedMessage' is *not* whitelisted
-// so it should not be hoisted.
+// This is the same as the basic test, but 'FormattedMessage' is *not* one of the
+// allowed tags so it should not be hoisted.
 var Foo = React.createClass({
   render: function () {
     return (

packages/babel-plugin-transform-react-constant-elements/test/fixtures/constant-elements/pure-expression-whitelist-deopt/output.js → packages/babel-plugin-transform-react-constant-elements/test/fixtures/allowMutablePropsOnTags/not-allowed-element/output.js

@@ -1,5 +1,5 @@
-// This is the same as the pure-expression-whitelist test, but 'FormattedMessage' is *not* whitelisted
-// so it should not be hoisted.
+// This is the same as the basic test, but 'FormattedMessage' is *not* one of the
+// allowed tags so it should not be hoisted.
 var Foo = React.createClass({
   render: function () {
     return <FormattedMessage id="someMessage.foo" defaultMessage={"Some text, " + "and some more too. {someValue}"} description="A test message for babel." values={{

packages/babel-plugin-transform-react-constant-elements/test/fixtures/constant-elements/for-loop/output.js

@@ -7,4 +7,3 @@ function render() {
 
   return nodes;
 }
-