Merge pull request #33 from ef4/fix-unbalanced-exit

Fix unbalanced exit

Author: ef4

Diff for: src/index.ts

@@ -1,14 +1,15 @@
 import type * as Babel from "@babel/core";
 import type { types as t, NodePath } from "@babel/core";
-import { createRequire } from "node:module";
 import { ImportUtil, type Importer } from "babel-import-util";
 import { globalId } from "./global-id.ts";
-const req = createRequire(import.meta.url);
-const { default: decoratorSyntax } = req("@babel/plugin-syntax-decorators");
+
+// @ts-expect-error no upstream types
+import decoratorSyntax from "@babel/plugin-syntax-decorators";
 
 interface State extends Babel.PluginPass {
   currentClassBodies: t.ClassBody[];
   currentObjectExpressions: {
+    node: t.ObjectExpression;
     decorated: [
       "field" | "method",
       t.Expression, // for the property name
@@ -54,8 +55,10 @@ function makeVisitor(babel: typeof Babel): Babel.Visitor<State> {
       enter(path, state) {
         state.currentClassBodies.unshift(path.node);
       },
-      exit(_path, state) {
-        state.currentClassBodies.shift();
+      exit(path, state) {
+        if (state.currentClassBodies[0] === path.node) {
+          state.currentClassBodies.shift();
+        }
       },
     },
     ClassExpression(path, state) {
@@ -241,12 +244,16 @@ function makeVisitor(babel: typeof Babel): Babel.Visitor<State> {
       }
     },
     ObjectExpression: {
-      enter(_path, state) {
+      enter(path, state) {
         state.currentObjectExpressions.unshift({
+          node: path.node,
           decorated: [],
         });
       },
       exit(path, state) {
+        if (state.currentObjectExpressions[0]?.node !== path.node) {
+          return;
+        }
         let { decorated } = state.currentObjectExpressions.shift()!;
         if (decorated.length > 0) {
           state.util.replaceWith(path, (i) =>

Diff for: tests/plugin-interop-test.ts

@@ -1,23 +1,22 @@
 import { module, test } from "qunit";
-import { builder, Builder } from "./helpers.ts";
+import { builder } from "./helpers.ts";
 import { type LegacyClassDecorator } from "../src/runtime.ts";
 import * as runtimeImpl from "../src/runtime.ts";
 import ourDecorators from "../src/index.ts";
 import { createRequire } from "node:module";
 import { mkdtemp, writeFile } from "node:fs/promises";
 import { join } from "node:path";
 import { tmpdir } from "node:os";
+import type * as Babel from "@babel/core";
 
 const require = createRequire(import.meta.url);
 const Colocation = require("@embroider/shared-internals/src/template-colocation-plugin");
 
-module("plugin-interop", (hooks) => {
-  let build: Builder;
-
-  hooks.before(async () => {
+module("plugin-interop", () => {
+  test("colocation", async (assert) => {
     let dir = await mkdtemp(join(tmpdir(), "decorator-transforms-"));
     await writeFile(join(dir, "example.hbs"), "");
-    build = builder(
+    let build = builder(
       [],
       [
         [
@@ -29,9 +28,6 @@ module("plugin-interop", (hooks) => {
       [],
       join(dir, "example.js")
     );
-  });
-
-  test("colocation", async (assert) => {
     let red: LegacyClassDecorator = (target) => {
       return class extends target {
         get red() {
@@ -77,4 +73,53 @@ module("plugin-interop", (hooks) => {
     assert.strictEqual(new Example().red, "#ff0000");
     assert.strictEqual(setComponents, 1);
   });
+
+  // This models a behavior that @embroider/macros uses when trying to optimize
+  // `getConfig().with.a.path`, since it uses parentPath to go higher and
+  // replace a larger part of the expression.
+  test("parentPath replace", async (assert) => {
+    function inserter(babel: typeof Babel): Babel.PluginObj {
+      let t = babel.types;
+      return {
+        visitor: {
+          CallExpression: {
+            enter(path) {
+              let callee = path.node.callee;
+              if (callee.type === "Identifier" && callee.name === "expandMe") {
+                path.parentPath.replaceWith(
+                  t.objectExpression([
+                    t.objectProperty(
+                      t.identifier("value"),
+                      t.stringLiteral("it worked")
+                    ),
+                  ])
+                );
+              }
+            },
+          },
+        },
+      };
+    }
+
+    let build = builder(
+      [],
+      [
+        [
+          ourDecorators,
+          { runtime: { import: "decorator-transforms/runtime" } },
+        ],
+        [inserter],
+      ],
+      []
+    );
+
+    let { default: Example } = await build.module(
+      `
+        export default expandMe().a;
+      `,
+      {}
+    );
+
+    assert.strictEqual(Example.value, "it worked");
+  });
 });