diff --git a/CommandHook/CommandHookModSystem.cs b/CommandHook/CommandHookModSystem.cs index c3b4539..8bb5b9f 100644 --- a/CommandHook/CommandHookModSystem.cs +++ b/CommandHook/CommandHookModSystem.cs @@ -31,29 +31,44 @@ public class CommandHookModSystem : ModSystem Instance = null; } - public void Register(string modId, string commandName, CommandRegistration registration) + public void Register( + string modId, + CommandRegistration registration, + params string[] commandNames + ) { - var builder = new Dictionary>(); + var builder = Thaw(); - foreach (var (cmd, mods) in registrations) + foreach (var commandName in commandNames) { - var inner = new Dictionary(); - foreach (var (id, reg) in mods) - inner[id] = reg; - builder[cmd] = inner; + if (!builder.TryGetValue(commandName, out var target)) + { + target = new Dictionary(); + builder[commandName] = target; + } + target[modId] = registration; } - if (!builder.TryGetValue(commandName, out var target)) - { - target = new Dictionary(); - builder[commandName] = target; - } - target[modId] = registration; + registrations = Freeze(builder); + } - registrations = builder.ToFrozenDictionary( - pair => pair.Key, - pair => pair.Value.ToFrozenDictionary() - ); + public void Unregister(string modId, params string[] commandNames) + { + var builder = Thaw(); + + if (commandNames.Length == 0) + { + foreach (var (_, mods) in builder) + mods.Remove(modId); + } + else + { + foreach (var commandName in commandNames) + if (builder.TryGetValue(commandName, out var mods)) + mods.Remove(modId); + } + + registrations = Freeze(builder); } public bool FireBefore(string commandName, ref CommandData data) @@ -71,4 +86,32 @@ public class CommandHookModSystem : ModSystem foreach (var (_, reg) in mods) reg.After?.Invoke(ref data, result); } + + private Dictionary> Thaw() + { + var builder = new Dictionary>(); + + foreach (var (cmd, mods) in registrations) + { + var inner = new Dictionary(); + foreach (var (id, reg) in mods) + inner[id] = reg; + builder[cmd] = inner; + } + + return builder; + } + + private static FrozenDictionary> Freeze( + Dictionary> builder + ) + { + var pruned = new Dictionary>(); + + foreach (var (cmd, mods) in builder) + if (mods.Count > 0) + pruned[cmd] = mods.ToFrozenDictionary(); + + return pruned.ToFrozenDictionary(); + } }