From 61fa69f39df94a8d38d3dc730cb77a6fedc8b12f Mon Sep 17 00:00:00 2001
From: Tomas Davidovic <tdavidovic@nvidia.com>
Date: Thu, 25 Sep 2025 06:11:47 +0200
Subject: [PATCH] Added code that ensures SlangSession does not return 2
 different SlangModules for one slang::IModule.

---
 src/sgl/device/shader.cpp | 19 +++++++++++++++++++
 src/sgl/device/shader.h   |  9 +++++++++
 2 files changed, 28 insertions(+)

diff --git a/src/sgl/device/shader.cpp b/src/sgl/device/shader.cpp
index 4f26d5e7..0fd88a1a 100644
--- a/src/sgl/device/shader.cpp
+++ b/src/sgl/device/shader.cpp
@@ -511,6 +511,8 @@ ref<SlangModule> SlangSession::load_module_from_source(
     module->load(build);
     module->store_built_data(build);
 
+    module = deduplicate_module(std::move(module));
+
     // Update cache of loaded modules.
     update_module_cache_and_dependencies();
 
@@ -625,6 +627,11 @@ void SlangSession::_unregister_module(SlangModule* module)
     auto existing = std::find(m_registered_modules.begin(), m_registered_modules.end(), module);
     SGL_ASSERT(existing != m_registered_modules.end());
     m_registered_modules.erase(existing);
+
+    if (auto it = m_sgl_to_cache_location_map.find(module); it != m_sgl_to_cache_location_map.end()) {
+        m_slang_to_sgl_module_map.erase(it->second);
+        m_sgl_to_cache_location_map.erase(it);
+    }
 }
 
 std::string SlangSession::to_string() const
@@ -727,6 +734,18 @@ bool SlangSession::write_module_to_cache(slang::IModule* module)
     return true;
 }
 
+ref<SlangModule> SlangSession::deduplicate_module(ref<SlangModule>&& module)
+{
+    auto it = m_slang_to_sgl_module_map.find(module->slang_module());
+    if (it != m_slang_to_sgl_module_map.end())
+        return ref<SlangModule>(it->second);
+
+    it = m_slang_to_sgl_module_map.insert(std::make_pair(module->slang_module(), module.get())).first;
+    m_sgl_to_cache_location_map[module.get()] = it;
+
+    return module;
+}
+
 std::string SlangSessionData::resolve_module_name(std::string_view module_name) const
 {
     // Return if module name is an absolute file path.
diff --git a/src/sgl/device/shader.h b/src/sgl/device/shader.h
index 60ecb4eb..40450581 100644
--- a/src/sgl/device/shader.h
+++ b/src/sgl/device/shader.h
@@ -335,12 +335,21 @@ class SGL_API SlangSession : public Object {
     /// Note: this is a vector, as order of creation matters.
     std::vector<SlangModule*> m_registered_modules;
 
+    /// Map to deduplicate sgl::SlangModule that reference same slang::IModule
+    std::map<slang::IModule*, SlangModule*> m_slang_to_sgl_module_map;
+
+    /// Caches the entry in the m_slang_to_sgl_module_map for faster deletion.
+    std::map<SlangModule*, std::map<slang::IModule*, SlangModule*>::iterator> m_sgl_to_cache_location_map;
+
     /// All created sgl programs (via link_program)
     std::set<ShaderProgram*> m_registered_programs;
 
     void update_module_cache_and_dependencies();
     bool write_module_to_cache(slang::IModule* module);
     void create_session(SlangSessionBuild& build);
+
+    // Returns either the same module, or a different `SlangModule` that references identical `slang::IModule`
+    ref<SlangModule> deduplicate_module(ref<SlangModule>&& module);
 };
 
 struct SlangModuleDesc {