Generic callback T no longer takes in the closure's captures (intermittent tslang crash)

.gitignore

@@ -15,4 +15,4 @@ libatomic_ops-*
 *.log
 
 # the JIT cache (tslang --emit=jit)
-__jit/
+__jit*/

tslang/include/TypeScript/MLIRLogic/MLIRHelper.h

@@ -6,6 +6,7 @@
 #include "llvm/Support/ConvertUTF.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/Support/xxhash.h"
 
 #include "parser.h"
 #include "node_factory.h"
@@ -159,7 +160,10 @@ class MLIRHelper
 
                 assert(line != 0 || column != 0);
 
-                auto hashCode = hash_value(fileName);
+                // of the file's name, not of the attribute (hash_value(StringAttr) hashes its storage
+                // address): an anonymous function's name was different in every run, and a table keyed
+                // by those names visited them in a different order each time
+                auto hashCode = (uint32_t)llvm::xxh3_64bits(fileName.getValue());
                 ssName << 'L' << line << 'C' << column << "FH" << hashCode;
             })
             .Case<mlir::NameLoc>([&](auto loc) {

tslang/lib/TypeScript/MLIRGenGenerics.cpp

@@ -526,7 +526,7 @@ namespace mlirgen
 
     mlir::Type MLIRGenImpl::instantiateSpecializedFunctionTypeHelper(mlir::Location location, mlir::Value functionRefValue,
                                                         mlir::Type recieverType, bool discoverReturnType,
-                                                        const GenContext &genContext)
+                                                        const GenContext &genContext, bool *capturesFirst)
     {
         auto currValue = functionRefValue;
         if (auto createBoundFunctionOp = currValue.getDefiningOp<mlir_ts::CreateBoundFunctionOp>())
@@ -547,7 +547,7 @@ namespace mlirgen
             SourceFileScope sourceFileScope(*this, functionGenericTypeInfo->sourceFile, functionGenericTypeInfo->fileName);
 
             return instantiateSpecializedFunctionTypeHelper(location, functionGenericTypeInfo->functionDeclaration,
-                                                            recieverType, discoverReturnType, genContext);
+                                                            recieverType, discoverReturnType, genContext, capturesFirst);
         }
 
         // a plausible trigger (a generic function assigned to a variable, then called through
@@ -561,7 +561,7 @@ namespace mlirgen
 
     mlir::Type MLIRGenImpl::instantiateSpecializedFunctionTypeHelper(mlir::Location location, FunctionLikeDeclarationBase funcDecl,
                                                         mlir::Type recieverType, bool discoverReturnType,
-                                                        const GenContext &genContext)
+                                                        const GenContext &genContext, bool *capturesFirst)
     {
         GenContext funcGenContext(genContext);
         funcGenContext.receiverFuncType = recieverType;
@@ -580,6 +580,15 @@ namespace mlirgen
             return mlir::Type();
         }
 
+        // a closure's type has its captures as an input ahead of its parameters
+        // (discoverFunctionReturnTypeAndCapturedVars)
+        if (capturesFirst)
+        {
+            auto funcType = funcOp->getFuncType();
+            *capturesFirst = funcOp->getHasCapturedVars() && funcType &&
+                             funcType.getNumInputs() > funcOp->getParams().size();
+        }
+
         return funcOp->getFuncType();
     }
 
@@ -789,8 +798,9 @@ namespace mlirgen
             LLVM_DEBUG(llvm::dbgs()
                             << "\n!! param type for arrow func[" << paramIndex << "]: " << recreatedParamType << "\n";);
 
+            auto capturesFirst = false;
             auto newArrowFuncType = instantiateSpecializedFunctionTypeHelper(location, argOp, recreatedParamType,
-                                                                                true, genericTypeGenContext);
+                                                                                true, genericTypeGenContext, &capturesFirst);
 
             LLVM_DEBUG(llvm::dbgs() << "\n!! instantiate specialized arrow type function: "
                                     << newArrowFuncType << "\n";);
@@ -800,9 +810,13 @@ namespace mlirgen
                 return {mlir::failure(), true};
             }
 
-            // ...and so does the arrow's type built again here
+            // ...and so does the arrow's type built again here: bound already, or a closure whose
+            // captures were just discovered (the arrow still a bare generic reference), it has them as its
+            // first input, which paired with `T` made it `T | captures` - in an order of the members' type
+            // addresses, so the arrow's parameter was taken for its captures in some runs and not others
             if (auto funcType = dyn_cast<mlir_ts::FunctionType>(newArrowFuncType);
-                funcType && argOp.getDefiningOp<mlir_ts::CreateBoundFunctionOp>() && funcType.getNumInputs() > 0)
+                funcType && (argOp.getDefiningOp<mlir_ts::CreateBoundFunctionOp>() || capturesFirst) &&
+                funcType.getNumInputs() > 0)
             {
                 newArrowFuncType = mlir_ts::FunctionType::get(builder.getContext(), funcType.getInputs().drop_front(),
                                                               funcType.getResults(), funcType.isVarArg());

tslang/lib/TypeScript/MLIRGenImpl.h

@@ -1614,11 +1614,11 @@ class MLIRGenImpl
 
     mlir::Type instantiateSpecializedFunctionTypeHelper(mlir::Location location, mlir::Value functionRefValue,
                                                         mlir::Type recieverType, bool discoverReturnType,
-                                                        const GenContext &genContext);
+                                                        const GenContext &genContext, bool *capturesFirst = nullptr);
 
     mlir::Type instantiateSpecializedFunctionTypeHelper(mlir::Location location, FunctionLikeDeclarationBase funcDecl,
                                                         mlir::Type recieverType, bool discoverReturnType,
-                                                        const GenContext &genContext);
+                                                        const GenContext &genContext, bool *capturesFirst = nullptr);
 
     void rollbackPostponedErrorMessages(mlir::SmallVector<std::unique_ptr<mlir::Diagnostic>> *postponedMessages, size_t size);
 

tslang/test/tester/CMakeLists.txt

@@ -616,6 +616,7 @@ tslang_add_test(NAME test-compile-own-00-object-literal-contextual-lambda COMMAN
 tslang_add_test(NAME test-compile-00-interface-field-getter-owned COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00interface_field_getter_owned.ts")
 tslang_add_test(NAME test-compile-00-conditional-record-literal-temporary COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00conditional_record_literal_temporary.ts")
 tslang_add_test(NAME test-compile-00-conditional-nested-record-literal-temporary COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00conditional_nested_record_literal_temporary.ts")
+tslang_add_test(NAME test-compile-00-generic-callback-closure-late-class COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00generic_callback_closure_late_class.ts")
 tslang_add_test(NAME test-compile-00-interface-print COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00interface_print.ts")
 tslang_add_test(NAME test-compile-00-class-method-devirtualized COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00class_method_devirtualized.ts")
 tslang_add_test(NAME test-compile-00-call-arg-param-field-global COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00call_arg_param_field_global.ts")
@@ -1260,6 +1261,7 @@ tslang_add_test(NAME test-jit-own-00-object-literal-contextual-lambda COMMAND te
 tslang_add_test(NAME test-jit-00-interface-field-getter-owned COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00interface_field_getter_owned.ts")
 tslang_add_test(NAME test-jit-00-conditional-record-literal-temporary COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00conditional_record_literal_temporary.ts")
 tslang_add_test(NAME test-jit-00-conditional-nested-record-literal-temporary COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00conditional_nested_record_literal_temporary.ts")
+tslang_add_test(NAME test-jit-00-generic-callback-closure-late-class COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00generic_callback_closure_late_class.ts")
 tslang_add_test(NAME test-jit-00-interface-print COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00interface_print.ts")
 tslang_add_test(NAME test-jit-00-class-method-devirtualized COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00class_method_devirtualized.ts")
 tslang_add_test(NAME test-jit-00-call-arg-param-field-global COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00call_arg_param_field_global.ts")
@@ -2285,6 +2287,7 @@ set(TSLANG_CORPUS
     00interface_field_getter_owned.ts
     00conditional_record_literal_temporary.ts
     00conditional_nested_record_literal_temporary.ts
+    00generic_callback_closure_late_class.ts
     00interface_print.ts
     00class_method_devirtualized.ts
     00call_arg_param_field_global.ts

tslang/test/tester/tests/00generic_callback_closure_late_class.ts

@@ -0,0 +1,54 @@
+// A closure (it reads `this`) given to a generic function in a class method, the array's element a
+// class declared further down. While the method's return type was discovered, the arrow's type was
+// rebuilt with its captures as its first input, and inference paired them with `T`: `T` was
+// `Late | captures`. Here that failed every time ("cast from array type ... to array type
+// '!ts.array<!ts.union<... !ts.ref<!ts.tuple<{"this",...'"); with the default library's generic
+// `Array.filter` the union's members came in an order of their types' addresses, so tslang crashed
+// in one run in four (0xC0000005, native_ts_graphics' ray_tracing_extended_scene.ts).
+let seed = 1.0;
+function rnd() {
+    seed = (seed * 16807.0) % 2147483647.0;
+    return seed / 2147483647.0;
+}
+
+function keep<T>(items: T[], f: (v: T, i: number) => boolean): T[] {
+    const kept: T[] = [];
+    for (let i = 0; i < items.length; i++) {
+        if (f(items[i], i)) kept.push(items[i]);
+    }
+
+    return kept;
+}
+
+class Gen {
+    particles: Late[] = [];
+    lifetime = 2.0;
+
+    update() {
+        const remaining = keep(this.particles, (particle) => !(particle.duration > rnd() * this.lifetime));
+        this.particles = remaining;
+        return remaining.length;
+    }
+
+    firstDuration() {
+        return this.particles[0].duration;
+    }
+
+    longer(min: number) {
+        return keep(this.particles, (particle) => particle.duration >= min * this.lifetime).length;
+    }
+}
+
+class Late {
+    constructor(public duration: number) {}
+}
+
+function main() {
+    const g = new Gen();
+    g.particles.push(new Late(0.0), new Late(100.0), new Late(0.0));
+    assert(g.update() == 2, "the long-lived one is kept out");
+    assert(g.particles.length == 2 && g.firstDuration() == 0.0);
+    g.particles.push(new Late(7.0));
+    assert(g.longer(1.0) == 1, "a second closure over this");
+    print("done.");
+}

tslang/tslang/tslang.cpp

@@ -19,6 +19,7 @@
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/DLTI/DLTI.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.h"
 #ifdef ENABLE_ASYNC
 #include "mlir/Dialect/Async/IR/Async.h"
 #endif
@@ -338,6 +339,9 @@ std::unique_ptr<mlir::MLIRContext> createMLIRContext()
 {
     mlir::DialectRegistry registry;
     registerAllExtensions(registry);
+    // the llvm dialect promises an inliner interface that only registerAllDialects provides: without
+    // it a pass that inlines past lowering (the JIT's) failed the promise check in an assertions build
+    mlir::LLVM::registerInlinerInterface(registry);
 
     auto mlirContext = std::make_unique<mlir::MLIRContext>(registry);
     // Load our Dialect in this MLIR Context.