From d8edb14b617d1a1e72a3397d59697c8b45d16323 Mon Sep 17 00:00:00 2001
From: WakelessSloth56 <wsloth56@foxmail.com>
Date: Thu, 6 Oct 2022 11:17:57 +0800
Subject: [PATCH] feat(game.datagen): IHTagsProvider

---
 .../game/datagen/tag/IHTagsProvider.java      | 27 +++++++++++++
 .../provider/HEntityTypeTagsProvider.java     | 40 +++++++++----------
 2 files changed, 47 insertions(+), 20 deletions(-)
 create mode 100644 src/main/java/org/auioc/mcmod/arnicalib/game/datagen/tag/IHTagsProvider.java

diff --git a/src/main/java/org/auioc/mcmod/arnicalib/game/datagen/tag/IHTagsProvider.java b/src/main/java/org/auioc/mcmod/arnicalib/game/datagen/tag/IHTagsProvider.java
new file mode 100644
index 00000000..7ced5d46
--- /dev/null
+++ b/src/main/java/org/auioc/mcmod/arnicalib/game/datagen/tag/IHTagsProvider.java
@@ -0,0 +1,27 @@
+package org.auioc.mcmod.arnicalib.game.datagen.tag;
+
+import java.util.function.Predicate;
+import javax.annotation.Nullable;
+import net.minecraft.data.tags.TagsProvider.TagAppender;
+import net.minecraftforge.registries.IForgeRegistry;
+import net.minecraftforge.registries.IForgeRegistryEntry;
+
+public interface IHTagsProvider<T extends IForgeRegistryEntry<T>> {
+
+    @Nullable
+    default IForgeRegistry<T> getRegistry() {
+        return null;
+    }
+
+    default TagAppender<T> addFromRegistry(TagAppender<T> appender, Predicate<T> predicate) {
+        var registry = this.getRegistry();
+        if (registry == null) throw new UnsupportedOperationException("'addFromRegistry' method should not be called unless a registry is specified");
+        for (var value : registry.getValues()) {
+            if (predicate.test(value)) {
+                appender.add(value);
+            }
+        }
+        return appender;
+    }
+
+}
diff --git a/src/main/java/org/auioc/mcmod/arnicalib/mod/datagen/provider/HEntityTypeTagsProvider.java b/src/main/java/org/auioc/mcmod/arnicalib/mod/datagen/provider/HEntityTypeTagsProvider.java
index b8a45c99..44a8815c 100644
--- a/src/main/java/org/auioc/mcmod/arnicalib/mod/datagen/provider/HEntityTypeTagsProvider.java
+++ b/src/main/java/org/auioc/mcmod/arnicalib/mod/datagen/provider/HEntityTypeTagsProvider.java
@@ -1,8 +1,9 @@
 package org.auioc.mcmod.arnicalib.mod.datagen.provider;
 
 import static org.auioc.mcmod.arnicalib.game.tag.HEntityTypeTags.*;
-import java.util.function.Predicate;
+import javax.annotation.Nonnull;
 import org.auioc.mcmod.arnicalib.ArnicaLib;
+import org.auioc.mcmod.arnicalib.game.datagen.tag.IHTagsProvider;
 import org.auioc.mcmod.arnicalib.game.entity.EntityTypePredicates;
 import org.jetbrains.annotations.Nullable;
 import net.minecraft.data.DataGenerator;
@@ -10,8 +11,9 @@
 import net.minecraft.world.entity.EntityType;
 import net.minecraftforge.common.data.ExistingFileHelper;
 import net.minecraftforge.registries.ForgeRegistries;
+import net.minecraftforge.registries.IForgeRegistry;
 
-public class HEntityTypeTagsProvider extends EntityTypeTagsProvider {
+public class HEntityTypeTagsProvider extends EntityTypeTagsProvider implements IHTagsProvider<EntityType<?>> {
 
     public HEntityTypeTagsProvider(DataGenerator generator, @Nullable ExistingFileHelper existingFileHelper) {
         super(generator, ArnicaLib.MOD_ID, existingFileHelper);
@@ -22,6 +24,12 @@ public String getName() {
         return "HEntityTypeTags";
     }
 
+    @Nonnull
+    @Override
+    public IForgeRegistry<EntityType<?>> getRegistry() {
+        return ForgeRegistries.ENTITIES;
+    }
+
     @Override
     protected void addTags() {
         tag(UNDEFINED_MOBS);
@@ -47,25 +55,17 @@ protected void addTags() {
             EntityType.SQUID, EntityType.TROPICAL_FISH, EntityType.TURTLE
         );
 
-        add(tag(MISC_ENTITIES), EntityTypePredicates.IS_MISC);
-        add(tag(MONSTERS), EntityTypePredicates.IS_MONSTER);
-        add(tag(CREATURES), EntityTypePredicates.IS_CREATURE);
-        add(tag(AXOLOTLS), EntityTypePredicates.IS_AXOLOTLS);
-        add(tag(UNDERGROUND_WATER_CREATURES), EntityTypePredicates.IS_UNDERGROUND_WATER_CREATURE);
-        add(tag(WATER_CREATURES), EntityTypePredicates.IS_WATER_CREATURE);
-        add(tag(WATER_AMBIENT_ENTITIES), EntityTypePredicates.IS_WATER_AMBIENT);
-        add(tag(AMBIENT_ENTITIES), EntityTypePredicates.IS_AMBIENT);
-
-        add(tag(FRIENDLY_ENTITIES), EntityTypePredicates.IS_FRIENDLY);
-        add(tag(PERSISTENT_ENTITIES), EntityTypePredicates.IS_PERSISTENT);
-    }
+        addFromRegistry(tag(MISC_ENTITIES), EntityTypePredicates.IS_MISC);
+        addFromRegistry(tag(MONSTERS), EntityTypePredicates.IS_MONSTER);
+        addFromRegistry(tag(CREATURES), EntityTypePredicates.IS_CREATURE);
+        addFromRegistry(tag(AXOLOTLS), EntityTypePredicates.IS_AXOLOTLS);
+        addFromRegistry(tag(UNDERGROUND_WATER_CREATURES), EntityTypePredicates.IS_UNDERGROUND_WATER_CREATURE);
+        addFromRegistry(tag(WATER_CREATURES), EntityTypePredicates.IS_WATER_CREATURE);
+        addFromRegistry(tag(WATER_AMBIENT_ENTITIES), EntityTypePredicates.IS_WATER_AMBIENT);
+        addFromRegistry(tag(AMBIENT_ENTITIES), EntityTypePredicates.IS_AMBIENT);
 
-    private static void add(TagAppender<EntityType<?>> appender, Predicate<EntityType<?>> predicate) {
-        for (var entityType : ForgeRegistries.ENTITIES.getValues()) {
-            if (predicate.test(entityType)) {
-                appender.add(entityType);
-            }
-        }
+        addFromRegistry(tag(FRIENDLY_ENTITIES), EntityTypePredicates.IS_FRIENDLY);
+        addFromRegistry(tag(PERSISTENT_ENTITIES), EntityTypePredicates.IS_PERSISTENT);
     }
 
 }