The JIT cache no longer crashes a run that compiles two or more modules (#549)

The JIT cache compiles each module in an MLIRContext of its own, but the
built-in types (i32, f32, Opaque, ...) were kept in a function-local static
map, made once in the first module's context. Every later module in the run
got types from a context already destroyed: 0xC0000005 in the compiler. The
default library's declarations name such types, so every program importing
a .ts module crashed (native_ts_graphics finding 7). The map is now a member
of MLIRGenImpl, made on first use per module.

Regression test test-jit-cache-embedded-types runs from an empty cache
every time (jit-cache-empty.cmake); a test-runner test keeps its cache
beside the source and compiles nothing from its second run on.

Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

Author: ASDAlexander77

docs/bugs/native-ts-graphics-findings.md

@@ -29,7 +29,7 @@ keeps failing by design. The repros stay here as the record of what was reported
 | 4 | Escapes not processed in templates with substitutions | Wrong value, silent | `ntg_template_literal_escape.ts` | [#539](https://github.com/ASDAlexander77/TypeScriptCompiler/pull/539) |
 | 5 | `Math.max` of negative numbers returns 2.2e-308 | Wrong value, silent (default lib) | `ntg_math_max_negative.ts` | [DefaultLib#31](https://github.com/ASDAlexander77/TypeScriptCompilerDefaultLib/pull/31), and [#535](https://github.com/ASDAlexander77/TypeScriptCompiler/pull/535) for why its test passed |
 | 6 | `boolean && (Class \| null)` as a condition fails to compile | Compile error | `ntg_and_nullable_class.ts` | [#540](https://github.com/ASDAlexander77/TypeScriptCompiler/pull/540), and [#541](https://github.com/ASDAlexander77/TypeScriptCompiler/pull/541) for its misplaced first error |
-| 7 | The JIT cache crashes a run that compiles two or more modules (pre-alpha93) | Compiler crash, no message | `ntg_jit_cache_import.ts` | open |
+| 7 | The JIT cache crashes a run that compiles two or more modules (pre-alpha93) | Compiler crash, no message | `ntg_jit_cache_import.ts` | [#549](https://github.com/ASDAlexander77/TypeScriptCompiler/pull/549) |
 
 The silent ones (2 to 5) cost the most time: the programs compile and run, and only the output
 is wrong.
@@ -226,6 +226,14 @@ per run) every run works.
 imports `core/donut.ts`, so every example crashed under the JIT twice before running.
 `run_jit.bat` passes `--jit-cache=false`.
 
+**Fixed by #549:** the cache compiles each module in an MLIRContext of its own, and the built-in
+types (`i32`, `f32`, `Opaque`, ...) were kept in a `static` map, made once, in the first module's
+context. Every later module in the run got types from a context destroyed by then. The map is
+now per module. The default library's declarations name such types, which is why every program
+importing a module crashed; without it (the suite runs with `--no-default-lib`) a module only
+crashed if it named one itself. Regression test: `test-jit-cache-embedded-types`, which runs from
+an empty cache every time.
+
 ---
 
 ## Seen earlier but not reproduced in minimal form

tslang/lib/TypeScript/MLIRGenImpl.h

@@ -14055,6 +14055,11 @@ class MLIRGenImpl
     // read of it loads the field again and casts it, so no value outlives the region it is in
     llvm::ScopedHashTable<SafeTypeKeyType, mlir::Type> safeTypesMap;
 
+    // the built-in types by name (getEmbeddedType), made on first use; not static: they belong to
+    // this module's MLIRContext, and one run can compile several modules, each in a context of its
+    // own (the JIT cache) - a type kept from an earlier one points into freed memory
+    llvm::StringMap<mlir::Type> embeddedTypeMap;
+
     // helper to get line number
     Parser parser;
     ts::SourceFile sourceFile;

tslang/lib/TypeScript/MLIRGenTypes.cpp

@@ -1025,7 +1025,7 @@ namespace mlirgen
 
     mlir::Type MLIRGenImpl::getEmbeddedTypeBuiltins(mlir::StringRef name)
     {
-        static llvm::StringMap<mlir::Type> embeddedTypes {
+        if (embeddedTypeMap.empty()) embeddedTypeMap = {
             {"TemplateStringsArray", getArrayType(getStringType()) },
             {"const",getConstType() },
 #ifdef ENABLE_JS_BUILTIN_TYPES
@@ -1082,13 +1082,12 @@ namespace mlirgen
             {"Opaque", getOpaqueType()},
         };
 
-        auto type = embeddedTypes[name];
-        return type;
+        return embeddedTypeMap.lookup(name);
     }
 
     mlir::Type MLIRGenImpl::getEmbeddedTypeNoBuiltins(mlir::StringRef name)
     {
-        static llvm::StringMap<mlir::Type> embeddedTypes {
+        if (embeddedTypeMap.empty()) embeddedTypeMap = {
             {"TemplateStringsArray", getArrayType(getStringType()) },
             {"const",getConstType() },
 #ifdef ENABLE_JS_BUILTIN_TYPES
@@ -1146,8 +1145,7 @@ namespace mlirgen
             {"Opaque", getOpaqueType()},
         };
 
-        auto type = embeddedTypes[name];
-        return type;
+        return embeddedTypeMap.lookup(name);
     }    
 
     mlir::Type MLIRGenImpl::getEmbeddedTypeWithParam(mlir::StringRef name, NodeArray<TypeNode> &typeArguments,

tslang/test/tester/CMakeLists.txt

@@ -3252,6 +3252,15 @@ add_test(NAME test-jit-cache-reuse-and-invalidate
                  "-DSOURCE_DIR=${jit_cache_dir}"
                  "-DWORK_DIR=${CMAKE_CURRENT_BINARY_DIR}/jit-cache-reuse"
                  -P "${CMAKE_CURRENT_SOURCE_DIR}/jit-cache.cmake")
+# Two modules compiled in one run, both naming built-in types (i32, f32): the second one used to
+# get the types made in the first one's MLIRContext, destroyed by then, and crash the compiler.
+add_test(NAME test-jit-cache-embedded-types
+         COMMAND ${CMAKE_COMMAND}
+                 "-DTSLANG=$<TARGET_FILE:tslang>"
+                 "-DSOURCE=${jit_cache_dir}/embedded_types.ts"
+                 "-DWORK_DIR=${CMAKE_CURRENT_BINARY_DIR}/jit-cache-embedded-types"
+                 "-DEXPECTED=done."
+                 -P "${CMAKE_CURRENT_SOURCE_DIR}/jit-cache-empty.cmake")
 
 # `import { a as b }` and `import * as M` - each used to be "can't resolve name".
 set(import_bindings_dir "${PROJECT_SOURCE_DIR}/test/tester/import-bindings")

tslang/test/tester/jit-cache-empty.cmake

@@ -0,0 +1,33 @@
+# Runs a program under the JIT cache from an empty cache, so the run compiles the program and every
+# module it imports - several modules in one process, each in an MLIRContext of its own. (A test
+# run by test-runner keeps its cache in __jit beside the source: from the second run on, it
+# compiles nothing.)
+#
+# TSLANG, SOURCE (the program), WORK_DIR (the cache, emptied first), EXPECTED (its output).
+# -mm=none and --no-default-lib keep the run to the compiler alone.
+
+cmake_minimum_required(VERSION 3.17.3)
+
+foreach(var TSLANG SOURCE WORK_DIR EXPECTED)
+    if(NOT DEFINED ${var})
+        message(FATAL_ERROR "${var} is required")
+    endif()
+endforeach()
+
+file(REMOVE_RECURSE "${WORK_DIR}")
+file(MAKE_DIRECTORY "${WORK_DIR}")
+
+execute_process(COMMAND "${TSLANG}" --emit=jit -mm=none --no-default-lib "--jit-cache-dir=${WORK_DIR}" "${SOURCE}"
+    OUTPUT_VARIABLE out
+    ERROR_VARIABLE err
+    RESULT_VARIABLE status)
+string(STRIP "${out}" out)
+if(NOT status EQUAL 0 OR NOT out STREQUAL EXPECTED)
+    message(FATAL_ERROR "expected '${EXPECTED}': exit ${status}\n${out}\n${err}")
+endif()
+
+file(GLOB objects "${WORK_DIR}/*.o" "${WORK_DIR}/*.obj")
+list(LENGTH objects count)
+if(count LESS 2)
+    message(FATAL_ERROR "expected an object for the program and one for each module it imports, found: ${objects}")
+endif()

tslang/test/tester/jit-cache/embedded_types.ts

@@ -0,0 +1,13 @@
+// Each module the JIT cache compiles has an MLIRContext of its own. The built-in types (i32,
+// f32, Opaque, ...) were kept in a static map, built in the first module's context: the next
+// module got types from a context that no longer existed, and the compiler crashed (0xC0000005).
+import { widen, half } from './embedded_types_module'
+
+const a: i32 = 21;
+const b: i64 = widen(a) * 2;
+assert(b == 42);
+
+const c: f32 = 5;
+assert(half(c) == 2.5);
+
+print("done.");

tslang/test/tester/jit-cache/embedded_types_module.ts

@@ -0,0 +1,8 @@
+// Imported by embedded_types.ts (see there).
+export function widen(x: i32): i64 {
+    return x;
+}
+
+export function half(x: f32): f32 {
+    return x / 2;
+}