diff --git a/src/core/compile.ts b/src/core/compile.ts
index 1af1680..7a9c878 100644
--- a/src/core/compile.ts
+++ b/src/core/compile.ts
@@ -25,7 +25,7 @@ export async function compile(testCodePaths: string[], outputFolder: string, com
       const argv = compileFlags.split(" ");
       ascArgv = ascArgv.concat(argv);
     }
-    await ascMain(ascArgv);
+    await ascMain(ascArgv, false);
   };
 
   // Here, for-await is more efficient and less memory cost than Promise.all()
diff --git a/src/core/precompile.ts b/src/core/precompile.ts
index 23a3277..2267c20 100644
--- a/src/core/precompile.ts
+++ b/src/core/precompile.ts
@@ -91,7 +91,7 @@ async function transform(transformFunction: string, codePath: string, flags: str
     const argv = flags.split(" ");
     ascArgv = ascArgv.concat(argv);
   }
-  await ascMain(ascArgv);
+  await ascMain(ascArgv, true);
   collectCallback();
 }
 
diff --git a/src/utils/ascWrapper.ts b/src/utils/ascWrapper.ts
index 6936027..1de2de4 100644
--- a/src/utils/ascWrapper.ts
+++ b/src/utils/ascWrapper.ts
@@ -11,10 +11,13 @@ export class CompilationError extends Error {
   }
 }
 
-export async function ascMain(ascArgv: string[]) {
+export async function ascMain(ascArgv: string[], isTransform: boolean) {
   const stderr = createMemoryStream();
   const { error } = await compiler.compile(ascArgv.concat("--noColors"), { stderr });
   if (error) {
+    if (isTransform && error instanceof Error && error.message === "TransformDone") {
+      return;
+    }
     throw new CompilationError(stderr.toString());
   }
 }
diff --git a/transform/listFunctions.mts b/transform/listFunctions.mts
index 40a7d69..19c804d 100644
--- a/transform/listFunctions.mts
+++ b/transform/listFunctions.mts
@@ -68,6 +68,7 @@ class SourceFunctionTransform extends Transform {
     this.visitNode(entrySource);
     this.functionInfos.reverse();
     globalThis.functionInfos = this.functionInfos;
+    throw new Error("TransformDone");
   }
 
   visitNode(node: Node) {
diff --git a/transform/listTestNames.mts b/transform/listTestNames.mts
index 6dd1fed..f89a3d0 100644
--- a/transform/listTestNames.mts
+++ b/transform/listTestNames.mts
@@ -42,11 +42,11 @@ import {
   NamespaceDeclaration,
   VariableDeclaration,
   SwitchCase,
-  Program,
   IdentifierExpression,
   LiteralExpression,
   LiteralKind,
   StringLiteralExpression,
+  Parser,
 } from "assemblyscript";
 import assert from "node:assert";
 
@@ -55,14 +55,15 @@ class SourceFunctionTransform extends Transform {
   currentTestDescriptions: string[] = [];
   testFileName: string;
 
-  afterInitialize(program: Program) {
+  afterParse(parser: Parser) {
     // There will be two sources with SourceKind.UserEntry, ~lib/rt/index-incremental.ts should be filtered
-    const entrySource = program.sources.find(
+    const entrySource = parser.sources.find(
       (source) => source.sourceKind === SourceKind.UserEntry && !source.normalizedPath.startsWith("~lib/")
     );
     this.testFileName = entrySource.normalizedPath;
     this.visitNode(entrySource);
     globalThis.testNames = this.testNames;
+    throw new Error("TransformDone");
   }
 
   visitNode(node: Node) {