From 86cdb09bf209611effa715933a4b8842a6278c1a Mon Sep 17 00:00:00 2001
From: Sam Gammon <sam@elide.ventures>
Date: Thu, 13 Jun 2024 03:24:57 -0700
Subject: [PATCH] fix(typescript): fixup ts module loader

- fix: truffle boundary for native ts module loader
- chore: general typescript cleanup

Fixes and closes elide-dev/elide#824

Signed-off-by: Sam Gammon <sam@elide.ventures>
---
 .../lang/typescript/JSRealmPatcher.java       |  2 --
 .../lang/typescript/TypeScriptLanguage.java   | 32 +++++++++++++------
 2 files changed, 22 insertions(+), 12 deletions(-)

diff --git a/packages/graalvm-ts/src/main/java/elide/runtime/lang/typescript/JSRealmPatcher.java b/packages/graalvm-ts/src/main/java/elide/runtime/lang/typescript/JSRealmPatcher.java
index 58852ba419..77c7d9fa79 100644
--- a/packages/graalvm-ts/src/main/java/elide/runtime/lang/typescript/JSRealmPatcher.java
+++ b/packages/graalvm-ts/src/main/java/elide/runtime/lang/typescript/JSRealmPatcher.java
@@ -12,12 +12,10 @@
  */
 package elide.runtime.lang.typescript;
 
-import com.oracle.truffle.api.CompilerDirectives;
 import com.oracle.truffle.js.runtime.JSRealm;
 import java.lang.reflect.Field;
 
 class JSRealmPatcher {
-  @CompilerDirectives.TruffleBoundary
   public static void setTSModuleLoader(JSRealm jsRealm, TypeScriptModuleLoader newModuleLoader) {
     try {
       Field moduleLoaderField = JSRealm.class.getDeclaredField("moduleLoader");
diff --git a/packages/graalvm-ts/src/main/java/elide/runtime/lang/typescript/TypeScriptLanguage.java b/packages/graalvm-ts/src/main/java/elide/runtime/lang/typescript/TypeScriptLanguage.java
index 81755d7990..636df9d8d7 100644
--- a/packages/graalvm-ts/src/main/java/elide/runtime/lang/typescript/TypeScriptLanguage.java
+++ b/packages/graalvm-ts/src/main/java/elide/runtime/lang/typescript/TypeScriptLanguage.java
@@ -23,6 +23,7 @@
 import com.oracle.truffle.js.lang.JavaScriptLanguage;
 import com.oracle.truffle.js.runtime.JSEngine;
 import com.oracle.truffle.js.runtime.JSRealm;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.List;
 import org.graalvm.polyglot.SandboxPolicy;
 
@@ -61,19 +62,26 @@ public class TypeScriptLanguage extends TruffleLanguage<JSRealm> {
   public static final String IMPLEMENTATION_NAME = "TypeScript";
   public static final String ID = "ts";
   public static final String TYPESCRIPT_VERSION = "5.4.5";
-  private TypeScriptCompiler tsCompiler;
+
+  @SuppressWarnings("java:S3077")
+  @CompilerDirectives.CompilationFinal private volatile TypeScriptCompiler tsCompiler = null;
+  private final AtomicBoolean compilerInitialized = new AtomicBoolean(false);
   private Env env;
 
   @Override
-  protected JSRealm createContext(Env env) {
+  protected JSRealm createContext(Env currentEnv) {
     CompilerAsserts.neverPartOfCompilation();
     var js = JavaScriptLanguage.getCurrentLanguage();
-    LanguageInfo jsInfo = env.getInternalLanguages().get("js");
-    env.initializeLanguage(jsInfo);
+    LanguageInfo jsInfo = currentEnv.getInternalLanguages().get("js");
+    currentEnv.initializeLanguage(jsInfo);
     var jsEnv = JavaScriptLanguage.getCurrentEnv();
+
+    if (!compilerInitialized.get()) {
+      compilerInitialized.compareAndSet(false, true);
+      tsCompiler = new TypeScriptCompiler(jsEnv);
+      env = jsEnv;
+    }
     var ctx = JSEngine.createJSContext(js, jsEnv);
-    tsCompiler = new TypeScriptCompiler(jsEnv);
-    this.env = jsEnv;
     var realm = ctx.createRealm(jsEnv);
     JSRealmPatcher.setTSModuleLoader(realm, new TypeScriptModuleLoader(realm, tsCompiler));
     return realm;
@@ -98,7 +106,7 @@ protected CallTarget parse(ParsingRequest parsingRequest) {
     return wrapper.getCallTarget();
   }
 
-  private static class TSRootNode extends RootNode {
+  private class TSRootNode extends RootNode {
     private final RootNode delegate;
 
     protected TSRootNode(TruffleLanguage<?> language, RootNode delegate) {
@@ -106,11 +114,15 @@ protected TSRootNode(TruffleLanguage<?> language, RootNode delegate) {
       this.delegate = delegate;
     }
 
+    @TruffleBoundary
+    private void setModuleLoader() {
+      JSRealm realm = JSRealm.get(delegate);
+      JSRealmPatcher.setTSModuleLoader(realm, new TypeScriptModuleLoader(realm, tsCompiler));
+    }
+
     @Override
     public Object execute(VirtualFrame frame) {
-      // @TODO(sgammon): causes restricted types to be reached
-      // JSRealm realm = JSRealm.get(delegate);
-      // JSRealmPatcher.setTSModuleLoader(realm, new TypeScriptModuleLoader(realm, tsCompiler));
+      setModuleLoader();
       return delegate.execute(frame);
     }
   }