From 25d01e0a210c7608ae7d6087b30812fcbeff7bd1 Mon Sep 17 00:00:00 2001
From: ASDAlexander77 <duzhar@googlemail.com>
Date: Sun, 12 Jul 2026 19:26:45 +0100
Subject: [PATCH] Store function signatures instead of FuncOp handles in
 functionMap (review doc 2)

NamespaceInfo::functionMap cached raw mlir_ts::FuncOp handles; since
the discovery pass emits into a throwaway module (#213), any handle
registered during discovery is guaranteed to dangle once the module is
erased - consumers survived only because they read the (context-owned)
type early. The map now stores FunctionEntry { symbol name, function
type }, which is all any consumer used; the one site that needs a live
op (generic-instantiation short-circuit) resolves it through
theModule.lookupSymbol with a short/full-name fallback.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 .../TypeScript/MLIRLogic/MLIRGenStore.h       | 16 ++++++++++-
 tslang/lib/TypeScript/MLIRGen.cpp             | 28 +++++++++++--------
 2 files changed, 32 insertions(+), 12 deletions(-)

diff --git a/tslang/include/TypeScript/MLIRLogic/MLIRGenStore.h b/tslang/include/TypeScript/MLIRLogic/MLIRGenStore.h
index 77a4f99d1..01244ce8f 100644
--- a/tslang/include/TypeScript/MLIRLogic/MLIRGenStore.h
+++ b/tslang/include/TypeScript/MLIRLogic/MLIRGenStore.h
@@ -985,6 +985,20 @@ struct GenericClassInfo
     }
 };
 
+// what we know about a registered function; deliberately not the FuncOp itself -
+// discovery-pass ops are erased with the discovery module, so cached op handles dangle.
+// Resolve a live op through theModule.lookupSymbol when one is actually needed.
+struct FunctionEntry
+{
+    std::string name;
+    mlir_ts::FunctionType funcType;
+
+    explicit operator bool() const
+    {
+        return static_cast<bool>(funcType);
+    }
+};
+
 struct NamespaceInfo
 {
   public:
@@ -998,7 +1012,7 @@ struct NamespaceInfo
 
     llvm::StringMap<mlir_ts::FunctionType> functionTypeMap;
 
-    llvm::StringMap<mlir_ts::FuncOp> functionMap;
+    llvm::StringMap<FunctionEntry> functionMap;
 
     llvm::StringMap<GenericFunctionInfo::TypePtr> genericFunctionMap;
 
diff --git a/tslang/lib/TypeScript/MLIRGen.cpp b/tslang/lib/TypeScript/MLIRGen.cpp
index 12a8edec8..4bfd919f6 100644
--- a/tslang/lib/TypeScript/MLIRGen.cpp
+++ b/tslang/lib/TypeScript/MLIRGen.cpp
@@ -5312,7 +5312,7 @@ class MLIRGenImpl
         auto funcIt = getFunctionMap().find(name);
         if (funcIt != getFunctionMap().end())
         {
-            auto cachedFuncType = funcIt->second.getFunctionType();
+            auto cachedFuncType = funcIt->second.funcType;
             if (cachedFuncType.getNumResults() > 0)
             {
                 auto returnType = cachedFuncType.getResult(0);
@@ -6032,7 +6032,8 @@ class MLIRGenImpl
         auto name = funcProto->getNameWithoutNamespace();
         if (!getFunctionMap().count(name))
         {
-            getFunctionMap().insert({name, funcOp});
+            getFunctionMap().insert(
+                {name, FunctionEntry{funcOp.getName().str(), mlir::cast<mlir_ts::FunctionType>(funcOp.getFunctionType())}});
 
             LLVM_DEBUG(llvm::dbgs() << "\n!! reg. func: " << name << " type:" << funcOp.getFunctionType() << " function name: " << funcProto->getName()
                                     << " num inputs:" << mlir::cast<mlir_ts::FunctionType>(funcOp.getFunctionType()).getNumInputs()
@@ -6081,10 +6082,18 @@ class MLIRGenImpl
         {
             auto [fullFunctionName, functionName] = getNameOfFunction(functionLikeDeclarationBaseAST, funcDeclGenContext);
 
-            auto funcOp = lookupFunctionMap(functionName);
-            if (funcOp && theModule.lookupSymbol(functionName) 
+            auto funcEntry = lookupFunctionMap(functionName);
+            if (funcEntry && theModule.lookupSymbol(functionName)
                 || theModule.lookupSymbol(fullFunctionName))
             {
+                // resolve a live op from the module instead of returning a cached handle;
+                // the registered symbol is usually the full name
+                auto funcOp = theModule.lookupSymbol<mlir_ts::FuncOp>(functionName);
+                if (!funcOp)
+                {
+                    funcOp = theModule.lookupSymbol<mlir_ts::FuncOp>(fullFunctionName);
+                }
+
                 return {mlir::success(), funcOp, functionName, false};
             }
         }
@@ -15966,11 +15975,8 @@ class MLIRGenImpl
         auto fn = getFunctionMap().find(name);
         if (fn != getFunctionMap().end())
         {
-            auto funcOp = fn->getValue();
-            auto funcType = funcOp.getFunctionType();
-            auto funcName = funcOp.getName();
-
-            return resolveFunctionWithCapture(location, funcName, funcType, mlir::Value(), false, genContext);
+            auto &funcEntry = fn->getValue();
+            return resolveFunctionWithCapture(location, funcEntry.name, funcEntry.funcType, mlir::Value(), false, genContext);
         }
 
         return mlir::Value();
@@ -26273,12 +26279,12 @@ genContext);
         lookupLogic(functionTypeMap);
     }
 
-    auto getFunctionMap() -> llvm::StringMap<mlir_ts::FuncOp> &
+    auto getFunctionMap() -> llvm::StringMap<FunctionEntry> &
     {
         return currentNamespace->functionMap;
     }
 
-    auto lookupFunctionMap(StringRef name) -> mlir_ts::FuncOp
+    auto lookupFunctionMap(StringRef name) -> FunctionEntry
     {
         lookupLogic(functionMap);
     }