Fix puts/cf.assert signature conflict crashing JIT tests

tslang/include/TypeScript/gcwrapper.h

@@ -19,6 +19,10 @@ void _mlir__GC_free(void *ptr);
 
 size_t _mlir__GC_get_heap_size();
 
+void _mlir__GC_add_roots(void *low, void *highPlus1);
+
+void _mlir__GC_remove_roots(void *low, void *highPlus1);
+
 void _mlir__GC_win32_free_heap();
 
 void *_mlir__GC_malloc_explicitly_typed(size_t size, int64_t descr);

tslang/lib/TypeScript/LowerToLLVM.cpp

@@ -158,8 +158,10 @@ class PrintOpLowering : public TsLlvmPattern<mlir_ts::PrintOp>
         auto i8PtrType = th.getPtrType();
         auto ptrType = th.getPtrType();
 
-        // Get a symbol reference to the printf function, inserting it if necessary.
-        auto putsFuncOp = ch.getOrInsertFunction("puts", th.getFunctionType(rewriter.getI32Type(), ptrType, false));
+        // Get a symbol reference to the puts function, inserting it if necessary.
+        // Note: declared as returning void (not the libc i32) to match the signature MLIR's
+        // cf.assert-to-LLVM lowering expects when it reserves/looks up "puts" for its own use.
+        auto putsFuncOp = ch.getOrInsertFunction("puts", th.getFunctionType(th.getVoidType(), ptrType, false));
 
         auto strType = mlir_ts::StringType::get(rewriter.getContext());
 
@@ -518,6 +520,32 @@ class SetStringLengthOpLowering : public TsLlvmPattern<mlir_ts::SetStringLengthO
     }
 };
 
+// LLVM's CoroSplit cannot spill variable-size allocas into the coroutine frame
+// ("Coroutines cannot handle non static allocas yet"), so buffers inside a
+// not-yet-split coroutine (async_execute_fn marked by the MLIR async lowering)
+// must be heap-allocated instead.
+static bool isInsidePresplitCoroutine(mlir::Operation *op)
+{
+    for (auto *parent = op->getParentOp(); parent; parent = parent->getParentOp())
+    {
+        if (auto passthrough = parent->getAttrOfType<mlir::ArrayAttr>("passthrough"))
+        {
+            for (auto attr : passthrough)
+            {
+                if (auto strAttr = dyn_cast<mlir::StringAttr>(attr))
+                {
+                    if (strAttr.getValue() == "presplitcoroutine")
+                    {
+                        return true;
+                    }
+                }
+            }
+        }
+    }
+
+    return false;
+}
+
 class StringConcatOpLowering : public TsLlvmPattern<mlir_ts::StringConcatOp>
 {
   public:
@@ -552,7 +580,8 @@ class StringConcatOpLowering : public TsLlvmPattern<mlir_ts::StringConcatOp>
             size = rewriter.create<LLVM::AddOp>(loc, llvmIndexType, ValueRange{size, size1.getResult()});
         }
 
-        auto allocInStack = op.getAllocInStack().has_value() && op.getAllocInStack().value();
+        auto allocInStack = op.getAllocInStack().has_value() && op.getAllocInStack().value()
+                            && !isInsidePresplitCoroutine(op);
 
         mlir::Value newStringValue = allocInStack ? ch.Alloca(i8PtrTy, size, true)
                                                   : ch.MemoryAlloc(size);

tslang/lib/TypeScriptRuntime/TypeScriptRuntime.def

@@ -15,6 +15,8 @@ EXPORTS
     GC_realloc=_mlir__GC_realloc
     GC_free=_mlir__GC_free
     GC_get_heap_size=_mlir__GC_get_heap_size
+    GC_add_roots=_mlir__GC_add_roots
+    GC_remove_roots=_mlir__GC_remove_roots
     GC_malloc_explicitly_typed=_mlir__GC_malloc_explicitly_typed
     GC_make_descriptor=_mlir__GC_make_descriptor
 

tslang/lib/TypeScriptRuntime/gc.cpp

@@ -58,6 +58,16 @@ size_t _mlir__GC_get_heap_size()
     return GC_get_heap_size();
 }
 
+void _mlir__GC_add_roots(void *low, void *highPlus1)
+{
+    GC_add_roots(low, highPlus1);
+}
+
+void _mlir__GC_remove_roots(void *low, void *highPlus1)
+{
+    GC_remove_roots(low, highPlus1);
+}
+
 void _mlir__GC_win32_free_heap()
 {
 #ifdef WIN32

tslang/test/tester/CMakeLists.txt

@@ -623,14 +623,11 @@ add_test(NAME test-jit-00-safe-cast-field-access COMMAND test-runner -jit "${PRO
 add_test(NAME test-jit-00-safe-cast-bug COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00safe_cast_bug.ts")
 add_test(NAME test-jit-00-optional COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00optional.ts")
 add_test(NAME test-jit-01-optional COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/01optional.ts")
-if (NOT(WIN32))
-# TODO: crash
-#add_test(NAME test-jit-00-async-await COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00async_await.ts")
-# TODO: crash
-#add_test(NAME test-jit-00-for-await COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00for_await.ts")
-# TODO: in opt mode, error
-#add_test(NAME test-jit-00-for-await-yield COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00for_await_yield.ts")
-endif()
+add_test(NAME test-jit-00-async-await COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00async_await.ts")
+add_test(NAME test-jit-00-for-await COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00for_await.ts")
+# TODO: async generator / for-await-of hangs at runtime under JIT
+# error: EXEC : LLVM error : Coroutines cannot handle non static allocas yet
+add_test(NAME test-jit-00-for-await-yield COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00for_await_yield.ts")
 if (NOT(WIN32))
 add_test(NAME test-jit-00-try-catch COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00try_catch.ts")
 add_test(NAME test-jit-01-try-catch COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/01try_catch.ts")
@@ -739,7 +736,7 @@ add_test(NAME test-jit-244-arraysome COMMAND test-runner  -jit "${PROJECT_SOURCE
 add_test(NAME test-jit-Grammar_and_types COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/Grammar_and_types.ts")
 add_test(NAME test-jit-StackTest COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00stack_test.ts")
 # TODO: find out why it is crashing
-#add_test(NAME test-jit-Raytrace COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/raytrace.ts")
+add_test(NAME test-jit-Raytrace COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/raytrace.ts")
 add_test(NAME test-jit-Path COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/path.ts")
 add_test(NAME test-jit-Print-Bug-01 COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/01print-bug.ts")
 add_test(NAME test-jit-arrayLiterals COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/arrayLiterals.ts")

tslang/tslang/jit.cpp

@@ -3,7 +3,11 @@
 #include "mlir/ExecutionEngine/ExecutionEngine.h"
 #include "mlir/IR/BuiltinOps.h"
 
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ExecutionEngine/SectionMemoryManager.h"
 #include "llvm/Support/DynamicLibrary.h"
+#include "llvm/Support/Memory.h"
+#include "llvm/Support/Path.h"
 
 #include <cstdio>
 #ifdef _WIN32
@@ -92,6 +96,63 @@ int loadLibrary(mlir::SmallString<256> &libPath, llvm::StringMap<void *> &export
     return 0;
 }
 
+// JIT'd globals live in RTDyld-allocated RW data sections (plain VirtualAlloc/mmap
+// memory), which the Boehm GC does not scan: for AOT binaries the loader-mapped
+// data segment is registered as a root set automatically, but in the JIT an object
+// reachable only from a global (e.g. a static class member) is collected on the
+// first GC cycle and its memory recycled, corrupting later accesses. Register every
+// RW data section with the GC as a root range. GC_add_roots/GC_remove_roots are
+// resolved dynamically from the already-loaded TypeScriptRuntime library, so this
+// stays inert when the GC is disabled or not present.
+class GCRootsSectionMemoryMapper : public llvm::SectionMemoryManager::MemoryMapper
+{
+    // AllocationPurpose lives in the enclosing SectionMemoryManager, not in the
+    // MemoryMapper base class, so it is not found by unqualified lookup here
+    using AllocationPurpose = llvm::SectionMemoryManager::AllocationPurpose;
+    using GCRootsFn = void (*)(void *, void *);
+
+  public:
+    llvm::sys::MemoryBlock allocateMappedMemory(AllocationPurpose purpose, size_t numBytes,
+                                                const llvm::sys::MemoryBlock *const nearBlock, unsigned flags,
+                                                std::error_code &errCode) override
+    {
+        auto block = llvm::sys::Memory::allocateMappedMemory(numBytes, nearBlock, flags, errCode);
+        if (purpose == AllocationPurpose::RWData && block.base() != nullptr)
+        {
+            if (auto addRoots = reinterpret_cast<GCRootsFn>(
+                    llvm::sys::DynamicLibrary::SearchForAddressOfSymbol("GC_add_roots")))
+            {
+                addRoots(block.base(), static_cast<char *>(block.base()) + block.allocatedSize());
+                gcRootBlocks.push_back(block.base());
+            }
+        }
+
+        return block;
+    }
+
+    std::error_code protectMappedMemory(const llvm::sys::MemoryBlock &block, unsigned flags) override
+    {
+        return llvm::sys::Memory::protectMappedMemory(block, flags);
+    }
+
+    std::error_code releaseMappedMemory(llvm::sys::MemoryBlock &block) override
+    {
+        if (llvm::is_contained(gcRootBlocks, block.base()))
+        {
+            if (auto removeRoots = reinterpret_cast<GCRootsFn>(
+                    llvm::sys::DynamicLibrary::SearchForAddressOfSymbol("GC_remove_roots")))
+            {
+                removeRoots(block.base(), static_cast<char *>(block.base()) + block.allocatedSize());
+            }
+        }
+
+        return llvm::sys::Memory::releaseMappedMemory(block);
+    }
+
+  private:
+    mlir::SmallVector<void *> gcRootBlocks;
+};
+
 int runJit(int argc, char **argv, mlir::ModuleOp module, CompileOptions &compileOptions)
 {
     // to avoid false positive memory leak reports in release builds
@@ -135,8 +196,17 @@ int runJit(int argc, char **argv, mlir::ModuleOp module, CompileOptions &compile
 #define LIB_NAME "lib"
 #define LIB_EXT "so"
 #endif
+    // Only locate a GC runtime when the user has not already supplied one via
+    // --shared-libs: loading a second TypeScriptRuntime copy from a different path
+    // (e.g. TSLANG_LIB_PATH) creates two independent GC instances, so roots
+    // registered in one are invisible to collections running in the other and
+    // GC-heap pointers cross between the two heaps.
+    auto hasTypeScriptRuntime = llvm::any_of(clSharedLibs, [](const std::string &libPath) {
+        return llvm::sys::path::stem(libPath).contains_insensitive("TypeScriptRuntime");
+    });
+
     std::string pathTypeScriptLib("../lib/" LIB_NAME "TypeScriptRuntime." LIB_EXT);
-    if (!disableGC.getValue())
+    if (!disableGC.getValue() && !hasTypeScriptRuntime)
     {
         auto absPath3 = makeAbsolutePath(mergeWithDefaultLibPath(getTslangLibPath(), LIB_NAME "TypeScriptRuntime." LIB_EXT));
         if (absPath3.empty())
@@ -175,7 +245,11 @@ int runJit(int argc, char **argv, mlir::ModuleOp module, CompileOptions &compile
     mlir::SmallVector<mlir::StringRef> sharedLibPaths;
     sharedLibPaths.append(begin(clSharedLibs), end(clSharedLibs));
 
+    // must outlive the engine: sections are released from the engine's destructor
+    GCRootsSectionMemoryMapper gcRootsMemoryMapper;
+
     mlir::ExecutionEngineOptions engineOptions;
+    engineOptions.sectionMemoryMapper = &gcRootsMemoryMapper;
     engineOptions.transformer = optPipeline;
     engineOptions.enableObjectDump = dumpObjectFile;
     engineOptions.enableGDBNotificationListener = !enableOpt;
@@ -264,13 +338,29 @@ int runJit(int argc, char **argv, mlir::ModuleOp module, CompileOptions &compile
     }
 
     // The JIT program has finished. On Windows/COFF the MLIR/ORC LLJIT platform
-    // registers process-level atexit glue that faults during teardown, so the
-    // normal CRT exit path crashes after a successful run. Flush stdio and exit
-    // the process directly, bypassing the CRT atexit handlers.
+    // registers process-level atexit glue that faults during teardown (exception
+    // 0x80000003), so the normal CRT exit path crashes after a successful run.
+    // ExitProcess is not safe either: it runs DLL_PROCESS_DETACH, and
+    // TypeScriptRuntime.dll's static AsyncRuntime destructor then waits on its
+    // thread pool whose worker threads ExitProcess has already terminated,
+    // deadlocking forever when async work is still pending. TerminateProcess
+    // skips all teardown, so nothing can crash or deadlock — but it also skips
+    // the detach-time stdio flush, so flush every CRT the JIT'd code may have
+    // bound to (ucrtbase buffers its streams separately from our static CRT).
     fflush(stdout);
     fflush(stderr);
 #ifdef _WIN32
-    ExitProcess(0);
+    for (auto crtName : {"ucrtbase.dll", "ucrtbased.dll"})
+    {
+        if (HMODULE crt = GetModuleHandleA(crtName))
+        {
+            if (auto crtFflush = (int(__cdecl *)(FILE *))GetProcAddress(crt, "fflush"))
+            {
+                crtFflush(nullptr); // fflush(NULL) flushes all of that CRT's output streams
+            }
+        }
+    }
+    TerminateProcess(GetCurrentProcess(), 0);
 #endif
     return 0;
 }