Fix IR memory leaks.

1, make IRModule class own a memory pool for all IR object allocations
2. For now, we allow IR objects to own other (externally) heap allocated objects, such as String, List and RefPtrs by tracking all IR objects that  has been allocated for the IRModule in a list named `IRModule::irObjectsToFree`. and call destructor for all these objects upon the destruction of the IRModule. In the long term, we should eliminate the use of all these externally allocated types in IR system and get rid of this tracking and explicit destructor calls.
3. remove non-generic `createValueImpl` functions and retain only generic versions in IRBulider so we can properly call the constructor of the IR types to set up virtual tables correctly for destructor dispatching.
4. add `MemoryPool` class for allocation of the IR objects.
5. Make sure we are disposing IRSpecContexts when we are done with the specialized IR module.
6. Add `_CrtDumpMemoryLeaks()` calls to check memory leaks upon destruction of a Slang session. If we are to support multiple sessions at a time, this call should probably be replaced with the more advanced MemoryState versions of the memory leak checker.

11 files changed, 141 insertions, 61 deletions

diff --git a/slang.h b/slang.h
index 422de04cc..724e49d13 100644
--- a/slang.h
+++ b/slang.h
@@ -1237,6 +1237,7 @@ namespace slang
 #include "source/slang/dxc-support.cpp"
 #include "source/slang/emit.cpp"
 #include "source/slang/ir.cpp"
+#include "source/slang/memory_pool.cpp"
 #include "source/slang/ir-legalize-types.cpp"
 #include "source/slang/ir-ssa.cpp"
 #include "source/slang/legalize-types.cpp"
diff --git a/source/slang/compiler.h b/source/slang/compiler.h
index fab1e19c8..7ead5d07e 100644
--- a/source/slang/compiler.h
+++ b/source/slang/compiler.h
@@ -179,7 +179,7 @@ namespace Slang
         // This will only be valid/non-null after semantic
         // checking and IR generation are complete, so it
         // is not safe to use this field without testing for NULL.
-        IRModule* irModule;
+        RefPtr<IRModule> irModule;
     };
 
     // A request to generate output in some target format
@@ -225,7 +225,7 @@ namespace Slang
         RefPtr<ModuleDecl>  moduleDecl;
 
         // The IR for the module
-        IRModule* irModule = nullptr;
+        RefPtr<IRModule> irModule = nullptr;
     };
 
     class Session;
diff --git a/source/slang/emit.cpp b/source/slang/emit.cpp
index 672798359..b8fbc350e 100644
--- a/source/slang/emit.cpp
+++ b/source/slang/emit.cpp
@@ -8158,6 +8158,7 @@ String emitEntryPoint(
         // TODO: do we want to emit directly from IR, or translate the
         // IR back into AST for emission?
         visitor.emitIRModule(&context, irModule);
+        destroyIRSpecializationState(irSpecializationState);
     }
 
     String code = sharedContext.sb.ProduceString();
@@ -8189,7 +8190,7 @@ String emitEntryPoint(
     finalResultBuilder << code;
 
     String finalResult = finalResultBuilder.ProduceString();
-
+    
     return finalResult;
 }
 
diff --git a/source/slang/ir-insts.h b/source/slang/ir-insts.h
index e2c1c2e03..1cfe83b97 100644
--- a/source/slang/ir-insts.h
+++ b/source/slang/ir-insts.h
@@ -610,22 +610,26 @@ struct IRBuilder
         UInt            caseArgCount,
         IRValue* const* caseArgs);
 
-
-    IRDecoration* addDecorationImpl(
-        IRValue*        value,
-        UInt            decorationSize,
-        IRDecorationOp  op);
-
     template<typename T>
     T* addDecoration(IRValue* value, IRDecorationOp op)
     {
-        return (T*) addDecorationImpl(value, sizeof(T), op);
+        assert(getModule());
+        auto decorationSize = sizeof(T);
+        auto decoration = (T*)getModule()->memoryPool.allocZero(decorationSize);
+        new(decoration)T();
+       
+        decoration->op = op;
+
+        decoration->next = value->firstDecoration;
+        value->firstDecoration = decoration;
+        getModule()->irObjectsToFree.Add(decoration);
+        return decoration;
     }
 
     template<typename T>
     T* addDecoration(IRValue* value)
     {
-        return (T*) addDecorationImpl(value, sizeof(T), IRDecorationOp(T::kDecorationOp));
+        return addDecoration<T>(value, IRDecorationOp(T::kDecorationOp));
     }
 
     IRHighLevelDeclDecoration* addHighLevelDeclDecoration(IRValue* value, Decl* decl);
diff --git a/source/slang/ir.cpp b/source/slang/ir.cpp
index 441db09a6..8667a67bd 100644
--- a/source/slang/ir.cpp
+++ b/source/slang/ir.cpp
@@ -543,26 +543,19 @@ namespace Slang
         value->sourceLoc = sourceLocInfo->sourceLoc;
     }
 
-    static IRValue* createValueImpl(
-        IRBuilder*  /*builder*/,
-        UInt        size,
-        IROp        op,
-        IRType*     type)
-    {
-        IRValue* value = (IRValue*) malloc(size);
-        memset(value, 0, size);
-        value->op = op;
-        value->type = type;
-        return value;
-    }
-
     template<typename T>
     static T* createValue(
         IRBuilder*  builder,
         IROp        op,
         IRType*     type)
     {
-        return (T*) createValueImpl(builder, sizeof(T), op, type);
+        assert(builder->getModule());
+        T* value = (T*)builder->getModule()->memoryPool.allocZero(sizeof(T));
+        new(value)T();
+        value->op = op;
+        value->type = type;
+        builder->getModule()->irObjectsToFree.Add(value);
+        return value;
     }
 
 
@@ -571,7 +564,8 @@ namespace Slang
     // In this case `argCount` and `args` represnt the
     // arguments *after* the type (which is a mandatory
     // argument for all instructions).
-    static IRInst* createInstImpl(
+    template<typename T>
+    static T* createInstImpl(
         IRBuilder*      builder,
         UInt            size,
         IROp            op,
@@ -581,8 +575,9 @@ namespace Slang
         UInt            varArgCount = 0,
         IRValue* const* varArgs = nullptr)
     {
-        IRInst* inst = (IRInst*) malloc(size);
-        memset(inst, 0, size);
+        assert(builder->getModule());
+        T* inst = (T*)builder->getModule()->memoryPool.allocZero(size);
+        new(inst)T();
 
         inst->argCount = (uint32_t)(fixedArgCount + varArgCount);
 
@@ -611,7 +606,7 @@ namespace Slang
             }
             operand++;
         }
-
+        builder->getModule()->irObjectsToFree.Add(inst);
         return inst;
     }
 
@@ -623,7 +618,7 @@ namespace Slang
         UInt            argCount,
         IRValue* const* args)
     {
-        return (T*)createInstImpl(
+        return createInstImpl<T>(
             builder,
             sizeof(T),
             op,
@@ -638,7 +633,7 @@ namespace Slang
         IROp            op,
         IRType*         type)
     {
-        return (T*)createInstImpl(
+        return createInstImpl<T>(
             builder,
             sizeof(T),
             op,
@@ -654,7 +649,7 @@ namespace Slang
         IRType*         type,
         IRValue*        arg)
     {
-        return (T*)createInstImpl(
+        return createInstImpl<T>(
             builder,
             sizeof(T),
             op,
@@ -672,7 +667,7 @@ namespace Slang
         IRValue*        arg2)
     {
         IRValue* args[] = { arg1, arg2 };
-        return (T*)createInstImpl(
+        return createInstImpl<T>(
             builder,
             sizeof(T),
             op,
@@ -689,7 +684,7 @@ namespace Slang
         UInt            argCount,
         IRValue* const* args)
     {
-        return (T*)createInstImpl(
+        return createInstImpl<T>(
             builder,
             sizeof(T) + argCount * sizeof(IRUse),
             op,
@@ -708,7 +703,7 @@ namespace Slang
         UInt            varArgCount,
         IRValue* const* varArgs)
     {
-        return (T*)createInstImpl(
+        return createInstImpl<T>(
             builder,
             sizeof(T) + varArgCount * sizeof(IRUse),
             op,
@@ -731,7 +726,7 @@ namespace Slang
         IRValue* fixedArgs[] = { arg1 };
         UInt fixedArgCount = sizeof(fixedArgs) / sizeof(fixedArgs[0]);
 
-        return (T*)createInstImpl(
+        return createInstImpl<T>(
             builder,
             sizeof(T) + varArgCount * sizeof(IRUse),
             op,
@@ -1705,22 +1700,6 @@ namespace Slang
         return inst;
     }
 
-    IRDecoration* IRBuilder::addDecorationImpl(
-        IRValue*        inst,
-        UInt            decorationSize,
-        IRDecorationOp  op)
-    {
-        auto decoration = (IRDecoration*) malloc(decorationSize);
-        memset(decoration, 0, decorationSize);
-
-        decoration->op = op;
-
-        decoration->next = inst->firstDecoration;
-        inst->firstDecoration = decoration;
-
-        return decoration;
-    }
-
     IRHighLevelDeclDecoration* IRBuilder::addHighLevelDeclDecoration(IRValue* inst, Decl* decl)
     {
         auto decoration = addDecoration<IRHighLevelDeclDecoration>(inst, kIRDecorationOp_HighLevelDecl);
@@ -2666,9 +2645,6 @@ namespace Slang
 #else
         // Run destructor to be sure...
         this->~IRValue();
-
-        // And then free the memory
-        free((void*) this);
 #endif
     }
 
@@ -3729,7 +3705,7 @@ namespace Slang
         CodeGenTarget target;
 
         // The specialized module we are building
-        IRModule*   module;
+        RefPtr<IRModule>   module;
 
         // The original, unspecialized module we are copying
         IRModule*   originalModule;
@@ -5726,11 +5702,9 @@ namespace Slang
         for (auto workItem : witnessTablesToSpecailize)
         {
             int diff = 0;
-            specializeWitnessTable(context->shared, workItem.srcTable, 
+            specializeWitnessTable(context->shared, workItem.srcTable,
                 workItem.specDeclRef.SubstituteImpl(SubstitutionSet(nullptr, nullptr, globalParamSubst), &diff), workItem.dstTable);
         }
         return globalParamSubst;
     }
-
-
 }
diff --git a/source/slang/ir.h b/source/slang/ir.h
index 0dc9daa9a..4cdd1f1de 100644
--- a/source/slang/ir.h
+++ b/source/slang/ir.h
@@ -10,6 +10,7 @@
 #include "../core/basic.h"
 
 #include "source-loc.h"
+#include "memory_pool.h"
 
 namespace Slang {
 
@@ -114,15 +115,26 @@ enum IRDecorationOp : uint16_t
     kIRDecorationOp_TargetIntrinsic,
 };
 
+// represents an object allocated in an IR memory pool 
+struct IRObject
+{
+    bool isDestroyed = false;
+    virtual ~IRObject()
+    {
+        isDestroyed = true;
+    }
+};
+
 // A "decoration" that gets applied to an instruction.
 // These usually don't affect semantics, but are useful
 // for preserving high-level source information.
-struct IRDecoration
+struct IRDecoration : public IRObject
 {
     // Next decoration attached to the same instruction
     IRDecoration* next;
 
     IRDecorationOp op;
+    virtual ~IRDecoration() = default;
 };
 
 // Use AST-level types directly to represent the
@@ -132,7 +144,7 @@ typedef Type IRType;
 struct IRBlock;
 
 // Base class for values in the IR
-struct IRValue
+struct IRValue : public IRObject
 {
     // The operation that this value represents
     IROp op;
@@ -536,6 +548,8 @@ struct IRModule : RefObject
 {
     // The compilation session in use.
     Session*    session;
+    MemoryPool  memoryPool;
+    List<IRObject*> irObjectsToFree; // list of ir objects to run destructor upon destruction
 
     // A list of all the functions and other
     // global values declared in this module.
@@ -544,6 +558,14 @@ struct IRModule : RefObject
 
     IRGlobalValue*  getFirstGlobalValue() { return firstGlobalValue; }
     IRGlobalValue*  getlastGlobalValue() { return lastGlobalValue; }
+
+    ~IRModule()
+    {
+        for (auto val : irObjectsToFree)
+            if (!val->isDestroyed)
+                val->~IRObject();
+        irObjectsToFree = List<IRObject*>();
+    }
 };
 
 void printSlangIRAssembly(StringBuilder& builder, IRModule* module);
diff --git a/source/slang/memory_pool.cpp b/source/slang/memory_pool.cpp
new file mode 100644
index 000000000..884c9cfe4
--- /dev/null
+++ b/source/slang/memory_pool.cpp
@@ -0,0 +1,52 @@
+#include "memory_pool.h"
+
+namespace Slang
+{
+    const size_t kPoolSegmentSize = 4 << 20; // use 4MB segments
+
+    struct MemoryPoolSegment
+    {
+        unsigned char* data = nullptr;
+        size_t allocPtr = 0;
+        MemoryPoolSegment* nextSegment = nullptr;
+    };
+
+    MemoryPool::~MemoryPool()
+    {
+        while (curSegment)
+        {
+            auto nxtSegment = curSegment->nextSegment;
+            free(curSegment->data);
+            delete curSegment;
+            curSegment = nxtSegment;
+        }
+    }
+
+    void newSegment(MemoryPool* pool)
+    {
+        auto seg = new MemoryPoolSegment();
+        seg->nextSegment = pool->curSegment;
+        seg->data = (unsigned char*)malloc(kPoolSegmentSize);
+        pool->curSegment = seg;
+    }
+
+    void * MemoryPool::alloc(size_t size)
+    {
+        assert(size < kPoolSegmentSize);
+        // ensure there is a segment available
+        if (!curSegment)
+            newSegment(this);
+        if (curSegment->allocPtr + size > kPoolSegmentSize)
+            newSegment(this);
+        // alloc memory from current segment
+        void* rs = curSegment->data + curSegment->allocPtr;
+        curSegment->allocPtr += size;
+        return rs;
+    }
+    void * MemoryPool::allocZero(size_t size)
+    {
+        auto rs = alloc(size);
+        memset(rs, 0, size);
+        return rs;
+    }
+}
diff --git a/source/slang/memory_pool.h b/source/slang/memory_pool.h
new file mode 100644
index 000000000..ee3ee1aa3
--- /dev/null
+++ b/source/slang/memory_pool.h
@@ -0,0 +1,19 @@
+#ifndef SLANG_MEMORY_POOL_H
+#define SLANG_MEMORY_POOL_H
+
+#include "../core/basic.h"
+
+namespace Slang
+{
+    struct MemoryPoolSegment;
+
+    struct MemoryPool : public RefObject
+    {
+        MemoryPoolSegment* curSegment = nullptr;
+        ~MemoryPool();
+        void* alloc(size_t size);
+        void* allocZero(size_t size);
+    };
+}
+
+#endif
\ No newline at end of file
diff --git a/source/slang/slang.cpp b/source/slang/slang.cpp
index 4238681f3..0f5d8e17a 100644
--- a/source/slang/slang.cpp
+++ b/source/slang/slang.cpp
@@ -751,6 +751,9 @@ SLANG_API void spDestroySession(
 {
     if(!session) return;
     delete SESSION(session);
+#ifdef _MSC_VER
+    _CrtDumpMemoryLeaks();
+#endif
 }
 
 SLANG_API void spAddBuiltins(
diff --git a/source/slang/slang.vcxproj b/source/slang/slang.vcxproj
index 1ff88020f..0154dd98f 100644
--- a/source/slang/slang.vcxproj
+++ b/source/slang/slang.vcxproj
@@ -186,6 +186,7 @@
     <ClInclude Include="lookup.h" />
     <ClInclude Include="lower-to-ir.h" />
     <ClInclude Include="mangle.h" />
+    <ClInclude Include="memory_pool.h" />
     <ClInclude Include="modifier-defs.h" />
     <ClInclude Include="name.h" />
     <ClInclude Include="object-meta-begin.h" />
@@ -225,6 +226,7 @@
     <ClCompile Include="lookup.cpp" />
     <ClCompile Include="lower-to-ir.cpp" />
     <ClCompile Include="mangle.cpp" />
+    <ClCompile Include="memory_pool.cpp" />
     <ClCompile Include="name.cpp" />
     <ClCompile Include="options.cpp" />
     <ClCompile Include="parameter-binding.cpp" />
diff --git a/source/slang/slang.vcxproj.filters b/source/slang/slang.vcxproj.filters
index e0bf59025..e64721f6d 100644
--- a/source/slang/slang.vcxproj.filters
+++ b/source/slang/slang.vcxproj.filters
@@ -45,6 +45,7 @@
     <ClInclude Include="mangle.h" />
     <ClInclude Include="legalize-types.h" />
     <ClInclude Include="ir-ssa.h" />
+    <ClInclude Include="memory_pool.h" />
   </ItemGroup>
   <ItemGroup>
     <ClCompile Include="check.cpp" />
@@ -75,6 +76,7 @@
     <ClCompile Include="ir-legalize-types.cpp" />
     <ClCompile Include="legalize-types.cpp" />
     <ClCompile Include="ir-ssa.cpp" />
+    <ClCompile Include="memory_pool.cpp" />
   </ItemGroup>
   <ItemGroup>
     <CustomBuild Include="core.meta.slang" />