From e830c34aa7774005c41667d3d13233220a9dafa4 Mon Sep 17 00:00:00 2001
From: Eideren <contact@eideren.com>
Date: Tue, 21 Jan 2025 21:27:24 +0100
Subject: [PATCH] Core: ServiceRegistry, GetOrCreate<T> and Remove(instance)

---
 sources/core/Stride.Core/IService.cs         | 11 ++++++
 sources/core/Stride.Core/IServiceRegistry.cs | 16 ++++++++
 sources/core/Stride.Core/ServiceRegistry.cs  | 39 ++++++++++++++++++--
 3 files changed, 62 insertions(+), 4 deletions(-)
 create mode 100644 sources/core/Stride.Core/IService.cs

diff --git a/sources/core/Stride.Core/IService.cs b/sources/core/Stride.Core/IService.cs
new file mode 100644
index 0000000000..e57ee1b0b9
--- /dev/null
+++ b/sources/core/Stride.Core/IService.cs
@@ -0,0 +1,11 @@
+// Copyright (c) .NET Foundation and Contributors (https://dotnetfoundation.org/ & https://stride3d.net)
+// Distributed under the MIT license. See the LICENSE.md file in the project root for more information.
+
+using Stride.Core.Annotations;
+
+namespace Stride.Core;
+
+public interface IService
+{
+    [NotNull] public static abstract IService NewInstance([NotNull] IServiceRegistry services);
+}
diff --git a/sources/core/Stride.Core/IServiceRegistry.cs b/sources/core/Stride.Core/IServiceRegistry.cs
index 06340c94b0..30f1decf47 100644
--- a/sources/core/Stride.Core/IServiceRegistry.cs
+++ b/sources/core/Stride.Core/IServiceRegistry.cs
@@ -52,6 +52,7 @@ public interface IServiceRegistry
         /// <summary>
         /// Gets the service object of the specified type.
         /// </summary>
+        /// <remarks>The generic type provided must match the generic type of your initial call to <see cref="AddService{T}"/></remarks>
         /// <typeparam name="T">The type of the service to retrieve.</typeparam>
         /// <returns>A service of the requested type, or [null] if not found.</returns>
         [CanBeNull]
@@ -60,7 +61,22 @@ public interface IServiceRegistry
         /// <summary>
         /// Removes the object providing a specified service.
         /// </summary>
+        /// <remarks>The generic type provided must match the generic type of your initial call to <see cref="AddService{T}"/></remarks>
         /// <typeparam name="T">The type of the service to remove.</typeparam>
         void RemoveService<T>() where T : class;
+
+        /// <summary>
+        /// Removes the following object from services if it was registered as one.
+        /// </summary>
+        /// <remarks>The generic type provided must match the generic type of your initial call to <see cref="AddService{T}"/></remarks>
+        /// <returns>True if the argument was a service, false otherwise</returns>
+        /// <typeparam name="T">The type of the service to remove.</typeparam>
+        bool RemoveService<T>([NotNull] T serviceObject) where T : class;
+
+        /// <summary>
+        /// Gets the service object of the specified type, create one if it didn't exist before.
+        /// </summary>
+        /// <typeparam name="T">The type of the service to retrieve.</typeparam>
+        [NotNull] T GetOrCreate<T>() where T : class, IService;
     }
 }
diff --git a/sources/core/Stride.Core/ServiceRegistry.cs b/sources/core/Stride.Core/ServiceRegistry.cs
index 4cde923327..8ce8348477 100644
--- a/sources/core/Stride.Core/ServiceRegistry.cs
+++ b/sources/core/Stride.Core/ServiceRegistry.cs
@@ -72,9 +72,8 @@ public void AddService<T>(T service)
             var type = typeof(T);
             lock (registeredService)
             {
-                if (registeredService.ContainsKey(type))
+                if (registeredService.TryAdd(type, service) == false)
                     throw new ArgumentException("Service is already registered with this type", nameof(type));
-                registeredService.Add(type, service);
             }
             OnServiceAdded(new ServiceEventArgs(type, service));
         }
@@ -91,13 +90,45 @@ public void RemoveService<T>()
             object oldService;
             lock (registeredService)
             {
-                if (registeredService.TryGetValue(type, out oldService))
-                    registeredService.Remove(type);
+                registeredService.Remove(type, out oldService);
             }
             if (oldService != null)
                 OnServiceRemoved(new ServiceEventArgs(type, oldService));
         }
 
+        /// <inheritdoc />
+        public bool RemoveService<T>(T serviceObject) where T : class
+        {
+            lock (registeredService)
+            {
+                if (ReferenceEquals(GetService<T>(), serviceObject))
+                {
+                    RemoveService<T>();
+                    return true;
+                }
+                else
+                {
+                    return false;
+                }
+            }
+        }
+
+        /// <inheritdoc />
+        public T GetOrCreate<T>() where T : class, IService
+        {
+            lock (registeredService)
+            {
+                var t = GetService<T>();
+                if (t is null)
+                {
+                    t = (T)T.NewInstance(this);
+                    AddService(t);
+                }
+
+                return t;
+            }
+        }
+
         private void OnServiceAdded(ServiceEventArgs e)
         {
             ServiceAdded?.Invoke(this, e);