Data Store Improvements

LabApi/Features/Stores/CustomDataStore.cs

@@ -1,5 +1,6 @@
 using System;
 using System.Collections.Generic;
+using System.Diagnostics.CodeAnalysis;
 using System.Linq;
 using LabApi.Features.Wrappers;
 using NorthwoodLib.Pools;
@@ -11,7 +12,7 @@ namespace LabApi.Features.Stores;
 /// </summary>
 public abstract class CustomDataStore
 {
-    private static readonly Dictionary<Type, Dictionary<Player, CustomDataStore>> StoreInstances = new ();
+    private static readonly Dictionary<Type, Dictionary<Player, CustomDataStore>> StoreInstances = new();
 
     /// <summary>
     /// Gets the <see cref="Player"/> that this instance is associated with.
@@ -25,7 +26,6 @@ public abstract class CustomDataStore
     protected CustomDataStore(Player owner)
     {
         Owner = owner;
-        InternalOnInstanceCreated();
     }
 
     /// <summary>
@@ -39,9 +39,6 @@ public static TStore GetOrAdd<TStore>(Player player)
     {
         Type type = typeof(TStore);
 
-        if (!CustomDataStoreManager.IsRegistered<TStore>())
-            CustomDataStoreManager.RegisterStore<TStore>();
-
         if (!StoreInstances.TryGetValue(type, out Dictionary<Player, CustomDataStore>? playerStores))
         {
             playerStores = new Dictionary<Player, CustomDataStore>();
@@ -53,37 +50,105 @@ public static TStore GetOrAdd<TStore>(Player player)
 
         store = (TStore)Activator.CreateInstance(type, player);
         playerStores[player] = store;
+        store.InternalOnInstanceCreated();
 
         return (TStore)store;
     }
 
     /// <summary>
-    /// Called when a new instance of the <see cref="CustomDataStore"/> is created.
+    /// Gets the <see cref="CustomDataStore"/> for the specified <see cref="Player"/> if it exists.
     /// </summary>
-    protected virtual void OnInstanceCreated() { }
+    /// <param name="player">The <see cref="Player"/> to get the <see cref="CustomDataStore"/> for.</param>
+    /// <typeparam name="TStore">The type of the <see cref="CustomDataStore"/>.</typeparam>
+    /// <returns>The <see cref="CustomDataStore"/> for the specified <see cref="Player"/> or null if it doesn't exist.</returns>
+    public static TStore? Get<TStore>(Player player)
+        where TStore : CustomDataStore
+    {
+        Type type = typeof(TStore);
+
+        if (!StoreInstances.TryGetValue(type, out Dictionary<Player, CustomDataStore>? playerStores))
+            return null;
+
+        if (playerStores.TryGetValue(player, out CustomDataStore? store))
+            return (TStore)store;
+
+        return null;
+    }
 
     /// <summary>
-    /// Called when an instance of the <see cref="CustomDataStore"/> is going to be destroyed.
+    /// Tries to get the <see cref="CustomDataStore"/> for the specified <see cref="Player"/>.
     /// </summary>
-    protected virtual void OnInstanceDestroyed() { }
+    /// <param name="player">The <see cref="Player"/> to get the <see cref="CustomDataStore"/> for.</param>
+    /// <param name="store">The <see cref="CustomDataStore"/> for the specified <see cref="Player"/>.</param>
+    /// <typeparam name="TStore">The type of the <see cref="CustomDataStore"/>.</typeparam>
+    /// <returns>Whether the <see cref="CustomDataStore"/> was successfully retrieved.</returns>
+    public static bool TryGet<TStore>(Player player, [NotNullWhen(true)] out TStore? store)
+        where TStore : CustomDataStore
+    {
+        store = Get<TStore>(player);
+        return store != null;
+    }
 
     /// <summary>
     /// Destroys the <see cref="CustomDataStore"/> for the specified <see cref="Player"/>.
     /// </summary>
     /// <param name="player">The <see cref="Player"/> to destroy the <see cref="CustomDataStore"/> for.</param>
     /// <typeparam name="TStore">The type of the <see cref="CustomDataStore"/>.</typeparam>
-    internal static void Destroy<TStore>(Player player)
+    /// <returns>Whether the <see cref="CustomDataStore"/> was successfully destroyed.</returns>
+    public static bool Destroy<TStore>(Player player)
         where TStore : CustomDataStore
     {
         if (!StoreInstances.TryGetValue(typeof(TStore), out Dictionary<Player, CustomDataStore>? playerStores))
-            return;
+            return false;
 
         if (!playerStores.TryGetValue(player, out CustomDataStore? store))
-            return;
+            return false;
 
         store.Destroy();
+        return true;
+    }
+
+    /// <summary>
+    /// Checks if the <see cref="CustomDataStore"/> for the specified <see cref="Player"/> exists.
+    /// </summary>
+    /// <param name="player"> The <see cref="Player"/> to check the <see cref="CustomDataStore"/> for.</param>
+    /// <typeparam name="TStore">The type of the <see cref="CustomDataStore"/></typeparam>
+    /// <returns>True if the <see cref="CustomDataStore"/> exists for the specified <see cref="Player"/>, false if not.</returns>
+    public static bool Exists<TStore>(Player player)
+        where TStore : CustomDataStore
+    {
+        Type type = typeof(TStore);
+
+        if (!StoreInstances.TryGetValue(type, out Dictionary<Player, CustomDataStore>? playerStores))
+            return false;
+
+        return playerStores.ContainsKey(player);
     }
 
+    /// <summary>
+    /// Gets all instances of the <see cref="CustomDataStore"/> for the specified type.
+    /// </summary>
+    /// <typeparam name="TStore">The type of the <see cref="CustomDataStore"/>.</typeparam>
+    /// <returns>An <see cref="IEnumerable{TStore}"/> of all instances of the <see cref="CustomDataStore"/>.</returns>
+    public static IEnumerable<(Player Player, TStore Store)> GetAll<TStore>()
+        where TStore : CustomDataStore
+    {
+        if (!StoreInstances.TryGetValue(typeof(TStore), out Dictionary<Player, CustomDataStore>? playerStores))
+            return Enumerable.Empty<(Player, TStore)>();
+
+        return playerStores.Select(entry => (entry.Key, (TStore)entry.Value));
+    }
+
+    /// <summary>
+    /// Called when a new instance of the <see cref="CustomDataStore"/> is created.
+    /// </summary>
+    protected virtual void OnInstanceCreated() { }
+
+    /// <summary>
+    /// Called when an instance of the <see cref="CustomDataStore"/> is going to be destroyed.
+    /// </summary>
+    protected virtual void OnInstanceDestroyed() { }
+
     /// <summary>
     /// Destroys all instances of the <see cref="CustomDataStore"/> for the specified type.
     /// </summary>
@@ -102,7 +167,7 @@ internal static void DestroyAll<TStore>()
     /// <summary>
     /// Destroys this instance of the <see cref="CustomDataStore"/>.
     /// </summary>
-    internal void Destroy()
+    private void Destroy()
     {
         OnInstanceDestroyed();
         StoreInstances[this.GetType()].Remove(Owner);
@@ -127,10 +192,21 @@ protected CustomDataStore(Player owner)
     {
     }
 
-    /// <summary>
-    /// Gets the <see cref="CustomDataStore"/> for the specified <see cref="Player"/>.
-    /// </summary>
-    /// <param name="player">The <see cref="Player"/> to get the <see cref="CustomDataStore"/> for.</param>
-    /// <returns>The <see cref="CustomDataStore"/> for the specified <see cref="Player"/>.</returns>
-    public static TStore Get(Player player) => GetOrAdd<TStore>(player);
+    /// <inheritdoc cref="CustomDataStore.GetOrAdd{TStore}"/>
+    public static TStore GetOrAdd(Player player) => GetOrAdd<TStore>(player);
+
+    /// <inheritdoc cref="CustomDataStore.Get{TStore}"/>
+    public static TStore? Get(Player player) => Get<TStore>(player);
+
+    /// <inheritdoc cref="CustomDataStore.TryGet{TStore}"/>
+    public static bool TryGet(Player player, [NotNullWhen(true)] out TStore? store) => TryGet<TStore>(player, out store);
+
+    /// <inheritdoc cref="CustomDataStore.Destroy{TStore}"/>
+    public static bool Destroy(Player player) => Destroy<TStore>(player);
+
+    /// <inheritdoc cref="CustomDataStore.GetAll{TStore}"/>
+    public static IEnumerable<(Player Player, TStore Store)> GetAll() => CustomDataStore.GetAll<TStore>();
+
+    /// <inheritdoc cref="CustomDataStore.Exists{TStore}"/>
+    public static bool Exists(Player player) => Exists<TStore>(player);
 }

LabApi/Features/Stores/CustomDataStoreManager.cs

@@ -1,6 +1,5 @@
 using System;
 using System.Collections.Generic;
-using System.Reflection;
 using LabApi.Features.Wrappers;
 
 namespace LabApi.Features.Stores;
@@ -10,79 +9,56 @@ namespace LabApi.Features.Stores;
 /// </summary>
 public static class CustomDataStoreManager
 {
-    private static readonly List<Type> RegisteredStores = new ();
-    private static readonly Dictionary<Type, MethodInfo> GetOrAddMethods = new ();
-    private static readonly Dictionary<Type, MethodInfo> DestroyMethods = new ();
-    private static readonly Dictionary<Type, MethodInfo> DestroyAllMethods = new ();
+    private static readonly Dictionary<Type, StoreHandler> Handlers = new();
 
     /// <summary>
     /// Registers a custom data store.
     /// </summary>
-    /// <typeparam name="T">The type of the custom data store.</typeparam>
+    /// <typeparam name="TStore">The type of the custom data store.</typeparam>
     /// <returns>Whether the store was successfully registered.</returns>
-    public static bool RegisterStore<T>()
-        where T : CustomDataStore
+    /// <remarks>Once registered, the store will be automatically created for all players when they join the server.</remarks>
+    public static bool RegisterStore<TStore>()
+        where TStore : CustomDataStore
     {
-        Type type = typeof(T);
-        if (RegisteredStores.Contains(type)) return false;
-
-        MethodInfo? getOrAddMethod = typeof(CustomDataStore).GetMethod(nameof(CustomDataStore.GetOrAdd), BindingFlags.Static | BindingFlags.NonPublic | BindingFlags.Public);
-        if (getOrAddMethod == null)
+        Type type = typeof(TStore);
+        if (Handlers.ContainsKey(type))
             return false;
 
-        getOrAddMethod = getOrAddMethod.MakeGenericMethod(type);
-        GetOrAddMethods.Add(type, getOrAddMethod);
-
-        MethodInfo? destroyMethod = typeof(CustomDataStore).GetMethod(nameof(CustomDataStore.Destroy), BindingFlags.Static | BindingFlags.NonPublic | BindingFlags.Public);
-        if (destroyMethod == null)
-            return false;
+        StoreHandler handler = new()
+        {
+            AddPlayer = player => CustomDataStore.GetOrAdd<TStore>(player),
+            RemovePlayer = player => CustomDataStore.Destroy<TStore>(player),
+            DestroyAll = CustomDataStore.DestroyAll<TStore>
+        };
 
-        destroyMethod = destroyMethod.MakeGenericMethod(type);
-        DestroyMethods.Add(type, destroyMethod);
-
-        MethodInfo? destroyAllMethod = typeof(CustomDataStore).GetMethod(nameof(CustomDataStore.DestroyAll), BindingFlags.Static | BindingFlags.NonPublic | BindingFlags.Public);
-        if (destroyAllMethod == null)
-            return false;
-
-        destroyAllMethod = destroyAllMethod.MakeGenericMethod(type);
-        DestroyAllMethods.Add(type, destroyAllMethod);
-
-        RegisteredStores.Add(type);
+        Handlers.Add(type, handler);
 
         return true;
     }
 
     /// <summary>
     /// Unregisters a custom data store.
     /// </summary>
-    /// <typeparam name="T">The type of the custom data store.</typeparam>
-    public static void UnregisterStore<T>()
-        where T : CustomDataStore
+    /// <typeparam name="TStore">The type of the custom data store.</typeparam>
+    public static void UnregisterStore<TStore>()
+        where TStore : CustomDataStore
     {
-        Type type = typeof(T);
-
-        if (DestroyAllMethods.TryGetValue(type, out MethodInfo? method))
-            method.Invoke(null, null);
+        Type type = typeof(TStore);
+        if (Handlers.TryGetValue(type, out StoreHandler? handler))
+            handler.DestroyAll();
 
-        DestroyAllMethods.Remove(type);
-        RegisteredStores.Remove(type);
-        GetOrAddMethods.Remove(type);
-        DestroyMethods.Remove(type);
+        Handlers.Remove(type);
     }
 
     internal static void AddPlayer(Player player)
     {
-        foreach (Type? storeType in RegisteredStores)
-            GetOrAddMethods[storeType].Invoke(null, new object[] { player });
+        foreach (StoreHandler? handler in Handlers.Values)
+            handler.AddPlayer(player);
     }
 
-    internal  static void RemovePlayer(Player player)
+    internal static void RemovePlayer(Player player)
     {
-        foreach (Type? storeType in RegisteredStores)
-            DestroyMethods[storeType].Invoke(null, new object[] { player });
+        foreach (StoreHandler? handler in Handlers.Values)
+            handler.RemovePlayer(player);
     }
-
-    internal static bool IsRegistered(Type type) => RegisteredStores.Contains(type);
-
-    internal static bool IsRegistered<T>() => IsRegistered(typeof(T));
 }

LabApi/Features/Stores/StoreHandler.cs

@@ -0,0 +1,13 @@
+using System;
+using LabApi.Features.Wrappers;
+
+namespace LabApi.Features.Stores;
+
+internal class StoreHandler
+{
+    internal required Action<Player> AddPlayer { get; init; }
+
+    internal required Action<Player> RemovePlayer { get; init; }
+
+    internal required Action DestroyAll { get; init; }
+}

LabApi/Features/Wrappers/Players/Player.cs

@@ -1503,6 +1503,17 @@ public TStore GetDataStore<TStore>()
     {
         return CustomDataStore.GetOrAdd<TStore>(this);
     }
+
+    /// <summary>
+    /// Checks if the <see cref="CustomDataStore"/> exists on the player.
+    /// </summary>
+    /// <typeparam name="TStore">The type of the <see cref="CustomDataStore"/></typeparam>
+    /// <returns>True if the <see cref="CustomDataStore"/> exists on the player, false if not.</returns>
+    public bool HasDataStore<TStore>()
+        where TStore : CustomDataStore
+    {
+        return CustomDataStore.Exists<TStore>(this);
+    }
 
     /// <summary>
     /// Handles the creation of a player in the server.