@usels/babel-plugin-legend-memo 0.0.1-beta.3

package/src/utils/hasGetCall.ts

@@ -0,0 +1,93 @@
+import type { NodePath } from '@babel/core';
+import type { PluginOptions } from '../types';
+import { getRootObject } from './getRootObject';
+
+function isGetCallNode(node: any, opts: PluginOptions): boolean {
+  if (node.type === 'CallExpression') {
+    const { callee, arguments: args } = node;
+    if (callee.type !== 'MemberExpression') return false;
+    if (callee.property.type !== 'Identifier') return false;
+    if (callee.property.name !== 'get') return false;
+    if (args.length !== 0) return false;
+    if (!opts.allGet) {
+      const root = getRootObject(callee.object);
+      if (!root || !root.name?.endsWith('$')) return false;
+    }
+    return true;
+  }
+  if (node.type === 'OptionalCallExpression') {
+    const { callee, arguments: args } = node;
+    if (callee.type !== 'OptionalMemberExpression') return false;
+    if (callee.property?.type !== 'Identifier') return false;
+    if (callee.property?.name !== 'get') return false;
+    if (args.length !== 0) return false;
+    if (!opts.allGet) {
+      const root = getRootObject(callee.object);
+      if (!root || !root.name?.endsWith('$')) return false;
+    }
+    return true;
+  }
+  return false;
+}
+
+/**
+ * Checks if a NodePath contains a zero-argument .get() call on a $-suffixed observable.
+ *
+ * IMPORTANT: path.traverse() visits CHILDREN, not the root node itself.
+ * We must check the root node separately.
+ *
+ * - Stops traversal at function boundaries (ArrowFunctionExpression, FunctionExpression, FunctionDeclaration)
+ * - Supports OptionalCallExpression: obs$?.get()
+ * - Ignores .get(key) calls with arguments (Map.get(key))
+ * - When opts.allGet is false (default), root object must end with '$'
+ */
+const FUNCTION_BOUNDARY_TYPES = new Set([
+  'ArrowFunctionExpression',
+  'FunctionExpression',
+  'FunctionDeclaration',
+]);
+
+export function hasGetCall(exprPath: NodePath, opts: PluginOptions): boolean {
+  // If the root node itself is a function boundary, do not descend into it.
+  // path.traverse() visits CHILDREN only, so the ArrowFunctionExpression visitor
+  // inside traverse() would NOT fire for the root — leading to false positives.
+  if (FUNCTION_BOUNDARY_TYPES.has(exprPath.node.type)) {
+    return false;
+  }
+
+  // Check the root node itself first (traverse does NOT visit it)
+  if (isGetCallNode(exprPath.node, opts)) {
+    return true;
+  }
+
+  let found = false;
+
+  exprPath.traverse({
+    // Stop at function boundaries — don't wrap event handlers or hook callbacks
+    ArrowFunctionExpression(innerPath) {
+      innerPath.skip();
+    },
+    FunctionExpression(innerPath) {
+      innerPath.skip();
+    },
+    FunctionDeclaration(innerPath) {
+      innerPath.skip();
+    },
+
+    OptionalCallExpression(innerPath) {
+      if (isGetCallNode(innerPath.node, opts)) {
+        found = true;
+        innerPath.stop();
+      }
+    },
+
+    CallExpression(innerPath) {
+      if (isGetCallNode(innerPath.node, opts)) {
+        found = true;
+        innerPath.stop();
+      }
+    },
+  });
+
+  return found;
+}

package/src/utils/index.ts

@@ -0,0 +1,9 @@
+export { getRootObject } from './getRootObject';
+export { hasGetCall } from './hasGetCall';
+export { hasAttributeGetCall } from './hasAttributeGetCall';
+export { isInsideReactiveContext } from './isInsideReactiveContext';
+export { isInsideObserverHOC } from './isInsideObserverHOC';
+export { isInsideAttribute } from './isInsideAttribute';
+export { createAutoElement } from './createAutoElement';
+export { addAutoImport } from './addAutoImport';
+export { wrapChildrenAsFunction } from './wrapChildrenAsFunction';

package/src/utils/isInsideAttribute.ts

@@ -0,0 +1,12 @@
+import type { NodePath } from '@babel/core';
+
+/**
+ * Returns true if the path is inside a JSXAttribute.
+ * Used by JSXExpressionContainer visitor to skip attribute expressions
+ * (those are handled by the JSXElement visitor instead).
+ */
+export function isInsideAttribute(path: NodePath): boolean {
+  return (
+    path.findParent((p) => p.isJSXAttribute()) !== null
+  );
+}

package/src/utils/isInsideObserverHOC.ts

@@ -0,0 +1,24 @@
+import type { NodePath } from '@babel/core';
+
+/**
+ * Returns true if the path is inside an observer() HOC call expression.
+ * observer() wraps a function component, making it reactive — no Auto wrapping needed.
+ *
+ * Unlike reactive JSX components (Auto, For, etc.), observer() is a CallExpression
+ * wrapper, so it cannot be detected by isInsideReactiveContext.
+ *
+ * Default observerNames: ["observer"]
+ * Configurable via opts.observerNames: ["observer", "reactive", ...]
+ */
+export function isInsideObserverHOC(
+  path: NodePath,
+  observerNames: Set<string>,
+): boolean {
+  return (
+    path.findParent((p) => {
+      if (!p.isCallExpression()) return false;
+      const callee = p.node.callee;
+      return callee.type === 'Identifier' && observerNames.has(callee.name);
+    }) !== null
+  );
+}

package/src/utils/isInsideReactiveContext.ts

@@ -0,0 +1,23 @@
+import type { NodePath } from '@babel/core';
+
+/**
+ * Returns true if the path has a parent JSXElement whose tag name is in reactiveComponents.
+ * Used to skip wrapping inside Auto, For, Show, Memo, Computed, Switch, etc.
+ *
+ * Uses direct node.type comparison instead of t.isJSXIdentifier() to avoid
+ * needing to pass `t` as a parameter.
+ */
+export function isInsideReactiveContext(
+  path: NodePath,
+  reactiveComponents: Set<string>,
+): boolean {
+  return (
+    path.findParent((p) => {
+      if (!p.isJSXElement()) return false;
+      const name = p.node.openingElement.name;
+      return (
+        name.type === 'JSXIdentifier' && reactiveComponents.has(name.name)
+      );
+    }) !== null
+  );
+}

package/src/utils/wrapChildrenAsFunction.ts

@@ -0,0 +1,113 @@
+import type { types as BabelTypes } from '@babel/core';
+import type {
+  Expression,
+  JSXChild,
+  JSXElement,
+  JSXExpressionContainer,
+} from '@babel/types';
+
+/**
+ * Expression types that indicate children are already function-like.
+ * Wrapping these again would be incorrect.
+ *
+ * - ArrowFunctionExpression: {() => ...} — already wrapped
+ * - FunctionExpression: {function() {...}} — already wrapped
+ * - MemberExpression: {obj.render} — already a reference
+ * - Identifier: {renderFn} — already a reference
+ */
+const ALREADY_FUNCTION_TYPES = new Set([
+  'ArrowFunctionExpression',
+  'FunctionExpression',
+  'MemberExpression',
+  'Identifier',
+]);
+
+/** Strip whitespace-only JSXText nodes (newlines/spaces between tags) */
+function filterEmptyText(children: JSXChild[]): JSXChild[] {
+  return children.filter(
+    (c) => !(c.type === 'JSXText' && c.value.trim().length === 0),
+  );
+}
+
+/** Check if the (single) child is already a function/reference — no wrapping needed */
+function areChildrenAlreadyFunction(children: JSXChild[]): boolean {
+  if (children.length !== 1) return false;
+  const child = children[0];
+  if (child.type !== 'JSXExpressionContainer') return false;
+  const expr = (child as JSXExpressionContainer).expression;
+  return expr.type !== 'JSXEmptyExpression' && ALREADY_FUNCTION_TYPES.has(expr.type);
+}
+
+/**
+ * Build the arrow function body from filtered children.
+ *
+ * - Single JSXElement child  → use it directly as body: () => <Foo />
+ * - Single expression child  → unwrap from container: () => someExpr
+ * - Multiple children        → wrap in Fragment: () => <><A /><B /></>
+ */
+function buildArrowBody(
+  t: typeof BabelTypes,
+  children: JSXChild[],
+): Expression {
+  if (children.length === 1) {
+    const child = children[0];
+    if (child.type === 'JSXElement') return child as JSXElement;
+    if (child.type === 'JSXExpressionContainer') {
+      const expr = (child as JSXExpressionContainer).expression;
+      if (expr.type !== 'JSXEmptyExpression') return expr as Expression;
+    }
+  }
+  return t.jsxFragment(
+    t.jsxOpeningFragment(),
+    t.jsxClosingFragment(),
+    children as any[],
+  );
+}
+
+/**
+ * Wraps non-function children of a JSXElement in an arrow function.
+ *
+ * Before: <Memo>{count$.get()}</Memo>
+ * After:  <Memo>{() => count$.get()}</Memo>
+ *
+ * Returns the new JSXElement node, or null if no transformation is needed
+ * (children are already a function, or there are no meaningful children).
+ */
+export function wrapChildrenAsFunction(
+  t: typeof BabelTypes,
+  node: JSXElement,
+): JSXElement | null {
+  // Only handle simple JSXIdentifier element names (not <Foo.Bar> etc.)
+  if (node.openingElement.name.type !== 'JSXIdentifier') return null;
+
+  const children = filterEmptyText(node.children);
+
+  if (children.length === 0) return null;
+  if (areChildrenAlreadyFunction(children)) return null;
+
+  // Only wrap if first child is a JSXElement or a non-function expression
+  const firstChild = children[0];
+  const needsWrapping =
+    firstChild.type === 'JSXElement' ||
+    (firstChild.type === 'JSXExpressionContainer' &&
+      !ALREADY_FUNCTION_TYPES.has(
+        (firstChild as JSXExpressionContainer).expression.type,
+      ));
+
+  if (!needsWrapping) return null;
+
+  const elementName = node.openingElement.name.name;
+  const body = buildArrowBody(t, children);
+  const arrowFn = t.arrowFunctionExpression([], body);
+
+  return t.jsxElement(
+    t.jsxOpeningElement(
+      t.jsxIdentifier(elementName),
+      node.openingElement.attributes,
+      false,
+    ),
+    t.jsxClosingElement(t.jsxIdentifier(elementName)),
+    [t.jsxExpressionContainer(arrowFn)],
+    false,
+  );
+}

package/src/visitors/index.ts

@@ -0,0 +1,3 @@
+export { createProgramVisitor } from './program';
+export { createJSXElementVisitor } from './jsxElement';
+export { createJSXExpressionContainerVisitor } from './jsxExpressionContainer';

package/src/visitors/jsxElement.ts

@@ -0,0 +1,52 @@
+import type { NodePath, types as BabelTypes } from '@babel/core';
+import type { JSXElement } from '@babel/types';
+import type { PluginState } from '../types';
+import { hasAttributeGetCall } from '../utils/hasAttributeGetCall';
+import { isInsideReactiveContext } from '../utils/isInsideReactiveContext';
+import { isInsideObserverHOC } from '../utils/isInsideObserverHOC';
+import { createAutoElement } from '../utils/createAutoElement';
+import { wrapChildrenAsFunction } from '../utils/wrapChildrenAsFunction';
+
+export function createJSXElementVisitor(t: typeof BabelTypes) {
+  return function JSXElement(
+    path: NodePath<JSXElement>,
+    state: PluginState,
+  ): void {
+    // STEP 1 (NEW): If this element is Memo/Show/Computed (or configured),
+    // auto-wrap non-function children in () => — equivalent to @legendapp/state/babel
+    const elementName =
+      path.node.openingElement.name.type === 'JSXIdentifier'
+        ? path.node.openingElement.name.name
+        : null;
+
+    if (elementName && state.autoWrapChildrenComponents.has(elementName)) {
+      const wrapped = wrapChildrenAsFunction(t, path.node);
+      if (wrapped !== null) {
+        path.replaceWith(wrapped);
+        // path.node now points to the new element — continue to check attributes
+      }
+    }
+
+    // 1. Skip if already inside a reactive context (Auto, For, Show, Memo, etc.)
+    if (isInsideReactiveContext(path, state.reactiveComponents)) return;
+
+    // 2. Skip if inside an observer() HOC
+    if (isInsideObserverHOC(path, state.observerNames)) return;
+
+    // 3. Check if any non-special attributes contain .get() calls
+    if (!hasAttributeGetCall(path, state.opts)) return;
+
+    // 4. Wrap the entire JSXElement in <Auto>{() => element}</Auto>
+    const autoElement = createAutoElement(
+      t,
+      path.node,
+      state.autoComponentName,
+    );
+    path.replaceWith(autoElement);
+
+    state.autoImportNeeded = true;
+
+    // 5. Skip traversal of the new node to prevent double-wrapping
+    path.skip();
+  };
+}

package/src/visitors/jsxExpressionContainer.ts

@@ -0,0 +1,42 @@
+import type { NodePath, types as BabelTypes } from '@babel/core';
+import type { JSXExpressionContainer } from '@babel/types';
+import type { PluginState } from '../types';
+import { hasGetCall } from '../utils/hasGetCall';
+import { isInsideReactiveContext } from '../utils/isInsideReactiveContext';
+import { isInsideObserverHOC } from '../utils/isInsideObserverHOC';
+import { isInsideAttribute } from '../utils/isInsideAttribute';
+import { createAutoElement } from '../utils/createAutoElement';
+
+export function createJSXExpressionContainerVisitor(t: typeof BabelTypes) {
+  return function JSXExpressionContainer(
+    path: NodePath<JSXExpressionContainer>,
+    state: PluginState,
+  ): void {
+    // 1. Skip if inside a JSXAttribute — JSXElement visitor handles those
+    if (isInsideAttribute(path)) return;
+
+    // 2. Skip if already inside a reactive context (Auto, For, Show, Memo, etc.)
+    if (isInsideReactiveContext(path, state.reactiveComponents)) return;
+
+    // 3. Skip if inside an observer() HOC
+    if (isInsideObserverHOC(path, state.observerNames)) return;
+
+    // 4. Skip JSXEmptyExpression {}
+    const expression = path.node.expression;
+    if (expression.type === 'JSXEmptyExpression') return;
+
+    // 5. Check if the expression contains a .get() call
+    const exprPath = path.get('expression') as NodePath;
+    if (!hasGetCall(exprPath, state.opts)) return;
+
+    // 6. Wrap the expression in <Auto>{() => expression}</Auto>
+    const autoElement = createAutoElement(
+      t,
+      expression,
+      state.autoComponentName,
+    );
+    path.replaceWith(autoElement);
+
+    state.autoImportNeeded = true;
+  };
+}

package/src/visitors/program.ts

@@ -0,0 +1,43 @@
+import type { NodePath, types as BabelTypes } from '@babel/core';
+import type { Program } from '@babel/types';
+import type { PluginState } from '../types';
+import { addAutoImport } from '../utils/addAutoImport';
+
+export function createProgramVisitor(t: typeof BabelTypes) {
+  return {
+    enter(path: NodePath<Program>, state: PluginState) {
+      const opts = state.opts ?? {};
+      state.autoImportNeeded = false;
+      state.autoImportSource = opts.importSource ?? '@legendapp/state/react';
+      state.autoComponentName = opts.componentName ?? 'Memo';
+      state.reactiveComponents = new Set([
+        state.autoComponentName,
+        'For',
+        'Show',
+        'Memo',
+        'Computed',
+        'Switch',
+        ...(opts.reactiveComponents ?? []),
+      ]);
+      state.observerNames = new Set([
+        'observer',
+        ...(opts.observerNames ?? []),
+      ]);
+      state.autoWrapChildrenComponents =
+        opts.wrapReactiveChildren !== false
+          ? new Set([
+              'Memo',
+              'Show',
+              'Computed',
+              ...(opts.wrapReactiveChildrenComponents ?? []),
+            ])
+          : new Set(opts.wrapReactiveChildrenComponents ?? []);
+    },
+
+    exit(path: NodePath<Program>, state: PluginState) {
+      if (state.autoImportNeeded) {
+        addAutoImport(path, t, state);
+      }
+    },
+  };
+}

package/tests/attribute.test.ts

@@ -0,0 +1,136 @@
+import pluginTester from 'babel-plugin-tester';
+import plugin from '../src';
+
+const babelOptions = {
+  plugins: ['@babel/plugin-syntax-jsx'],
+  configFile: false,
+  babelrc: false,
+};
+
+pluginTester({
+  plugin,
+  pluginOptions: {},
+  babelOptions,
+  title: 'JSX attribute transforms',
+  tests: {
+    'wraps element with single attribute .get()': {
+      code: `
+        function App() {
+          return <Component value={obs$.get()} />;
+        }
+      `,
+      output: `
+        import { Memo } from "@legendapp/state/react";
+        function App() {
+          return <Memo>{() => <Component value={obs$.get()} />}</Memo>;
+        }
+      `,
+    },
+
+    'wraps self-closing element with attribute .get()': {
+      code: `
+        function App() {
+          return <Input disabled={isDisabled$.get()} />;
+        }
+      `,
+      output: `
+        import { Memo } from "@legendapp/state/react";
+        function App() {
+          return <Memo>{() => <Input disabled={isDisabled$.get()} />}</Memo>;
+        }
+      `,
+    },
+
+    'wraps element with multiple attribute .get() into single Auto': {
+      code: `
+        function App() {
+          return <Component value={obs$.get()} label={name$.get()} />;
+        }
+      `,
+      output: `
+        import { Memo } from "@legendapp/state/react";
+        function App() {
+          return (
+            <Memo>{() => <Component value={obs$.get()} label={name$.get()} />}</Memo>
+          );
+        }
+      `,
+    },
+
+    'wraps entire element when both attribute and children have .get()': {
+      code: `
+        function App() {
+          return <div className={theme$.get()}>{count$.get()}</div>;
+        }
+      `,
+      output: `
+        import { Memo } from "@legendapp/state/react";
+        function App() {
+          return (
+            <Memo>{() => <div className={theme$.get()}>{count$.get()}</div>}</Memo>
+          );
+        }
+      `,
+    },
+
+    'wraps element with spread attribute containing .get()': {
+      code: `
+        function App() {
+          return <Component {...obs$.get()} />;
+        }
+      `,
+      output: `
+        import { Memo } from "@legendapp/state/react";
+        function App() {
+          return <Memo>{() => <Component {...obs$.get()} />}</Memo>;
+        }
+      `,
+    },
+
+    'does NOT wrap element with only key={obs$.get()}': {
+      code: `
+        function App() {
+          return <li key={item$.id.get()}>content</li>;
+        }
+      `,
+    },
+
+    'does NOT wrap element with only ref={obs$.get()}': {
+      code: `
+        function App() {
+          return <div ref={domRef$.get()}>content</div>;
+        }
+      `,
+    },
+
+    'wraps element when key + other attribute both have .get()': {
+      code: `
+        function App() {
+          return <li key={item$.id.get()} className={theme$.get()}>content</li>;
+        }
+      `,
+      output: `
+        import { Memo } from "@legendapp/state/react";
+        function App() {
+          return (
+            <Memo>
+              {() => (
+                <li key={item$.id.get()} className={theme$.get()}>
+                  content
+                </li>
+              )}
+            </Memo>
+          );
+        }
+      `,
+    },
+
+    'does NOT wrap element when attribute has .get() with args (Map.get)': {
+      code: `
+        function App() {
+          return <Component value={map.get("key")} />;
+        }
+      `,
+    },
+  },
+});