From cbd1e5c837c0c28a81ddf2e7dc2e2a82bf9a30cc Mon Sep 17 00:00:00 2001
From: Wiktor Kocielski <withaust@gmail.com>
Date: Sun, 25 Jun 2023 23:39:59 +0300
Subject: [PATCH 1/3] Add NetworkReplicationNode::GetObject

---
 .../Engine/Networking/NetworkReplicationHierarchy.cpp | 11 +++++++++++
 .../Engine/Networking/NetworkReplicationHierarchy.h   |  8 ++++++++
 2 files changed, 19 insertions(+)

diff --git a/Source/Engine/Networking/NetworkReplicationHierarchy.cpp b/Source/Engine/Networking/NetworkReplicationHierarchy.cpp
index 32b7978d0..364e4f7dd 100644
--- a/Source/Engine/Networking/NetworkReplicationHierarchy.cpp
+++ b/Source/Engine/Networking/NetworkReplicationHierarchy.cpp
@@ -63,6 +63,17 @@ bool NetworkReplicationNode::RemoveObject(ScriptingObject* obj)
     return !Objects.Remove(obj);
 }
 
+bool NetworkReplicationNode::GetObject(ScriptingObject* obj, NetworkReplicationHierarchyObject& result)
+{
+    const int32 index = Objects.Find(obj);
+    if (index != -1)
+    {
+        result = Objects[index];
+        return true;
+    }
+    return false;
+}
+
 bool NetworkReplicationNode::DirtyObject(ScriptingObject* obj)
 {
     const int32 index = Objects.Find(obj);
diff --git a/Source/Engine/Networking/NetworkReplicationHierarchy.h b/Source/Engine/Networking/NetworkReplicationHierarchy.h
index 36701c15d..1260859a5 100644
--- a/Source/Engine/Networking/NetworkReplicationHierarchy.h
+++ b/Source/Engine/Networking/NetworkReplicationHierarchy.h
@@ -200,6 +200,14 @@ API_CLASS(Abstract, Namespace = "FlaxEngine.Networking") class FLAXENGINE_API Ne
     /// <returns>True on successful removal, otherwise false.</returns>
     API_FUNCTION() virtual bool RemoveObject(ScriptingObject* obj);
 
+    /// <summary>
+    /// Gets object from the hierarchy.
+    /// </summary>
+    /// <param name="obj">The object to get.</param>
+    /// <param name="result">The hierarchy object to retrieve.</param>
+    /// <returns>True on successful retrieval, otherwise false.</returns>
+    API_FUNCTION() bool GetObject(ScriptingObject* obj, NetworkReplicationHierarchyObject& result);
+
     /// <summary>
     /// Force replicates the object during the next update. Resets any internal tracking state to force the synchronization.
     /// </summary>

From 65397621e90138ce9e2b6992cbb1fcd409954321 Mon Sep 17 00:00:00 2001
From: Wiktor Kocielski <withaust@gmail.com>
Date: Mon, 26 Jun 2023 17:12:48 +0300
Subject: [PATCH 2/3] Make GetObject virtual and implement for Grid

---
 .../Networking/NetworkReplicationHierarchy.cpp       | 12 ++++++++++++
 .../Engine/Networking/NetworkReplicationHierarchy.h  |  3 ++-
 2 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/Source/Engine/Networking/NetworkReplicationHierarchy.cpp b/Source/Engine/Networking/NetworkReplicationHierarchy.cpp
index 364e4f7dd..e95a06423 100644
--- a/Source/Engine/Networking/NetworkReplicationHierarchy.cpp
+++ b/Source/Engine/Networking/NetworkReplicationHierarchy.cpp
@@ -186,6 +186,18 @@ bool NetworkReplicationGridNode::RemoveObject(ScriptingObject* obj)
     return false;
 }
 
+bool NetworkReplicationGridNode::GetObject(ScriptingObject* obj, NetworkReplicationHierarchyObject& result)
+{
+    for (const auto& e : _children)
+    {
+        if (e.Value.Node->GetObject(obj, result))
+        {
+            return true;
+        }
+    }
+    return false;
+}
+
 void NetworkReplicationGridNode::Update(NetworkReplicationHierarchyUpdateResult* result)
 {
     CHECK(result);
diff --git a/Source/Engine/Networking/NetworkReplicationHierarchy.h b/Source/Engine/Networking/NetworkReplicationHierarchy.h
index 1260859a5..ae482dfa6 100644
--- a/Source/Engine/Networking/NetworkReplicationHierarchy.h
+++ b/Source/Engine/Networking/NetworkReplicationHierarchy.h
@@ -206,7 +206,7 @@ API_CLASS(Abstract, Namespace = "FlaxEngine.Networking") class FLAXENGINE_API Ne
     /// <param name="obj">The object to get.</param>
     /// <param name="result">The hierarchy object to retrieve.</param>
     /// <returns>True on successful retrieval, otherwise false.</returns>
-    API_FUNCTION() bool GetObject(ScriptingObject* obj, NetworkReplicationHierarchyObject& result);
+    API_FUNCTION() virtual bool GetObject(ScriptingObject* obj, NetworkReplicationHierarchyObject& result);
 
     /// <summary>
     /// Force replicates the object during the next update. Resets any internal tracking state to force the synchronization.
@@ -255,6 +255,7 @@ public:
 
     void AddObject(NetworkReplicationHierarchyObject obj) override;
     bool RemoveObject(ScriptingObject* obj) override;
+    bool GetObject(ScriptingObject* obj, NetworkReplicationHierarchyObject& result) override;
     void Update(NetworkReplicationHierarchyUpdateResult* result) override;
 };
 

From 60137059889fa164c3897288ee3a3fffe5824e8d Mon Sep 17 00:00:00 2001
From: Wiktor Kocielski <withaust@gmail.com>
Date: Mon, 26 Jun 2023 17:32:19 +0300
Subject: [PATCH 3/3] Implement object to cell caching for the grid

---
 .../NetworkReplicationHierarchy.cpp           | 34 ++++++++++++-------
 .../Networking/NetworkReplicationHierarchy.h  |  1 +
 2 files changed, 23 insertions(+), 12 deletions(-)

diff --git a/Source/Engine/Networking/NetworkReplicationHierarchy.cpp b/Source/Engine/Networking/NetworkReplicationHierarchy.cpp
index e95a06423..5097b6a26 100644
--- a/Source/Engine/Networking/NetworkReplicationHierarchy.cpp
+++ b/Source/Engine/Networking/NetworkReplicationHierarchy.cpp
@@ -167,6 +167,7 @@ void NetworkReplicationGridNode::AddObject(NetworkReplicationHierarchyObject obj
         cell->MinCullDistance = obj.CullDistance;
     }
     cell->Node->AddObject(obj);
+    _objectToCell[obj.Object] = coord;
 
     // Cache minimum culling distance for a whole cell to skip it at once
     cell->MinCullDistance = Math::Min(cell->MinCullDistance, obj.CullDistance);
@@ -174,26 +175,35 @@ void NetworkReplicationGridNode::AddObject(NetworkReplicationHierarchyObject obj
 
 bool NetworkReplicationGridNode::RemoveObject(ScriptingObject* obj)
 {
-    for (const auto& e : _children)
+    Int3 coord;
+
+    if (!_objectToCell.TryGet(obj, coord))
     {
-        if (e.Value.Node->RemoveObject(obj))
-        {
-            // TODO: remove empty cells?
-            // TODO: update MinCullDistance for cell?
-            return true;
-        }
+        return false;
+    }
+
+    if (_children[coord].Node->RemoveObject(obj))
+    {
+        _objectToCell.Remove(obj);
+        // TODO: remove empty cells?
+        // TODO: update MinCullDistance for cell?
+        return true;
     }
     return false;
 }
 
 bool NetworkReplicationGridNode::GetObject(ScriptingObject* obj, NetworkReplicationHierarchyObject& result)
 {
-    for (const auto& e : _children)
+    Int3 coord;
+
+    if (!_objectToCell.TryGet(obj, coord))
     {
-        if (e.Value.Node->GetObject(obj, result))
-        {
-            return true;
-        }
+        return false;
+    }
+
+    if (_children[coord].Node->GetObject(obj, result))
+    {
+        return true;
     }
     return false;
 }
diff --git a/Source/Engine/Networking/NetworkReplicationHierarchy.h b/Source/Engine/Networking/NetworkReplicationHierarchy.h
index ae482dfa6..b054de453 100644
--- a/Source/Engine/Networking/NetworkReplicationHierarchy.h
+++ b/Source/Engine/Networking/NetworkReplicationHierarchy.h
@@ -246,6 +246,7 @@ private:
     };
 
     Dictionary<Int3, Cell> _children;
+    Dictionary<ScriptingObject*, Int3> _objectToCell;
 
 public:
     /// <summary>