From 025dc974ad1b4efeef0cb11b490e1c2f6bfa4969 Mon Sep 17 00:00:00 2001
From: Jason <11360596+jpenilla@users.noreply.github.com>
Date: Wed, 23 Nov 2022 14:03:54 -0700
Subject: [PATCH] Remove from command collection on unregister (#408)

---
 .../commandframework/CommandManager.java      |  2 +-
 .../cloud/commandframework/CommandTree.java   | 19 ++++++++++++++++---
 2 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/cloud-core/src/main/java/cloud/commandframework/CommandManager.java b/cloud-core/src/main/java/cloud/commandframework/CommandManager.java
index 26760e43..07b2510b 100644
--- a/cloud-core/src/main/java/cloud/commandframework/CommandManager.java
+++ b/cloud-core/src/main/java/cloud/commandframework/CommandManager.java
@@ -545,7 +545,7 @@ public abstract class CommandManager<C> {
         this.commandRegistrationHandler.unregisterRootCommand((StaticArgument<?>) node.getValue());
 
         // We then delete it from the tree.
-        this.commandTree.deleteRecursively(node, true);
+        this.commandTree.deleteRecursively(node, true, this.commands::remove);
 
         // And lastly we re-build the entire tree.
         this.commandTree.verifyAndRegister();
diff --git a/cloud-core/src/main/java/cloud/commandframework/CommandTree.java b/cloud-core/src/main/java/cloud/commandframework/CommandTree.java
index aa881bd6..8b901e49 100644
--- a/cloud-core/src/main/java/cloud/commandframework/CommandTree.java
+++ b/cloud-core/src/main/java/cloud/commandframework/CommandTree.java
@@ -58,6 +58,7 @@ import java.util.Objects;
 import java.util.Optional;
 import java.util.Queue;
 import java.util.Set;
+import java.util.function.Consumer;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 import org.apiguardian.api.API;
@@ -952,15 +953,27 @@ public final class CommandTree<C> {
         return null;
     }
 
-    void deleteRecursively(final @NonNull Node<@Nullable CommandArgument<C, ?>> node, final boolean root) {
+    void deleteRecursively(
+        final @NonNull Node<@Nullable CommandArgument<C, ?>> node,
+        final boolean root,
+        final Consumer<Command<C>> op
+    ) {
         for (final Node<@Nullable CommandArgument<C, ?>> child : new ArrayList<>(node.children)) {
-            this.deleteRecursively(child, false);
+            this.deleteRecursively(child, false, op);
         }
 
+        final @Nullable CommandArgument<C, ?> value = node.getValue();
+        final @Nullable Command<C> owner = value == null ? null : value.getOwningCommand();
+        if (owner != null) {
+            op.accept(owner);
+        }
         this.removeNode(node, root);
     }
 
-    private boolean removeNode(final @NonNull Node<@Nullable CommandArgument<C, ?>> node, final boolean root) {
+    private boolean removeNode(
+        final @NonNull Node<@Nullable CommandArgument<C, ?>> node,
+        final boolean root
+    ) {
         if (root) {
             // root command node - remove it from the root tree
             return this.internalTree.removeChild(node);