From 0c167f38b73ae40bd72b9553043c9a6bc39fe473 Mon Sep 17 00:00:00 2001
From: Wojtek Figat <wojtek@figat.pl>
Date: Wed, 5 Apr 2023 22:43:05 +0200
Subject: [PATCH] Add `TaskGraphSystem.RemoveDependency` and automatically
 cleanup dependencies on system destroy

---
 Source/Engine/Threading/TaskGraph.cpp | 20 ++++++++++++++++++++
 Source/Engine/Threading/TaskGraph.h   |  9 +++++++++
 2 files changed, 29 insertions(+)

diff --git a/Source/Engine/Threading/TaskGraph.cpp b/Source/Engine/Threading/TaskGraph.cpp
index 26a611a8a..92c8f1bc6 100644
--- a/Source/Engine/Threading/TaskGraph.cpp
+++ b/Source/Engine/Threading/TaskGraph.cpp
@@ -18,11 +18,31 @@ TaskGraphSystem::TaskGraphSystem(const SpawnParams& params)
 {
 }
 
+TaskGraphSystem::~TaskGraphSystem()
+{
+    // Cleanup any outstanding dependencies
+    for (auto* e : _reverseDependencies)
+        e->_dependencies.Remove(this);
+}
+
 void TaskGraphSystem::AddDependency(TaskGraphSystem* system)
 {
+    CHECK(system);
+    if (_dependencies.Contains(system))
+        return;
+    system->_reverseDependencies.Add(this);
     _dependencies.Add(system);
 }
 
+void TaskGraphSystem::RemoveDependency(TaskGraphSystem* system)
+{
+    CHECK(system);
+    if (!_dependencies.Contains(system))
+        return;
+    system->_reverseDependencies.Remove(this);
+    _dependencies.Remove(system);
+}
+
 void TaskGraphSystem::PreExecute(TaskGraph* graph)
 {
 }
diff --git a/Source/Engine/Threading/TaskGraph.h b/Source/Engine/Threading/TaskGraph.h
index f6c130f4d..96e7f440e 100644
--- a/Source/Engine/Threading/TaskGraph.h
+++ b/Source/Engine/Threading/TaskGraph.h
@@ -16,6 +16,7 @@ DECLARE_SCRIPTING_TYPE(TaskGraphSystem);
     friend TaskGraph;
 private:
     Array<TaskGraphSystem*, InlinedAllocation<16>> _dependencies;
+    Array<TaskGraphSystem*, InlinedAllocation<16>> _reverseDependencies;
 
 public:
     /// <summary>
@@ -24,12 +25,20 @@ public:
     API_FIELD() int32 Order = 0;
 
 public:
+    ~TaskGraphSystem();
+
     /// <summary>
     /// Adds the dependency on the system execution. Before this system can be executed the given dependant system has to be executed first.
     /// </summary>
     /// <param name="system">The system to depend on.</param>
     API_FUNCTION() void AddDependency(TaskGraphSystem* system);
 
+    /// <summary>
+    /// Removes the dependency on the system execution.
+    /// </summary>
+    /// <param name="system">The system to not depend on anymore.</param>
+    API_FUNCTION() void RemoveDependency(TaskGraphSystem* system);
+
     /// <summary>
     /// Called before executing any systems of the graph. Can be used to initialize data (synchronous).
     /// </summary>