fix: avoid traversing on Program:enter (facebook#786)

* fix: avoid traversing on Program:enter
* chore: add a test for pre-plugins

packages/babel-plugin/__tests__/stylex-transform-pre-plugin-test.js

@@ -0,0 +1,149 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ *
+ */
+
+'use strict';
+
+jest.autoMockOff();
+
+import * as t from '@babel/types';
+const { transformSync } = require('@babel/core');
+const stylexPlugin = require('../src/index');
+
+// A simple babel plugin that looks for inline() calls
+// hoists them out to stylex.create calls
+// and replaces them with references.
+function inlineDemoPlugin() {
+  const collectedStyles = [];
+  let stylesIdentifier;
+  return {
+    visitor: {
+      Program: {
+        enter(path, _state) {
+          // get unique identifier
+          stylesIdentifier = path.scope.generateUidIdentifier('styles');
+          path.traverse({
+            CallExpression(path, _state) {
+              if (
+                path.node.callee.type === 'Identifier' &&
+                path.node.callee.name === 'inline'
+              ) {
+                const args = path.node.arguments;
+                if (args.length !== 1) {
+                  return;
+                }
+                const index = collectedStyles.length;
+                collectedStyles.push(args[0]);
+                path.replaceWith(
+                  t.memberExpression(
+                    stylesIdentifier,
+                    t.identifier(`$${index}`),
+                  ),
+                );
+              }
+            },
+          });
+          path.node.body.push(
+            t.variableDeclaration('const', [
+              t.variableDeclarator(
+                stylesIdentifier,
+                t.callExpression(
+                  t.memberExpression(
+                    t.identifier('stylex'),
+                    t.identifier('create'),
+                  ),
+                  [
+                    t.objectExpression(
+                      collectedStyles.map((style, index) =>
+                        t.objectProperty(t.identifier(`$${index}`), style),
+                      ),
+                    ),
+                  ],
+                ),
+              ),
+            ]),
+          );
+        },
+      },
+    },
+  };
+}
+
+function transform(source, opts = {}) {
+  return transformSync(source, {
+    filename: opts.filename,
+    parserOpts: {
+      flow: 'all',
+    },
+    babelrc: false,
+    plugins: [
+      ['babel-plugin-syntax-hermes-parser', { flow: 'detect' }],
+      inlineDemoPlugin,
+      [
+        stylexPlugin,
+        {
+          runtimeInjection: true,
+          unstable_moduleResolution: { type: 'haste' },
+          ...opts,
+        },
+      ],
+    ],
+  }).code;
+}
+
+describe('[transform] stylex.create()', () => {
+  test('transforms style object', () => {
+    expect(
+      transform(`
+        import stylex from 'stylex';
+
+        function Demo() {
+          return (
+            <div>
+              <button {...stylex.props(
+                styles.default,
+                inline({
+                  backgroundColor: 'pink',
+                  color: 'white',
+                })
+              )}>
+                Hello
+              </button>
+            </div>
+          );
+        }
+
+        const styles = stylex.create({
+          default: {
+            appearance: 'none',
+            borderWidth: '0',
+            borderStyle: 'none',
+          }
+        });
+      `),
+    ).toMatchInlineSnapshot(`
+      "import _inject from "@stylexjs/stylex/lib/stylex-inject";
+      var _inject2 = _inject;
+      import stylex from 'stylex';
+      function Demo() {
+        return <div>
+                    <button {...{
+            className: "xjyslct xc342km xng3xce x6tqnqi x1awj2ng"
+          }}>
+                      Hello
+                    </button>
+                  </div>;
+      }
+      _inject2(".xjyslct{appearance:none}", 3000);
+      _inject2(".xc342km{border-width:0}", 2000);
+      _inject2(".xng3xce{border-style:none}", 2000);
+      _inject2(".x6tqnqi{background-color:pink}", 3000);
+      _inject2(".x1awj2ng{color:white}", 3000);"
+    `);
+  });
+});

packages/babel-plugin/src/index.js

@@ -69,28 +69,48 @@ function styleXTransform(): PluginObj<> {
               }
             }
           }
-
-          path.traverse({
-            // Look for stylex-related function calls and transform them.
-            // e.g.
-            //   stylex.create(...)
-            //   stylex.keyframes(...)
-            CallExpression(path: NodePath<t.CallExpression>) {
-              if (pathUtils.isVariableDeclarator(path.parentPath)) {
-                // # Look for `stylex.keyframes` calls
-                //   Needs to be handled *before* `stylex.create` as the `create` call
-                //   may use the generated animation name.
-                transformStyleXKeyframes(path.parentPath, state);
-              }
-              transformStyleXDefineVars(path, state);
-              transformStyleXCreateTheme(path, state);
-              transformStyleXCreate(path, state);
-            },
-          });
         },
         // After all other visitors are done, we can remove `styles=stylex.create(...)`
         // variables entirely if they're not needed.
         exit: (path: NodePath<t.Program>) => {
+          path.traverse({
+            Identifier(path: NodePath<t.Identifier>) {
+              // Look for variables bound to `stylex.create` calls that are used
+              // outside of `stylex(...)` calls
+              if (pathUtils.isReferencedIdentifier(path)) {
+                const { name } = path.node;
+                if (state.styleMap.has(name)) {
+                  const parentPath = path.parentPath;
+                  if (pathUtils.isMemberExpression(parentPath)) {
+                    const { property, computed } = parentPath.node;
+                    if (property.type === 'Identifier' && !computed) {
+                      state.markComposedNamespace([name, property.name, true]);
+                    } else if (property.type === 'StringLiteral' && computed) {
+                      state.markComposedNamespace([name, property.value, true]);
+                    } else if (property.type === 'NumericLiteral' && computed) {
+                      state.markComposedNamespace([
+                        name,
+                        String(property.value),
+                        true,
+                      ]);
+                    } else {
+                      state.markComposedNamespace([name, true, true]);
+                    }
+                  } else {
+                    state.markComposedNamespace([name, true, true]);
+                  }
+                }
+              }
+            },
+            CallExpression(path: NodePath<t.CallExpression>) {
+              // Don't traverse the children of `stylex(...)` calls.
+              // This is important for detecting which `stylex.create()` calls
+              // should be kept.
+              skipStylexMergeChildren(path, state);
+              skipStylexPropsChildren(path, state);
+              skipStylexAttrsChildren(path, state);
+            },
+          });
           path.traverse({
             CallExpression(path: NodePath<t.CallExpression>) {
               transformStylexCall(path, state);
@@ -217,41 +237,15 @@ function styleXTransform(): PluginObj<> {
       },
 
       CallExpression(path: NodePath<t.CallExpression>) {
-        // Don't traverse the children of `stylex(...)` calls.
-        // This is important for detecting which `stylex.create()` calls
-        // should be kept.
-        skipStylexMergeChildren(path, state);
-        skipStylexPropsChildren(path, state);
-        skipStylexAttrsChildren(path, state);
-      },
-
-      Identifier(path: NodePath<t.Identifier>) {
-        // Look for variables bound to `stylex.create` calls that are used
-        // outside of `stylex(...)` calls
-        if (pathUtils.isReferencedIdentifier(path)) {
-          const { name } = path.node;
-          if (state.styleMap.has(name)) {
-            const parentPath = path.parentPath;
-            if (pathUtils.isMemberExpression(parentPath)) {
-              const { property, computed } = parentPath.node;
-              if (property.type === 'Identifier' && !computed) {
-                state.markComposedNamespace([name, property.name, true]);
-              } else if (property.type === 'StringLiteral' && computed) {
-                state.markComposedNamespace([name, property.value, true]);
-              } else if (property.type === 'NumericLiteral' && computed) {
-                state.markComposedNamespace([
-                  name,
-                  String(property.value),
-                  true,
-                ]);
-              } else {
-                state.markComposedNamespace([name, true, true]);
-              }
-            } else {
-              state.markComposedNamespace([name, true, true]);
-            }
-          }
+        if (pathUtils.isVariableDeclarator(path.parentPath)) {
+          // # Look for `stylex.keyframes` calls
+          //   Needs to be handled *before* `stylex.create` as the `create` call
+          //   may use the generated animation name.
+          transformStyleXKeyframes(path.parentPath, state);
         }
+        transformStyleXDefineVars(path, state);
+        transformStyleXCreateTheme(path, state);
+        transformStyleXCreate(path, state);
       },
     },
   };