HnRenderDelegate: store lights in a light type -> prim map

Author: TheMostDiligent

Hydrogent/interface/HnRenderDelegate.hpp

@@ -512,8 +512,9 @@ class HnRenderDelegate final : public pxr::HdRenderDelegate
     std::unordered_set<HnMaterial*> m_Materials;
     HnMaterial*                     m_FallbackMaterial = nullptr;
 
-    std::mutex                   m_LightsMtx;
-    std::unordered_set<HnLight*> m_Lights;
+    std::mutex m_LightsMtx;
+    // Light type -> light prim
+    std::unordered_multimap<pxr::TfToken, HnLight*, pxr::TfToken::HashFunctor> m_Lights;
 
     Uint32 m_MeshResourcesVersion     = ~0u;
     Uint32 m_MaterialResourcesVersion = ~0u;

Hydrogent/src/HnRenderDelegate.cpp

@@ -654,7 +654,7 @@ pxr::HdSprim* HnRenderDelegate::CreateSprim(const pxr::TfToken& TypeId,
         HnLight* Light = HnLight::Create(SPrimId, TypeId);
         {
             std::lock_guard<std::mutex> Guard{m_LightsMtx};
-            m_Lights.emplace(Light);
+            m_Lights.emplace(TypeId, Light);
         }
         SPrim = Light;
     }
@@ -710,7 +710,16 @@ void HnRenderDelegate::DestroySprim(pxr::HdSprim* SPrim)
     else if (dynamic_cast<HnLight*>(SPrim) != nullptr)
     {
         std::lock_guard<std::mutex> Guard{m_LightsMtx};
-        m_Lights.erase(static_cast<HnLight*>(SPrim));
+
+        auto it_range = m_Lights.equal_range(static_cast<HnLight*>(SPrim)->GetTypeId());
+        for (auto it = it_range.first; it != it_range.second; ++it)
+        {
+            if (it->second == SPrim)
+            {
+                m_Lights.erase(it);
+                break;
+            }
+        }
     }
 
     delete SPrim;
@@ -901,19 +910,21 @@ void HnRenderDelegate::CommitResources(pxr::HdChangeTracker* tracker)
             std::lock_guard<std::mutex> Guard{m_LightsMtx};
 
             HnLight* DomeLight = nullptr;
-            for (HnLight* pLight : m_Lights)
+
+            auto dome_lights_range = m_Lights.equal_range(pxr::HdPrimTypeTokens->domeLight);
+            for (auto it = dome_lights_range.first; it != dome_lights_range.second; ++it)
             {
-                if (pLight->GetTypeId() == pxr::HdPrimTypeTokens->domeLight)
+                HnLight* pLight = it->second;
+                VERIFY_EXPR(pLight->GetTypeId() == pxr::HdPrimTypeTokens->domeLight);
+
+                if (DomeLight == nullptr)
                 {
-                    if (DomeLight == nullptr)
-                    {
-                        pLight->PrecomputeIBLCubemaps(*this);
-                        DomeLight = pLight;
-                    }
-                    else
-                    {
-                        LOG_WARNING_MESSAGE("Only one dome light is supported. ", pLight->GetId(), " will be ignored");
-                    }
+                    pLight->PrecomputeIBLCubemaps(*this);
+                    DomeLight = pLight;
+                }
+                else
+                {
+                    LOG_WARNING_MESSAGE("Only one dome light is supported. ", pLight->GetId(), " will be ignored");
                 }
             }
             m_LightResourcesVersion = LightResourcesVersion;

Hydrogent/src/Tasks/HnBeginFrameTask.cpp

@@ -396,8 +396,9 @@ void HnBeginFrameTask::Prepare(pxr::HdTaskContext* TaskCtx,
         const auto&  Lights                 = RenderDelegate->GetLights();
 
         Uint32 ShadowCastingLightIdx = 0;
-        for (HnLight* Light : Lights)
+        for (auto light_it : Lights)
         {
+            HnLight* Light = light_it.second;
             if (Light->ShadowsEnabled() && Light->IsVisible() && ShadowCastingLightIdx < NumShadowCastingLights)
             {
                 Light->SetFrameAttribsIndex(ShadowCastingLightIdx++);
@@ -439,9 +440,10 @@ void HnBeginFrameTask::UpdateFrameConstants(IDeviceContext* pCtx,
         const TextureDesc& ShadowAtlasDesc = ShadowMapMgr->GetAtlasDesc();
 
         const auto& Lights = RenderDelegate->GetLights();
-        for (const HnLight* Light : Lights)
+        for (auto light_it : Lights)
         {
-            const Int32 ShadowCastingLightIdx = Light->GetFrameAttribsIndex();
+            const HnLight* Light                 = light_it.second;
+            const Int32    ShadowCastingLightIdx = Light->GetFrameAttribsIndex();
             VERIFY_EXPR((ShadowCastingLightIdx < 0) == (!Light->ShadowsEnabled() || !Light->IsVisible() || ShadowCastingLightIdx >= static_cast<Int32>(NumShadowCastingLights)));
             if (ShadowCastingLightIdx < 0)
                 continue;
@@ -574,8 +576,9 @@ void HnBeginFrameTask::UpdateFrameConstants(IDeviceContext* pCtx,
         }
 
         int LightCount = 0;
-        for (HnLight* Light : RenderDelegate->GetLights())
+        for (auto light_it : RenderDelegate->GetLights())
         {
+            HnLight* Light = light_it.second;
             if (!Light->IsVisible())
                 continue;
 

Hydrogent/src/Tasks/HnRenderShadowsTask.cpp

@@ -257,8 +257,9 @@ void HnRenderShadowsTask::Prepare(pxr::HdTaskContext* TaskCtx,
     const auto& Lights = RenderDelegate->GetLights();
 
     // Sort all shadow lights with dirty shadow maps by shadow map slice
-    for (HnLight* Light : Lights)
+    for (auto light_it : Lights)
     {
+        HnLight* Light = light_it.second;
         if (!Light->ShadowsEnabled())
             continue;