refactor: rework registration system with listener interface and rebuild pattern

CommandHook/CommandHookModSystem.cs

@@ -10,6 +10,9 @@ public class CommandHookModSystem : ModSystem
 {
     internal static CommandHookModSystem? Instance;
 
+    private readonly List<ICommandHookListener> listeners = new();
+    private readonly List<ICommandHookListener> wildcards = new();
+
     private FrozenDictionary<string, FrozenDictionary<string, CommandRegistration>> registrations =
         FrozenDictionary<string, FrozenDictionary<string, CommandRegistration>>.Empty;
 
@@ -17,58 +20,106 @@ public class CommandHookModSystem : ModSystem
 
     public override bool ShouldLoad(EnumAppSide forSide) => forSide == EnumAppSide.Server;
 
+    public override double ExecuteOrder() => 0.0;
+
     public override void StartServerSide(ICoreServerAPI api)
     {
         Instance = this;
         harmony = new Harmony(Mod.Info.ModID);
         harmony.PatchAll();
+
+        api.Event.ServerRunPhase(EnumServerRunPhase.RunGame, Rebuild);
+
         Mod.Logger.Notification("Loaded");
     }
 
     public override void Dispose()
     {
         harmony?.UnpatchAll(Mod.Info.ModID);
+        listeners.Clear();
+        wildcards.Clear();
         Instance = null;
     }
 
-    public void Register(
-        string modId,
-        CommandRegistration registration,
-        params string[] commandNames
-    )
+    public void Register(ICommandHookListener listener)
     {
-        var builder = Thaw();
+        var commands = listener.Commands;
 
-        foreach (var commandName in commandNames)
+        if (commands == null || commands.Count == 0)
         {
-            if (!builder.TryGetValue(commandName, out var target))
-            {
-                target = new Dictionary<string, CommandRegistration>();
-                builder[commandName] = target;
-            }
-            target[modId] = registration;
+            Unregister(listener);
+            return;
         }
 
-        registrations = Freeze(builder);
+        for (int i = 0; i < listeners.Count; i++)
+        {
+            if (listeners[i].ModId != listener.ModId)
+                continue;
+
+            if (CommandListEquals(listeners[i].Commands, commands))
+                return;
+
+            listeners[i] = listener;
+            SyncWildcard(listener, commands);
+            Rebuild();
+            return;
+        }
+
+        listeners.Add(listener);
+        SyncWildcard(listener, commands);
     }
 
-    public void Unregister(string modId, params string[] commandNames)
+    public void Unregister(ICommandHookListener listener)
     {
-        var builder = Thaw();
+        for (int i = 0; i < listeners.Count; i++)
+        {
+            if (listeners[i].ModId != listener.ModId)
+                continue;
 
-        if (commandNames.Length == 0)
-        {
-            foreach (var (_, mods) in builder)
-                mods.Remove(modId);
+            listeners.RemoveAt(i);
+            wildcards.RemoveAll(w => w.ModId == listener.ModId);
+            Rebuild();
+            return;
         }
-        else
+    }
+
+    public void Rebuild()
+    {
+        var builder = new Dictionary<string, Dictionary<string, CommandRegistration>>();
+
+        foreach (var listener in listeners)
         {
-            foreach (var commandName in commandNames)
-                if (builder.TryGetValue(commandName, out var mods))
-                    mods.Remove(modId);
+            var commands = listener.Commands;
+            if (commands == null || commands.Count == 0)
+                continue;
+
+            bool isWildcard = commands.Count == 1 && commands[0] == "*";
+            if (isWildcard)
+                continue;
+
+            foreach (var cmd in commands)
+            {
+                if (!builder.TryGetValue(cmd, out var mods))
+                {
+                    mods = new Dictionary<string, CommandRegistration>();
+                    builder[cmd] = mods;
+                }
+                mods[listener.ModId] = listener.Registration;
+            }
         }
 
-        registrations = Freeze(builder);
+        foreach (var wildcard in wildcards)
+        foreach (var mods in builder.Values)
+            mods[wildcard.ModId] = wildcard.Registration;
+
+        var pruned = new Dictionary<string, FrozenDictionary<string, CommandRegistration>>(
+            builder.Count
+        );
+        foreach (var (cmd, mods) in builder)
+            if (mods.Count > 0)
+                pruned[cmd] = mods.ToFrozenDictionary();
+
+        registrations = pruned.ToFrozenDictionary();
     }
 
     public bool FireBefore(string commandName, ref CommandData data)
@@ -87,31 +138,23 @@ public class CommandHookModSystem : ModSystem
                 reg.After?.Invoke(ref data, result);
     }
 
-    private Dictionary<string, Dictionary<string, CommandRegistration>> Thaw()
+    private void SyncWildcard(ICommandHookListener listener, IReadOnlyList<string> commands)
     {
-        var builder = new Dictionary<string, Dictionary<string, CommandRegistration>>();
-
-        foreach (var (cmd, mods) in registrations)
-        {
-            var inner = new Dictionary<string, CommandRegistration>();
-            foreach (var (id, reg) in mods)
-                inner[id] = reg;
-            builder[cmd] = inner;
-        }
-
-        return builder;
+        bool isWildcard = commands.Count == 1 && commands[0] == "*";
+        wildcards.RemoveAll(w => w.ModId == listener.ModId);
+        if (isWildcard)
+            wildcards.Add(listener);
     }
 
-    private static FrozenDictionary<string, FrozenDictionary<string, CommandRegistration>> Freeze(
-        Dictionary<string, Dictionary<string, CommandRegistration>> builder
-    )
+    private static bool CommandListEquals(IReadOnlyList<string> a, IReadOnlyList<string> b)
     {
-        var pruned = new Dictionary<string, FrozenDictionary<string, CommandRegistration>>();
+        if (a.Count != b.Count)
+            return false;
 
-        foreach (var (cmd, mods) in builder)
-            if (mods.Count > 0)
-                pruned[cmd] = mods.ToFrozenDictionary();
+        for (int i = 0; i < a.Count; i++)
+            if (a[i] != b[i])
+                return false;
 
-        return pruned.ToFrozenDictionary();
+        return true;
     }
 }