fix(transform): prevent running actions against different instances of handlers

packages/esbuild/src/index.ts

@@ -21,6 +21,7 @@ import {
   TransformCacheCollection,
   createFileReporter,
 } from '@wyw-in-js/transform';
+import { asyncResolverFactory } from '@wyw-in-js/shared';
 
 type EsbuildPluginOptions = {
   debug?: IFileReporterOptions | false | null | undefined;
@@ -42,32 +43,40 @@ export default function wywInJS({
 }: EsbuildPluginOptions = {}): Plugin {
   let options = esbuildOptions;
   const cache = new TransformCacheCollection();
+
+  const createAsyncResolver = asyncResolverFactory(
+    async (
+      resolved: {
+        errors: unknown[];
+        path: string;
+      },
+      token: string
+    ): Promise<string> => {
+      if (resolved.errors.length > 0) {
+        throw new Error(`Cannot resolve ${token}`);
+      }
+
+      return resolved.path.replace(/\\/g, posix.sep);
+    },
+    (what, importer) => [
+      what,
+      {
+        resolveDir: isAbsolute(importer)
+          ? dirname(importer)
+          : join(process.cwd(), dirname(importer)),
+        kind: 'import-statement',
+      },
+    ]
+  );
+
   return {
     name: 'wyw-in-js',
     setup(build) {
       const cssLookup = new Map<string, string>();
 
       const { emitter, onDone } = createFileReporter(debug ?? false);
 
-      const asyncResolve = async (
-        token: string,
-        importer: string
-      ): Promise<string> => {
-        const context = isAbsolute(importer)
-          ? dirname(importer)
-          : join(process.cwd(), dirname(importer));
-
-        const result = await build.resolve(token, {
-          resolveDir: context,
-          kind: 'import-statement',
-        });
-
-        if (result.errors.length > 0) {
-          throw new Error(`Cannot resolve ${token}`);
-        }
-
-        return result.path.replace(/\\/g, posix.sep);
-      };
+      const asyncResolve = createAsyncResolver(build.resolve);
 
       build.onEnd(() => {
         onDone(process.cwd());

packages/rollup/src/index.ts

@@ -5,9 +5,14 @@
  */
 
 import { createFilter } from '@rollup/pluginutils';
-import type { Plugin } from 'rollup';
+import type { Plugin, ResolvedId } from 'rollup';
 
-import { logger, slugify, syncResolve } from '@wyw-in-js/shared';
+import {
+  asyncResolverFactory,
+  logger,
+  slugify,
+  syncResolve,
+} from '@wyw-in-js/shared';
 import type { PluginOptions, Preprocessor, Result } from '@wyw-in-js/transform';
 import {
   getFileIdx,
@@ -34,6 +39,32 @@ export default function wywInJS({
   const cache = new TransformCacheCollection();
   const emptyConfig = {};
 
+  const createAsyncResolver = asyncResolverFactory(
+    async (resolved: ResolvedId | null, what, importer, stack) => {
+      if (resolved) {
+        if (resolved.external) {
+          // If module is marked as external, Rollup will not resolve it,
+          // so we need to resolve it ourselves with default resolver
+          return syncResolve(what, importer, stack);
+        }
+
+        // Vite adds param like `?v=667939b3` to cached modules
+        const resolvedId = resolved.id.split('?')[0];
+
+        if (resolvedId.startsWith('\0')) {
+          // \0 is a special character in Rollup that tells Rollup to not include this in the bundle
+          // https://rollupjs.org/guide/en/#outputexports
+          return null;
+        }
+
+        return resolvedId;
+      }
+
+      throw new Error(`Could not resolve ${what}`);
+    },
+    (what, importer) => [what, importer]
+  );
+
   const plugin: Plugin = {
     name: 'wyw-in-js',
     load(id: string) {
@@ -54,39 +85,6 @@ export default function wywInJS({
 
       log('init %s', id);
 
-      const asyncResolve = async (
-        what: string,
-        importer: string,
-        stack: string[]
-      ) => {
-        const resolved = await this.resolve(what, importer);
-        if (resolved) {
-          if (resolved.external) {
-            // If module is marked as external, Rollup will not resolve it,
-            // so we need to resolve it ourselves with default resolver
-            const resolvedId = syncResolve(what, importer, stack);
-            log("resolve: ✅ '%s'@'%s -> %O\n%s", what, importer, resolved);
-            return resolvedId;
-          }
-
-          log("resolve: ✅ '%s'@'%s -> %O\n%s", what, importer, resolved);
-
-          // Vite adds param like `?v=667939b3` to cached modules
-          const resolvedId = resolved.id.split('?')[0];
-
-          if (resolvedId.startsWith('\0')) {
-            // \0 is a special character in Rollup that tells Rollup to not include this in the bundle
-            // https://rollupjs.org/guide/en/#outputexports
-            return null;
-          }
-
-          return resolvedId;
-        }
-
-        log("resolve: ❌ '%s'@'%s", what, importer);
-        throw new Error(`Could not resolve ${what}`);
-      };
-
       const transformServices = {
         options: {
           filename: id,
@@ -100,7 +98,7 @@ export default function wywInJS({
       const result = await transform(
         transformServices,
         code,
-        asyncResolve,
+        createAsyncResolver(this.resolve),
         emptyConfig
       );
 

packages/shared/src/asyncResolverFactory.ts

@@ -0,0 +1,35 @@
+export const asyncResolverFactory = <
+  TResolved,
+  const TResolverArgs extends readonly unknown[],
+  TResolve extends (...args: TResolverArgs) => Promise<TResolved>,
+>(
+  onResolve: (
+    resolved: TResolved,
+    what: string,
+    importer: string,
+    stack: string[]
+  ) => Promise<string | null>,
+  mapper: (what: string, importer: string, stack: string[]) => TResolverArgs
+) => {
+  const memoizedSyncResolve = new WeakMap<
+    TResolve,
+    (what: string, importer: string, stack: string[]) => Promise<string | null>
+  >();
+
+  return (resolveFn: TResolve) => {
+    if (!memoizedSyncResolve.has(resolveFn)) {
+      const fn = (
+        what: string,
+        importer: string,
+        stack: string[]
+      ): Promise<string | null> =>
+        resolveFn(...mapper(what, importer, stack)).then((resolved) =>
+          onResolve(resolved, what, importer, stack)
+        );
+
+      memoizedSyncResolve.set(resolveFn, fn);
+    }
+
+    return memoizedSyncResolve.get(resolveFn)!;
+  };
+};

packages/shared/src/index.ts

@@ -1,6 +1,7 @@
 export { Debugger } from 'debug';
 
 export { asyncResolveFallback, syncResolve } from './asyncResolveFallback';
+export { asyncResolverFactory } from './asyncResolverFactory';
 export { hasEvalMeta } from './hasEvalMeta';
 export { findPackageJSON } from './findPackageJSON';
 export { isBoxedPrimitive } from './isBoxedPrimitive';

packages/transform/src/transform.ts

@@ -25,6 +25,7 @@ import {
 } from './transform/generators/resolveImports';
 import { withDefaultServices } from './transform/helpers/withDefaultServices';
 import type {
+  Handler,
   Handlers,
   IResolveImportsAction,
   Services,
@@ -39,6 +40,16 @@ type PartialServices = Partial<Omit<Services, 'options'>> & {
 
 type AllHandlers<TMode extends 'async' | 'sync'> = Handlers<TMode>;
 
+const memoizedSyncResolve = new WeakMap<
+  (what: string, importer: string, stack: string[]) => string | null,
+  Handler<'sync', IResolveImportsAction>
+>();
+
+const memoizedAsyncResolve = new WeakMap<
+  (what: string, importer: string, stack: string[]) => Promise<string | null>,
+  Handler<'async' | 'sync', IResolveImportsAction>
+>();
+
 export function transformSync(
   partialServices: PartialServices,
   originalCode: string,
@@ -78,13 +89,20 @@ export function transformSync(
 
   const workflowAction = entrypoint.createAction('workflow', undefined);
 
+  if (!memoizedSyncResolve.has(syncResolve)) {
+    memoizedSyncResolve.set(
+      syncResolve,
+      function resolveImports(this: IResolveImportsAction) {
+        return syncResolveImports.call(this, syncResolve);
+      }
+    );
+  }
+
   try {
     const result = syncActionRunner(workflowAction, {
       ...baseHandlers,
       ...customHandlers,
-      resolveImports() {
-        return syncResolveImports.call(this, syncResolve);
-      },
+      resolveImports: memoizedSyncResolve.get(syncResolve)!,
     });
 
     entrypoint.log('%s is ready', entrypoint.name);
@@ -160,13 +178,21 @@ export async function transform(
 
   const workflowAction = entrypoint.createAction('workflow', undefined);
 
+  if (!memoizedAsyncResolve.has(asyncResolve)) {
+    const resolveImports = function resolveImports(
+      this: IResolveImportsAction
+    ) {
+      return asyncResolveImports.call(this, asyncResolve);
+    };
+
+    memoizedAsyncResolve.set(asyncResolve, resolveImports);
+  }
+
   try {
     const result = await asyncActionRunner(workflowAction, {
       ...baseHandlers,
       ...customHandlers,
-      resolveImports(this: IResolveImportsAction) {
-        return asyncResolveImports.call(this, asyncResolve);
-      },
+      resolveImports: memoizedAsyncResolve.get(asyncResolve)!,
     });
 
     entrypoint.log('%s is ready', entrypoint.name);

packages/transform/src/transform/actions/BaseAction.ts

@@ -51,6 +51,8 @@ export class BaseAction<TAction extends ActionQueueItem>
     TypeOfResult<TAction>
   >[] = [];
 
+  private handler: null | unknown = null;
+
   public constructor(
     public readonly type: TAction['type'],
     public readonly services: Services,
@@ -127,6 +129,12 @@ export class BaseAction<TAction extends ActionQueueItem>
   >(handler: THandler) {
     type IterationResult = AnyIteratorResult<TMode, TypeOfResult<TAction>>;
 
+    if (this.handler && this.handler !== handler) {
+      throw new Error(`action handler is already set`);
+    }
+
+    this.handler = handler;
+
     if (!this.activeScenario) {
       this.activeScenario = handler.call(this);
       this.activeScenarioNextResults = [];

packages/transform/src/transform/actions/__tests__/BaseAction.test.ts

@@ -308,5 +308,21 @@ describe('BaseAction', () => {
         expect(e).toBe(error);
       }
     });
+
+    it('should throw an error if run against multiple handlers', () => {
+      const handler1: Handler<'sync', ITransformAction> = function* handler() {
+        return emptyResult;
+      };
+
+      const handler2: Handler<'sync', ITransformAction> = function* handler() {
+        return emptyResult;
+      };
+
+      action.run(handler1);
+
+      expect(() => action.run(handler2)).toThrowError(
+        'action handler is already set'
+      );
+    });
   });
 });