Refactor GenContext cleanup logic to use stack-allocated structures and improve memory management (#204)

Author: ASDAlexander77

tslang/include/TypeScript/MLIRLogic/MLIRGenContext.h

@@ -75,11 +75,12 @@ struct GenContext
         receiverFuncType = mlir::Type();
     }
 
-    // TODO: you are using "theModule.getBody()->clear();", do you need this hack anymore?
+    // erases the IR collected during a dummy run (cleanup lists, dummy funcOp) and detaches from it;
+    // the lists, passResult and state are owned by the stack frame that set them, not by GenContext
     void clean()
     {
         if (cleanUps)
-        {    
+        {
             for (auto op : *cleanUps)
             {
                 op->dropAllDefinedValueUses();
@@ -88,7 +89,7 @@ struct GenContext
                 op->erase();
             }
 
-            delete cleanUps;
+            cleanUps->clear();
             cleanUps = nullptr;
         }
 
@@ -102,30 +103,16 @@ struct GenContext
                 op->erase();
             }
 
-            delete cleanUpOps;
-            cleanUpOps = nullptr;               
+            cleanUpOps->clear();
+            cleanUpOps = nullptr;
         }
 
-        if (passResult)
-        {
-            delete passResult;
-            passResult = nullptr;
-        }
-
-        cleanState();
+        passResult = nullptr;
+        state = nullptr;
 
         cleanFuncOp();
     }
 
-    void cleanState()
-    {
-        if (state)
-        {
-            delete state;
-            state = nullptr;
-        }
-    }
-
     void cleanFuncOp()
     {
         if (funcOp)

tslang/lib/TypeScript/MLIRGen.cpp

@@ -764,9 +764,6 @@ class MLIRGenImpl
         genContextPartial.dummyRun = true;
         genContextPartial.rootContext = &genContextPartial;
         genContextPartial.postponedMessages = &postponedMessages;
-        // TODO: no need to clean up here as whole module will be removed
-        //genContextPartial.cleanUps = new mlir::SmallVector<mlir::Block *>();
-        //genContextPartial.cleanUpOps = new mlir::SmallVector<mlir::Operation *>();
 
         for (auto includeFile : includeFiles)
         {
@@ -781,8 +778,6 @@ class MLIRGenImpl
 
         auto notResolved = processStatements(module->statements, genContextPartial);
 
-        genContextPartial.clean();
-
         // clean up: erase only the ops this discovery pass added, preserving any content that
         // was already generated before a nested discovery (see preExistingOps above).
         clearTempModule();
@@ -5557,16 +5552,22 @@ class MLIRGenImpl
             llvm::ScopedHashTableScope<StringRef, VariableDeclarationDOM::TypePtr> 
                 fullNameGlobalsMapScope(fullNameGlobalsMap);
 
+            // owned here; GenContext borrows pointers to them (see GenContext::clean)
+            SmallVector<mlir::Block *> cleanUpsList;
+            SmallVector<mlir::Operation *> cleanUpOpsList;
+            PassResult passResultData;
+            int discoverState = 1;
+
             GenContext genContextWithPassResult{};
             genContextWithPassResult.funcOp = dummyFuncOp;
             genContextWithPassResult.thisType = genContext.thisType;
             genContextWithPassResult.thisClassType = genContext.thisClassType;
             genContextWithPassResult.allowPartialResolve = true;
             genContextWithPassResult.dummyRun = true;
-            genContextWithPassResult.cleanUps = new SmallVector<mlir::Block *>();
-            genContextWithPassResult.cleanUpOps = new SmallVector<mlir::Operation *>();
-            genContextWithPassResult.passResult = new PassResult();
-            genContextWithPassResult.state = new int(1);
+            genContextWithPassResult.cleanUps = &cleanUpsList;
+            genContextWithPassResult.cleanUpOps = &cleanUpOpsList;
+            genContextWithPassResult.passResult = &passResultData;
+            genContextWithPassResult.state = &discoverState;
             genContextWithPassResult.allocateVarsInContextThis =
                 (functionLikeDeclarationBaseAST->internalFlags & InternalFlags::VarsInObjectContext) ==
                 InternalFlags::VarsInObjectContext;
@@ -6091,7 +6092,8 @@ class MLIRGenImpl
         auto funcGenContext = GenContext(funcDeclGenContext);
         funcGenContext.clearScopeVars();
         funcGenContext.funcOp = funcOp;
-        funcGenContext.state = new int(1);
+        int funcState = 1;
+        funcGenContext.state = &funcState;
         // if funcGenContext.passResult is null and allocateVarsInContextThis is true, this type should contain fully
         // defined object with local variables as fields
         funcGenContext.allocateVarsInContextThis =
@@ -6124,8 +6126,6 @@ class MLIRGenImpl
                 functionLikeDeclarationBaseAST, funcProto->getNameWithoutNamespace(), funcOp, funcProto, funcGenContext);
         }
 
-        funcGenContext.cleanState();
-
         if (mlir::failed(resultFromBody))
         {
             return {mlir::failure(), funcOp, "", false};