From 0a11e8ffc5241df5a474f47d57fc9368983f4019 Mon Sep 17 00:00:00 2001
From: AMZN-Olex <5432499+AMZN-Olex@users.noreply.github.com>
Date: Tue, 5 Oct 2021 11:26:31 -0400
Subject: [PATCH] Optimization for looking up components.

Signed-off-by: AMZN-Olex <5432499+AMZN-Olex@users.noreply.github.com>
---
 .../NetworkHierarchyChildComponent.cpp        | 15 ++++----
 .../NetworkHierarchyRootComponent.cpp         | 36 +++++++++++++++----
 2 files changed, 36 insertions(+), 15 deletions(-)

diff --git a/Gems/Multiplayer/Code/Source/Components/NetworkHierarchyChildComponent.cpp b/Gems/Multiplayer/Code/Source/Components/NetworkHierarchyChildComponent.cpp
index 5d9b4a40ad..1f32387893 100644
--- a/Gems/Multiplayer/Code/Source/Components/NetworkHierarchyChildComponent.cpp
+++ b/Gems/Multiplayer/Code/Source/Components/NetworkHierarchyChildComponent.cpp
@@ -201,16 +201,13 @@ namespace Multiplayer
         {
             if (const AZ::Entity* childEntity = componentApplication->FindEntity(childEntityId))
             {
-                for (Component* component : childEntity->GetComponents())
+                if (auto* hierarchyChildComponent = childEntity->FindComponent<NetworkHierarchyChildComponent>())
                 {
-                    if (component->GetUnderlyingComponentType() == NetworkHierarchyChildComponent::TYPEINFO_Uuid())
-                    {
-                        static_cast<NetworkHierarchyChildComponent*>(component)->SetTopLevelHierarchyRootEntity(nullptr);
-                    }
-                    else if (component->GetUnderlyingComponentType() == NetworkHierarchyRootComponent::TYPEINFO_Uuid())
-                    {
-                        static_cast<NetworkHierarchyRootComponent*>(component)->SetTopLevelHierarchyRootEntity(nullptr);
-                    }
+                    hierarchyChildComponent->SetTopLevelHierarchyRootEntity(nullptr);
+                }
+                else if (auto* hierarchyRootComponent = childEntity->FindComponent<NetworkHierarchyRootComponent>())
+                {
+                    hierarchyRootComponent->SetTopLevelHierarchyRootEntity(nullptr);
                 }
             }
         }
diff --git a/Gems/Multiplayer/Code/Source/Components/NetworkHierarchyRootComponent.cpp b/Gems/Multiplayer/Code/Source/Components/NetworkHierarchyRootComponent.cpp
index 5011e3f17b..eeb61ef450 100644
--- a/Gems/Multiplayer/Code/Source/Components/NetworkHierarchyRootComponent.cpp
+++ b/Gems/Multiplayer/Code/Source/Components/NetworkHierarchyRootComponent.cpp
@@ -173,6 +173,29 @@ namespace Multiplayer
         }
     }
 
+    static AZStd::tuple<NetworkHierarchyRootComponent*, NetworkHierarchyChildComponent*> GetHierarchyComponents(const AZ::Entity* entity)
+    {
+        NetworkHierarchyChildComponent* childComponent = nullptr;
+        NetworkHierarchyRootComponent* rootComponent = nullptr;
+
+        for (AZ::Component* component : entity->GetComponents())
+        {
+            if (component->GetUnderlyingComponentType() == NetworkHierarchyChildComponent::TYPEINFO_Uuid())
+            {
+                childComponent = static_cast<NetworkHierarchyChildComponent*>(component);
+                break;
+            }
+
+            if (component->GetUnderlyingComponentType() == NetworkHierarchyRootComponent::TYPEINFO_Uuid())
+            {
+                rootComponent = static_cast<NetworkHierarchyRootComponent*>(component);
+                break;
+            }
+        }
+
+        return AZStd::tie(rootComponent, childComponent);
+    }
+
     void NetworkHierarchyRootComponent::OnParentChanged([[maybe_unused]] AZ::EntityId oldParent, AZ::EntityId newParent)
     {
         // If the parent is part of a hierarchy, it will detect this entity as a new child and rebuild hierarchy.
@@ -181,8 +204,8 @@ namespace Multiplayer
 
         if (AZ::Entity* parentEntity = AZ::Interface<AZ::ComponentApplicationRequests>::Get()->FindEntity(newParent))
         {
-            if (parentEntity->FindComponent<NetworkHierarchyRootComponent>() == nullptr &&
-                parentEntity->FindComponent<NetworkHierarchyChildComponent>() == nullptr)
+            auto [rootComponent, childComponent] = GetHierarchyComponents(parentEntity);
+            if (rootComponent == nullptr && childComponent == nullptr)
             {
                 RebuildHierarchy();
             }
@@ -257,8 +280,7 @@ namespace Multiplayer
 
             if (candidate)
             {
-                auto* hierarchyChildComponent = candidate->FindComponent<NetworkHierarchyChildComponent>();
-                auto* hierarchyRootComponent = candidate->FindComponent<NetworkHierarchyRootComponent>();
+                auto [hierarchyRootComponent, hierarchyChildComponent] = GetHierarchyComponents(candidate);
 
                 if ((hierarchyChildComponent && hierarchyChildComponent->IsHierarchyEnabled()) ||
                     (hierarchyRootComponent && hierarchyRootComponent->IsHierarchyEnabled()))
@@ -282,11 +304,13 @@ namespace Multiplayer
 
     void NetworkHierarchyRootComponent::SetRootForEntity(AZ::Entity* root, const AZ::Entity* childEntity)
     {
-        if (auto* hierarchyChildComponent = childEntity->FindComponent<NetworkHierarchyChildComponent>())
+        auto [hierarchyRootComponent, hierarchyChildComponent] = GetHierarchyComponents(childEntity);
+
+        if (hierarchyChildComponent)
         {
             hierarchyChildComponent->SetTopLevelHierarchyRootEntity(root);
         }
-        else if (auto* hierarchyRootComponent = childEntity->FindComponent<NetworkHierarchyRootComponent>())
+        else if (hierarchyRootComponent)
         {
             hierarchyRootComponent->SetTopLevelHierarchyRootEntity(root);
         }