add registry support for commands

Author: crosby-moe

src/main/java/meteordevelopment/meteorclient/commands/Command.java

@@ -23,7 +23,7 @@
 import java.util.List;
 
 public abstract class Command {
-    protected static final CommandRegistryAccess REGISTRY_ACCESS = CommandManager.createRegistryAccess(BuiltinRegistries.createWrapperLookup());
+    protected static CommandRegistryAccess REGISTRY_ACCESS = CommandManager.createRegistryAccess(BuiltinRegistries.createWrapperLookup());
     protected static final int SINGLE_SUCCESS = com.mojang.brigadier.Command.SINGLE_SUCCESS;
     protected static final MinecraftClient mc = MeteorClient.mc;
 

src/main/java/meteordevelopment/meteorclient/commands/Commands.java

@@ -7,9 +7,14 @@
 
 import com.mojang.brigadier.CommandDispatcher;
 import com.mojang.brigadier.exceptions.CommandSyntaxException;
+import meteordevelopment.meteorclient.MeteorClient;
 import meteordevelopment.meteorclient.commands.commands.*;
+import meteordevelopment.meteorclient.events.game.GameJoinedEvent;
 import meteordevelopment.meteorclient.pathing.PathManagers;
 import meteordevelopment.meteorclient.utils.PostInit;
+import meteordevelopment.orbit.EventHandler;
+import net.minecraft.client.network.ClientPlayNetworkHandler;
+import net.minecraft.command.CommandRegistryAccess;
 import net.minecraft.command.CommandSource;
 
 import java.util.ArrayList;
@@ -19,8 +24,8 @@
 import static meteordevelopment.meteorclient.MeteorClient.mc;
 
 public class Commands {
-    public static final CommandDispatcher<CommandSource> DISPATCHER = new CommandDispatcher<>();
     public static final List<Command> COMMANDS = new ArrayList<>();
+    public static CommandDispatcher<CommandSource> DISPATCHER = new CommandDispatcher<>();
 
     @PostInit(dependencies = PathManagers.class)
     public static void init() {
@@ -64,11 +69,12 @@ public static void init() {
         add(new LocateCommand());
 
         COMMANDS.sort(Comparator.comparing(Command::getName));
+
+        MeteorClient.EVENT_BUS.subscribe(Commands.class);
     }
 
     public static void add(Command command) {
         COMMANDS.removeIf(existing -> existing.getName().equals(command.getName()));
-        command.registerTo(DISPATCHER);
         COMMANDS.add(command);
     }
 
@@ -85,4 +91,32 @@ public static Command get(String name) {
 
         return null;
     }
+
+    /**
+     * Argument types that rely on Minecraft registries access those registries through a {@link CommandRegistryAccess}
+     * object. Since dynamic registries are specific to each server, we need to make a new CommandRegistryAccess object
+     * every time we join a server.
+     * <p>
+     * The command tree and by extension the {@link CommandDispatcher} also have to be rebuilt because:
+     * <ol>
+     * <li>Argument types that require registries use a registry wrapper object that is created and stored in the
+     *     argument type objects when the command tree is built.
+     * <li>Registry entries and keys are compared using referential equality. Even if the data encoded is the same,
+     *     registry wrapper objects' dynamic data becomes stale after joining another server.
+     * <li>The CommandDispatcher's node merging only adds missing children, it cannot replace stale argument type
+     *     objects.
+     * </ol>
+     *
+     * @author Crosby
+     */
+    @EventHandler
+    private static void onJoin(GameJoinedEvent event) {
+        ClientPlayNetworkHandler networkHandler = mc.getNetworkHandler();
+        Command.REGISTRY_ACCESS = CommandRegistryAccess.of(networkHandler.getRegistryManager(), networkHandler.getEnabledFeatures());
+
+        DISPATCHER = new CommandDispatcher<>();
+        for (Command command : COMMANDS) {
+            command.registerTo(DISPATCHER);
+        }
+    }
 }