diff --git a/docs/MLIRGen-refactoring-review.md b/docs/MLIRGen-refactoring-review.md index b3d0a49eb..9c0ff9eda 100644 --- a/docs/MLIRGen-refactoring-review.md +++ b/docs/MLIRGen-refactoring-review.md @@ -70,7 +70,7 @@ Refactor: move traversal state into `GenContext` (it is already threaded through `mlirGenSourceFile` runs the whole AST twice: a `dummyRun/allowPartialResolve` discovery pass whose ops are thrown away, then real codegen. Costs: 2× walk, op churn, and cleanup fragility (see §2; the nested-import bug lived here). Shorter-term hardening (cheap): -- Build discovery ops into a **dedicated throwaway `ModuleOp`** with its own builder instead of the real `theModule` — then cleanup is `discoveryModule.erase()` and can never touch real content. (This supersedes the current snapshot-and-erase in `mlirDiscoverAllDependencies`; the snapshot fix is correct but a separate module makes the invariant structural.) Note: requires routing the builder's insertion point for the whole discovery call tree, so do it as a deliberate change with tests, not a drive-by. +- Build discovery ops into a **dedicated throwaway `ModuleOp`** with its own builder instead of the real `theModule` — then cleanup is `discoveryModule.erase()` and can never touch real content. (This supersedes the current snapshot-and-erase in `mlirDiscoverAllDependencies`; the snapshot fix is correct but a separate module makes the invariant structural.) Note: requires routing the builder's insertion point for the whole discovery call tree, so do it as a deliberate change with tests, not a drive-by. *Status: done — `DiscoveryModuleScope` (same pattern as `TempModuleScope`) swaps `theModule`/builder for the duration of `mlirDiscoverAllDependencies` and erases the throwaway module on scope exit; the snapshot-and-selective-erase is gone. A side effect worth knowing: during a nested discovery the `theModule.lookupSymbol` early-outs (synthesized-function and specialized-generic checks) no longer see real-module content, so discovery may redo that work in the throwaway module — correctness is unaffected because registration maps insert-if-absent.* Longer term: replace the discovery codegen pass with a symbol-collection pass over the AST that populates the maps *without creating IR at all*. diff --git a/tslang/lib/TypeScript/MLIRGen.cpp b/tslang/lib/TypeScript/MLIRGen.cpp index 41adaf576..12a8edec8 100644 --- a/tslang/lib/TypeScript/MLIRGen.cpp +++ b/tslang/lib/TypeScript/MLIRGen.cpp @@ -747,16 +747,11 @@ class MLIRGenImpl llvm::ScopedHashTableScope fullNameGlobalsMapScope( fullNameGlobalsMap); - // Snapshot the ops already present in the module. When this discovery pass is nested - // (an 'import' of a local source file triggers mlirGenInclude during SourceGeneration), - // it must not wipe the real module content that was already generated (e.g. default-lib - // function bodies such as 'console.log'). Clearing the whole body left dangling - // SymbolRefs and produced "op 'console.log' does not reference a valid function". - llvm::SmallPtrSet preExistingOps; - for (auto &op : theModule.getBody()->getOperations()) - { - preExistingOps.insert(&op); - } + // Discovery emits into a throwaway module, so its cleanup can never disturb real module + // content. When this discovery pass is nested (an 'import' of a local source file triggers + // mlirGenInclude during SourceGeneration), the real module already holds generated content + // (e.g. default-lib function bodies such as 'console.log') that must survive. + DiscoveryModuleScope discoveryModuleScope(*this); // Process of discovery here GenContext genContextPartial{}; @@ -778,30 +773,8 @@ class MLIRGenImpl auto notResolved = processStatements(module->statements, genContextPartial); - // clean up: erase only the ops this discovery pass added, preserving any content that - // was already generated before a nested discovery (see preExistingOps above). + // clean up; the ops this pass created go away with the discovery module on scope exit clearTempModule(); - { - llvm::SmallVector addedOps; - for (auto &op : theModule.getBody()->getOperations()) - { - if (!preExistingOps.contains(&op)) - { - addedOps.push_back(&op); - } - } - - for (auto *op : addedOps) - { - op->dropAllReferences(); - } - - for (auto *op : llvm::reverse(addedOps)) - { - op->dropAllUses(); - op->erase(); - } - } // clear state for (auto &statement : module->statements) @@ -20551,6 +20524,34 @@ genContext); return mlir::success(); } + // RAII scope that redirects theModule and the builder into a fresh throwaway + // module for the discovery pass. On scope exit the discovery module is erased + // with everything the pass created, and theModule/builder are restored, so + // discovery cleanup is structurally unable to touch real module content. + class DiscoveryModuleScope + { + public: + DiscoveryModuleScope(MLIRGenImpl &mlirGenImpl) + : moduleGuard(mlirGenImpl.theModule), insertGuard(mlirGenImpl.builder) + { + discoveryModule = + mlir::ModuleOp::create(mlirGenImpl.theModule.getLoc(), mlir::StringRef("discovery_module")); + mlirGenImpl.theModule = discoveryModule; + mlirGenImpl.builder.setInsertionPointToStart(discoveryModule.getBody()); + } + + ~DiscoveryModuleScope() + { + // members restore theModule and the insertion point after the erase + discoveryModule.erase(); + } + + private: + MLIRValueGuard moduleGuard; + mlir::OpBuilder::InsertionGuard insertGuard; + mlir::ModuleOp discoveryModule; + }; + // RAII scope that redirects theModule and the builder into the temp module // for speculative evaluation and restores both when it goes out of scope. class TempModuleScope