Improve ClientBlockPredicateArgumentType

src/main/java/net/earthcomputer/clientcommands/command/arguments/ClientBlockPredicateArgumentType.java

@@ -1,21 +1,97 @@
 package net.earthcomputer.clientcommands.command.arguments;
 
+import com.mojang.brigadier.StringReader;
 import com.mojang.brigadier.context.CommandContext;
 import com.mojang.brigadier.exceptions.CommandSyntaxException;
+import com.mojang.brigadier.exceptions.DynamicCommandExceptionType;
+import com.mojang.brigadier.suggestion.Suggestions;
+import com.mojang.brigadier.suggestion.SuggestionsBuilder;
+import net.minecraft.block.Block;
+import net.minecraft.block.BlockState;
 import net.minecraft.block.pattern.CachedBlockPosition;
 import net.minecraft.client.MinecraftClient;
+import net.minecraft.command.argument.BlockArgumentParser;
 import net.minecraft.command.argument.BlockPredicateArgumentType;
 import net.minecraft.server.command.ServerCommandSource;
+import net.minecraft.tag.BlockTags;
+import net.minecraft.tag.Tag;
+import net.minecraft.text.TranslatableText;
+import net.minecraft.util.Identifier;
 
+import java.util.HashMap;
+import java.util.Map;
+import java.util.concurrent.CompletableFuture;
 import java.util.function.Predicate;
 
 public class ClientBlockPredicateArgumentType extends BlockPredicateArgumentType {
+    private static final DynamicCommandExceptionType UNKNOWN_TAG_EXCEPTION = new DynamicCommandExceptionType(arg -> new TranslatableText("arguments.block.tag.unknown", arg));
+
+    private boolean allowNbt = true;
+    private boolean allowTags = true;
 
     private ClientBlockPredicateArgumentType() {}
 
+    public static ClientBlockPredicateArgumentType blockPredicate() {
+        return new ClientBlockPredicateArgumentType();
+    }
+
+    public ClientBlockPredicateArgumentType disallowNbt() {
+        allowNbt = false;
+        return this;
+    }
+
+    public ClientBlockPredicateArgumentType disallowTags() {
+        allowTags = false;
+        return this;
+    }
+
+    @Override
+    public BlockPredicate parse(StringReader stringReader) throws CommandSyntaxException {
+        BlockArgumentParser blockParser = (new BlockArgumentParser(stringReader, allowTags)).parse(allowNbt);
+        BlockPredicate predicate;
+        if (blockParser.getBlockState() != null) {
+            BlockPredicateArgumentType.StatePredicate statePredicate = new BlockPredicateArgumentType.StatePredicate(blockParser.getBlockState(), blockParser.getBlockProperties().keySet(), blockParser.getNbtData());
+            predicate = tagManager -> statePredicate;
+        } else {
+            Identifier tagId = blockParser.getTagId();
+            predicate = tagManager -> {
+                Tag<Block> tag = tagManager.getBlocks().getTag(tagId);
+                if (tag == null) {
+                    throw UNKNOWN_TAG_EXCEPTION.create(String.valueOf(tagId));
+                } else {
+                    return new BlockPredicateArgumentType.TagPredicate(tag, blockParser.getProperties(), blockParser.getNbtData());
+                }
+            };
+        }
+
+        if (blockParser.getNbtData() == null) {
+            // optimization: if there is no NBT data, we can cache the blockstate results
+            return tagManager -> {
+                Predicate<CachedBlockPosition> oldPredicate = predicate.create(tagManager);
+                Map<BlockState, Boolean> cache = new HashMap<>();
+                return pos -> cache.computeIfAbsent(pos.getBlockState(), state -> oldPredicate.test(pos));
+            };
+        }
+
+        return predicate;
+    }
+
+    @Override
+    public <S> CompletableFuture<Suggestions> listSuggestions(CommandContext<S> context, SuggestionsBuilder builder) {
+        StringReader stringReader = new StringReader(builder.getInput());
+        stringReader.setCursor(builder.getStart());
+        BlockArgumentParser blockParser = new BlockArgumentParser(stringReader, allowTags);
+
+        try {
+            blockParser.parse(true);
+        } catch (CommandSyntaxException ignore) {
+        }
+
+        return blockParser.getSuggestions(builder, BlockTags.getTagGroup());
+    }
+
     public static Predicate<CachedBlockPosition> getBlockPredicate(CommandContext<ServerCommandSource> context, String arg) throws CommandSyntaxException {
         //noinspection ConstantConditions
         return context.getArgument(arg, BlockPredicate.class).create(MinecraftClient.getInstance().getNetworkHandler().getTagManager());
     }
-
 }

src/main/resources/clientcommands.aw

@@ -1,3 +1,6 @@
 accessWidener	v1	named
 extendable method net/minecraft/loot/condition/EntityPropertiesLootCondition <init> (Lnet/minecraft/predicate/entity/EntityPredicate;Lnet/minecraft/loot/context/LootContext$EntityTarget;)V
 extendable method net/minecraft/loot/condition/LocationCheckLootCondition <init> (Lnet/minecraft/predicate/entity/LocationPredicate;Lnet/minecraft/util/math/BlockPos;)V
+accessible class net/minecraft/command/argument/BlockPredicateArgumentType$StatePredicate
+accessible class net/minecraft/command/argument/BlockPredicateArgumentType$TagPredicate
+accessible method net/minecraft/command/argument/BlockPredicateArgumentType$TagPredicate <init> (Lnet/minecraft/tag/Tag;Ljava/util/Map;Lnet/minecraft/nbt/CompoundTag;)V