Cleanup global memory allocations on shutdown. (#4731)

Co-authored-by: Jay Kwak <82421531+jkwak-work@users.noreply.github.com>

7 files changed, 23 insertions, 5 deletions

diff --git a/source/compiler-core/slang-spirv-core-grammar.h b/source/compiler-core/slang-spirv-core-grammar.h
index 7fe8f6cd7..958aaaef2 100644
--- a/source/compiler-core/slang-spirv-core-grammar.h
+++ b/source/compiler-core/slang-spirv-core-grammar.h
@@ -16,7 +16,8 @@ namespace Slang
     struct SPIRVCoreGrammarInfo : public RefObject
     {
         static RefPtr<SPIRVCoreGrammarInfo> loadFromJSON(SourceView& source, DiagnosticSink& sink);
-        static RefPtr<SPIRVCoreGrammarInfo> getEmbeddedVersion();
+        static RefPtr<SPIRVCoreGrammarInfo>& getEmbeddedVersion();
+        static inline void freeEmbeddedGrammerInfo() { getEmbeddedVersion() = nullptr; }
 
         template<typename K, typename T>
         struct Lookup
diff --git a/source/core/slang-performance-profiler.cpp b/source/core/slang-performance-profiler.cpp
index faf06bb92..b480e1c8e 100644
--- a/source/core/slang-performance-profiler.cpp
+++ b/source/core/slang-performance-profiler.cpp
@@ -42,6 +42,10 @@ namespace Slang
         {
             data.clear();
         }
+        virtual void dispose() override
+        {
+            data = decltype(data)();
+        }
     };
 
     PerformanceProfiler* Slang::PerformanceProfiler::getProfiler()
diff --git a/source/core/slang-performance-profiler.h b/source/core/slang-performance-profiler.h
index 372a7b273..71b34d262 100644
--- a/source/core/slang-performance-profiler.h
+++ b/source/core/slang-performance-profiler.h
@@ -29,6 +29,7 @@ public:
     virtual void exitFunction(FuncProfileContext context) = 0;
     virtual void getResult(StringBuilder& out) = 0;
     virtual void clear() = 0;
+    virtual void dispose() = 0;
 public:
     static PerformanceProfiler* getProfiler();
 };
diff --git a/source/slang/slang-api.cpp b/source/slang/slang-api.cpp
index c4f668fb2..412cec0d4 100644
--- a/source/slang/slang-api.cpp
+++ b/source/slang/slang-api.cpp
@@ -1,9 +1,10 @@
 // slang-api.cpp
 
 #include "slang-compiler.h"
-
 #include "slang-repro.h"
-
+#include "slang-capability.h"
+#include "../core/slang-rtti-info.h"
+#include "../core/slang-performance-profiler.h"
 #include "../core/slang-shared-library.h"
 #include "../slang-record-replay/record/slang-global-session.h"
 #include "../slang-record-replay/util/record-utility.h"
@@ -139,6 +140,14 @@ SLANG_API SlangResult slang_createGlobalSession(
     return SLANG_OK;
 }
 
+SLANG_API void slang_shutdown()
+{
+    Slang::PerformanceProfiler::getProfiler()->dispose();
+    Slang::SPIRVCoreGrammarInfo::freeEmbeddedGrammerInfo();
+    Slang::RttiInfo::deallocateAll();
+    Slang::freeCapabilityDefs();
+}
+
 SLANG_API SlangResult slang_createGlobalSessionWithoutStdLib(
     SlangInt                apiVersion,
     slang::IGlobalSession** outGlobalSession)
diff --git a/source/slang/slang-capability.h b/source/slang/slang-capability.h
index 299e956e4..79b1f7068 100644
--- a/source/slang/slang-capability.h
+++ b/source/slang/slang-capability.h
@@ -371,6 +371,8 @@ const CapabilityAtomSet& getAtomSetOfStages();
 
 bool hasTargetAtom(const CapabilityAtomSet& setIn, CapabilityAtom& targetAtom);
 
+void freeCapabilityDefs();
+
 //#define UNIT_TEST_CAPABILITIES
 #ifdef UNIT_TEST_CAPABILITIES
 void TEST_CapabilitySet();
diff --git a/source/slang/slang.cpp b/source/slang/slang.cpp
index c6e460d02..6f5e79458 100644
--- a/source/slang/slang.cpp
+++ b/source/slang/slang.cpp
@@ -140,7 +140,7 @@ void Session::init()
 {
     SLANG_ASSERT(BaseTypeInfo::check());
 
-    
+
     _initCodeGenTransitionMap();
 
     ::memset(m_downstreamCompilerLocators, 0, sizeof(m_downstreamCompilerLocators));
diff --git a/source/slangc/main.cpp b/source/slangc/main.cpp
index e69a6655d..2907625ed 100644
--- a/source/slangc/main.cpp
+++ b/source/slangc/main.cpp
@@ -104,7 +104,7 @@ SLANG_TEST_TOOL_API SlangResult innerMain(StdWriters* stdWriters, slang::IGlobal
     SlangResult res = _compile(compileRequest, argc, argv);
     // Now that we are done, clean up after ourselves
     spDestroyCompileRequest(compileRequest);
-
+    
     return res;
 }
 
@@ -112,6 +112,7 @@ int MAIN(int argc, char** argv)
 {
     auto stdWriters = StdWriters::initDefaultSingleton();
     SlangResult res = innerMain(stdWriters, nullptr, argc, argv);
+    slang::shutdown();
     return (int)TestToolUtil::getReturnCode(res);
 }