From f92fb4373feb7e6a975537b98a847618f179ccb5 Mon Sep 17 00:00:00 2001
From: fullwall <fullwall25@gmail.com>
Date: Sat, 2 Nov 2024 19:05:36 +0800
Subject: [PATCH] Implement a lot more methods in custom entity registry to fix
 tags

---
 .../nms/v1_10_R1/util/NMSImpl.java            |  5 +-
 .../nms/v1_11_R1/util/NMSImpl.java            |  5 +-
 .../nms/v1_12_R1/util/NMSImpl.java            |  5 +-
 .../nms/v1_13_R2/util/NMSImpl.java            |  5 +-
 .../nms/v1_14_R1/util/NMSImpl.java            |  5 +-
 .../nms/v1_15_R1/util/NMSImpl.java            |  5 +-
 .../nms/v1_16_R3/util/NMSImpl.java            |  5 +-
 .../nms/v1_17_R1/util/NMSImpl.java            |  6 +-
 .../nms/v1_18_R2/util/NMSImpl.java            |  7 +-
 .../nms/v1_19_R3/util/NMSImpl.java            |  5 +-
 .../nms/v1_20_R4/util/NMSImpl.java            |  5 +-
 .../nms/v1_21_R2/entity/AllayController.java  |  1 -
 .../v1_21_R2/util/CustomEntityRegistry.java   | 90 +++++++++++++++++--
 .../nms/v1_21_R2/util/NMSImpl.java            |  5 +-
 .../nms/v1_8_R3/util/NMSImpl.java             |  5 +-
 15 files changed, 110 insertions(+), 49 deletions(-)

diff --git a/v1_10_R1/src/main/java/net/citizensnpcs/nms/v1_10_R1/util/NMSImpl.java b/v1_10_R1/src/main/java/net/citizensnpcs/nms/v1_10_R1/util/NMSImpl.java
index 8f66b3387..76f1dbd07 100644
--- a/v1_10_R1/src/main/java/net/citizensnpcs/nms/v1_10_R1/util/NMSImpl.java
+++ b/v1_10_R1/src/main/java/net/citizensnpcs/nms/v1_10_R1/util/NMSImpl.java
@@ -578,10 +578,9 @@ public class NMSImpl implements NMSBridge {
             return DEFAULT_SPEED;
         EntityLiving handle = NMSImpl.getHandle((LivingEntity) npc.getEntity());
         if (handle == null) {
+            return DEFAULT_SPEED;
         }
-        return DEFAULT_SPEED;
-        // return (float)
-        // handle.getAttributeInstance(GenericAttributes.d).getValue();
+        return (float) handle.getAttributeInstance(GenericAttributes.MOVEMENT_SPEED).getValue();
     }
 
     @Override
diff --git a/v1_11_R1/src/main/java/net/citizensnpcs/nms/v1_11_R1/util/NMSImpl.java b/v1_11_R1/src/main/java/net/citizensnpcs/nms/v1_11_R1/util/NMSImpl.java
index 493e2b836..04e52d801 100644
--- a/v1_11_R1/src/main/java/net/citizensnpcs/nms/v1_11_R1/util/NMSImpl.java
+++ b/v1_11_R1/src/main/java/net/citizensnpcs/nms/v1_11_R1/util/NMSImpl.java
@@ -598,10 +598,9 @@ public class NMSImpl implements NMSBridge {
             return DEFAULT_SPEED;
         EntityLiving handle = NMSImpl.getHandle((LivingEntity) npc.getEntity());
         if (handle == null) {
+            return DEFAULT_SPEED;
         }
-        return DEFAULT_SPEED;
-        // return (float)
-        // handle.getAttributeInstance(GenericAttributes.d).getValue();
+        return (float) handle.getAttributeInstance(GenericAttributes.MOVEMENT_SPEED).getValue();
     }
 
     @Override
diff --git a/v1_12_R1/src/main/java/net/citizensnpcs/nms/v1_12_R1/util/NMSImpl.java b/v1_12_R1/src/main/java/net/citizensnpcs/nms/v1_12_R1/util/NMSImpl.java
index b3b810f93..9f8ff7c4d 100644
--- a/v1_12_R1/src/main/java/net/citizensnpcs/nms/v1_12_R1/util/NMSImpl.java
+++ b/v1_12_R1/src/main/java/net/citizensnpcs/nms/v1_12_R1/util/NMSImpl.java
@@ -599,10 +599,9 @@ public class NMSImpl implements NMSBridge {
             return DEFAULT_SPEED;
         EntityLiving handle = NMSImpl.getHandle((LivingEntity) npc.getEntity());
         if (handle == null) {
+            return DEFAULT_SPEED;
         }
-        return DEFAULT_SPEED;
-        // return (float)
-        // handle.getAttributeInstance(GenericAttributes.d).getValue();
+        return (float) handle.getAttributeInstance(GenericAttributes.MOVEMENT_SPEED).getValue();
     }
 
     @Override
diff --git a/v1_13_R2/src/main/java/net/citizensnpcs/nms/v1_13_R2/util/NMSImpl.java b/v1_13_R2/src/main/java/net/citizensnpcs/nms/v1_13_R2/util/NMSImpl.java
index c8080720a..aeca2faaa 100644
--- a/v1_13_R2/src/main/java/net/citizensnpcs/nms/v1_13_R2/util/NMSImpl.java
+++ b/v1_13_R2/src/main/java/net/citizensnpcs/nms/v1_13_R2/util/NMSImpl.java
@@ -628,10 +628,9 @@ public class NMSImpl implements NMSBridge {
             return DEFAULT_SPEED;
         EntityLiving handle = NMSImpl.getHandle((LivingEntity) npc.getEntity());
         if (handle == null) {
+            return DEFAULT_SPEED;
         }
-        return DEFAULT_SPEED;
-        // return (float)
-        // handle.getAttributeInstance(GenericAttributes.d).getValue();
+        return (float) handle.getAttributeInstance(GenericAttributes.MOVEMENT_SPEED).getValue();
     }
 
     @Override
diff --git a/v1_14_R1/src/main/java/net/citizensnpcs/nms/v1_14_R1/util/NMSImpl.java b/v1_14_R1/src/main/java/net/citizensnpcs/nms/v1_14_R1/util/NMSImpl.java
index ed1894fe1..45299039b 100644
--- a/v1_14_R1/src/main/java/net/citizensnpcs/nms/v1_14_R1/util/NMSImpl.java
+++ b/v1_14_R1/src/main/java/net/citizensnpcs/nms/v1_14_R1/util/NMSImpl.java
@@ -667,10 +667,9 @@ public class NMSImpl implements NMSBridge {
             return DEFAULT_SPEED;
         EntityLiving handle = NMSImpl.getHandle((LivingEntity) npc.getEntity());
         if (handle == null) {
+            return DEFAULT_SPEED;
         }
-        return DEFAULT_SPEED;
-        // return (float)
-        // handle.getAttributeInstance(GenericAttributes.d).getValue();
+        return (float) handle.getAttributeInstance(GenericAttributes.MOVEMENT_SPEED).getValue();
     }
 
     @Override
diff --git a/v1_15_R1/src/main/java/net/citizensnpcs/nms/v1_15_R1/util/NMSImpl.java b/v1_15_R1/src/main/java/net/citizensnpcs/nms/v1_15_R1/util/NMSImpl.java
index e365cd2d7..d82ab53e1 100644
--- a/v1_15_R1/src/main/java/net/citizensnpcs/nms/v1_15_R1/util/NMSImpl.java
+++ b/v1_15_R1/src/main/java/net/citizensnpcs/nms/v1_15_R1/util/NMSImpl.java
@@ -682,10 +682,9 @@ public class NMSImpl implements NMSBridge {
             return DEFAULT_SPEED;
         EntityLiving handle = NMSImpl.getHandle((LivingEntity) npc.getEntity());
         if (handle == null) {
+            return DEFAULT_SPEED;
         }
-        return DEFAULT_SPEED;
-        // return (float)
-        // handle.getAttributeInstance(GenericAttributes.d).getValue();
+        return (float) handle.getAttributeInstance(GenericAttributes.MOVEMENT_SPEED).getValue();
     }
 
     @Override
diff --git a/v1_16_R3/src/main/java/net/citizensnpcs/nms/v1_16_R3/util/NMSImpl.java b/v1_16_R3/src/main/java/net/citizensnpcs/nms/v1_16_R3/util/NMSImpl.java
index 42b6ed8f4..6136b1535 100644
--- a/v1_16_R3/src/main/java/net/citizensnpcs/nms/v1_16_R3/util/NMSImpl.java
+++ b/v1_16_R3/src/main/java/net/citizensnpcs/nms/v1_16_R3/util/NMSImpl.java
@@ -703,10 +703,9 @@ public class NMSImpl implements NMSBridge {
             return DEFAULT_SPEED;
         EntityLiving handle = NMSImpl.getHandle((LivingEntity) npc.getEntity());
         if (handle == null) {
+            return DEFAULT_SPEED;
         }
-        return DEFAULT_SPEED;
-        // return (float)
-        // handle.getAttributeInstance(GenericAttributes.d).getValue();
+        return (float) handle.getAttributeInstance(GenericAttributes.MOVEMENT_SPEED).getValue();
     }
 
     @Override
diff --git a/v1_17_R1/src/main/java/net/citizensnpcs/nms/v1_17_R1/util/NMSImpl.java b/v1_17_R1/src/main/java/net/citizensnpcs/nms/v1_17_R1/util/NMSImpl.java
index 1e49ba17b..519e7cc3a 100644
--- a/v1_17_R1/src/main/java/net/citizensnpcs/nms/v1_17_R1/util/NMSImpl.java
+++ b/v1_17_R1/src/main/java/net/citizensnpcs/nms/v1_17_R1/util/NMSImpl.java
@@ -704,11 +704,11 @@ public class NMSImpl implements NMSBridge {
     public float getSpeedFor(NPC npc) {
         if (!npc.isSpawned() || !(npc.getEntity() instanceof org.bukkit.entity.LivingEntity))
             return DEFAULT_SPEED;
-        LivingEntity handle = NMSImpl.getHandle((org.bukkit.entity.LivingEntity) npc.getEntity());
+        LivingEntity handle = getHandle((org.bukkit.entity.LivingEntity) npc.getEntity());
         if (handle == null) {
+            return DEFAULT_SPEED;
         }
-        return DEFAULT_SPEED;
-        // return (float) handle.getAttribute(Attributes.d).getValue();
+        return (float) handle.getAttributeBaseValue(Attributes.MOVEMENT_SPEED);
     }
 
     @Override
diff --git a/v1_18_R2/src/main/java/net/citizensnpcs/nms/v1_18_R2/util/NMSImpl.java b/v1_18_R2/src/main/java/net/citizensnpcs/nms/v1_18_R2/util/NMSImpl.java
index 3287c7390..77f9ef586 100644
--- a/v1_18_R2/src/main/java/net/citizensnpcs/nms/v1_18_R2/util/NMSImpl.java
+++ b/v1_18_R2/src/main/java/net/citizensnpcs/nms/v1_18_R2/util/NMSImpl.java
@@ -711,12 +711,11 @@ public class NMSImpl implements NMSBridge {
     public float getSpeedFor(NPC npc) {
         if (!npc.isSpawned() || !(npc.getEntity() instanceof org.bukkit.entity.LivingEntity))
             return DEFAULT_SPEED;
-        LivingEntity handle = NMSImpl.getHandle((org.bukkit.entity.LivingEntity) npc.getEntity());
+        LivingEntity handle = getHandle((org.bukkit.entity.LivingEntity) npc.getEntity());
         if (handle == null) {
+            return DEFAULT_SPEED;
         }
-        return DEFAULT_SPEED;
-        // return (float)
-        // handle.getAttribute(Attributes.d).getValue();
+        return (float) handle.getAttributeBaseValue(Attributes.MOVEMENT_SPEED);
     }
 
     @Override
diff --git a/v1_19_R3/src/main/java/net/citizensnpcs/nms/v1_19_R3/util/NMSImpl.java b/v1_19_R3/src/main/java/net/citizensnpcs/nms/v1_19_R3/util/NMSImpl.java
index 8fb373b27..9cecae39f 100644
--- a/v1_19_R3/src/main/java/net/citizensnpcs/nms/v1_19_R3/util/NMSImpl.java
+++ b/v1_19_R3/src/main/java/net/citizensnpcs/nms/v1_19_R3/util/NMSImpl.java
@@ -762,10 +762,9 @@ public class NMSImpl implements NMSBridge {
             return DEFAULT_SPEED;
         LivingEntity handle = getHandle((org.bukkit.entity.LivingEntity) npc.getEntity());
         if (handle == null) {
+            return DEFAULT_SPEED;
         }
-        return DEFAULT_SPEED;
-        // return (float)
-        // handle.getAttribute(Attributes.d).getValue();
+        return (float) handle.getAttributeBaseValue(Attributes.MOVEMENT_SPEED);
     }
 
     @Override
diff --git a/v1_20_R4/src/main/java/net/citizensnpcs/nms/v1_20_R4/util/NMSImpl.java b/v1_20_R4/src/main/java/net/citizensnpcs/nms/v1_20_R4/util/NMSImpl.java
index b98e9e74e..10e44a1e0 100644
--- a/v1_20_R4/src/main/java/net/citizensnpcs/nms/v1_20_R4/util/NMSImpl.java
+++ b/v1_20_R4/src/main/java/net/citizensnpcs/nms/v1_20_R4/util/NMSImpl.java
@@ -791,10 +791,9 @@ public class NMSImpl implements NMSBridge {
             return DEFAULT_SPEED;
         LivingEntity handle = getHandle((org.bukkit.entity.LivingEntity) npc.getEntity());
         if (handle == null) {
+            return DEFAULT_SPEED;
         }
-        return DEFAULT_SPEED;
-        // return (float)
-        // handle.getAttribute(Attributes.d).getValue();
+        return (float) handle.getAttributeBaseValue(Attributes.MOVEMENT_SPEED);
     }
 
     @Override
diff --git a/v1_21_R2/src/main/java/net/citizensnpcs/nms/v1_21_R2/entity/AllayController.java b/v1_21_R2/src/main/java/net/citizensnpcs/nms/v1_21_R2/entity/AllayController.java
index e6f14b49d..7aaae632a 100644
--- a/v1_21_R2/src/main/java/net/citizensnpcs/nms/v1_21_R2/entity/AllayController.java
+++ b/v1_21_R2/src/main/java/net/citizensnpcs/nms/v1_21_R2/entity/AllayController.java
@@ -60,7 +60,6 @@ public class AllayController extends MobEntityController {
 
     public static class EntityAllayNPC extends Allay implements NPCHolder {
         private final CitizensNPC npc;
-
         private int taskId = -1;
 
         public EntityAllayNPC(EntityType<? extends Allay> types, Level level) {
diff --git a/v1_21_R2/src/main/java/net/citizensnpcs/nms/v1_21_R2/util/CustomEntityRegistry.java b/v1_21_R2/src/main/java/net/citizensnpcs/nms/v1_21_R2/util/CustomEntityRegistry.java
index 7f512bafc..6e8382b53 100644
--- a/v1_21_R2/src/main/java/net/citizensnpcs/nms/v1_21_R2/util/CustomEntityRegistry.java
+++ b/v1_21_R2/src/main/java/net/citizensnpcs/nms/v1_21_R2/util/CustomEntityRegistry.java
@@ -8,10 +8,12 @@ import java.util.Map.Entry;
 import java.util.Optional;
 import java.util.Set;
 import java.util.function.Supplier;
+import java.util.stream.Stream;
 
 import com.google.common.collect.BiMap;
 import com.google.common.collect.HashBiMap;
 import com.google.common.collect.Maps;
+import com.mojang.serialization.Codec;
 import com.mojang.serialization.Lifecycle;
 
 import net.citizensnpcs.util.NMS;
@@ -20,12 +22,16 @@ import net.minecraft.core.DefaultedRegistry;
 import net.minecraft.core.Holder;
 import net.minecraft.core.Holder.Reference;
 import net.minecraft.core.HolderGetter;
+import net.minecraft.core.HolderOwner;
+import net.minecraft.core.HolderSet.Named;
 import net.minecraft.core.IdMap;
 import net.minecraft.core.MappedRegistry;
+import net.minecraft.core.RegistrationInfo;
 import net.minecraft.core.Registry;
 import net.minecraft.resources.ResourceKey;
 import net.minecraft.resources.ResourceLocation;
 import net.minecraft.tags.TagKey;
+import net.minecraft.tags.TagLoader.LoadResult;
 import net.minecraft.util.RandomSource;
 import net.minecraft.world.entity.EntityType;
 
@@ -63,22 +69,32 @@ public class CustomEntityRegistry extends DefaultedMappedRegistry<EntityType<?>>
 
     @Override
     public EntityType byId(int var0) {
-        return this.wrapped.byId(var0);
+        return wrapped.byId(var0);
     }
 
     @Override
     public EntityType byIdOrThrow(int var0) {
-        return this.wrapped.byIdOrThrow(var0);
+        return wrapped.byIdOrThrow(var0);
+    }
+
+    @Override
+    public boolean canSerializeIn(HolderOwner<EntityType<?>> owner) {
+        return wrapped.canSerializeIn(owner);
     }
 
     @Override
     public boolean containsKey(ResourceKey<EntityType<?>> var0) {
-        return this.wrapped.containsKey(var0);
+        return wrapped.containsKey(var0);
     }
 
     @Override
     public boolean containsKey(ResourceLocation var0) {
-        return this.wrapped.containsKey(var0);
+        return wrapped.containsKey(var0);
+    }
+
+    @Override
+    public Reference<EntityType<?>> createIntrusiveHolder(EntityType<?> type) {
+        return wrapped.createIntrusiveHolder(type);
     }
 
     @Override
@@ -91,6 +107,11 @@ public class CustomEntityRegistry extends DefaultedMappedRegistry<EntityType<?>>
         return wrapped.entrySet();
     }
 
+    @Override
+    public Registry<EntityType<?>> freeze() {
+        return wrapped.freeze();
+    }
+
     @Override
     public MappedRegistry<EntityType<?>> get() {
         return wrapped;
@@ -109,16 +130,24 @@ public class CustomEntityRegistry extends DefaultedMappedRegistry<EntityType<?>>
     }
 
     @Override
-    public int getId(EntityType key) {
+    public Optional<Named<EntityType<?>>> get(TagKey<EntityType<?>> key) {
+        return wrapped.get(key);
+    }
+
+    @Override
+    public Optional<Reference<EntityType<?>>> getAny() {
+        return wrapped.getAny();
+    }
+
+    @Override
+    public int getId(EntityType<?> key) {
         if (entityIds.containsKey(key))
             return entityIds.get(key);
         return wrapped.getId(key);
     }
 
     @Override
-    public ResourceLocation getKey(EntityType value) {
-        if (entityClasses.containsKey(value))
-            return entityClasses.get(value);
+    public ResourceLocation getKey(EntityType<?> value) {
         return wrapped.getKey(value);
     }
 
@@ -144,6 +173,11 @@ public class CustomEntityRegistry extends DefaultedMappedRegistry<EntityType<?>>
         return wrapped.getResourceKey(var0);
     }
 
+    @Override
+    public Stream<Named<EntityType<?>>> getTags() {
+        return wrapped.getTags();
+    }
+
     @Override
     public EntityType<?> getValue(ResourceKey<EntityType<?>> key) {
         return wrapped.getValue(key);
@@ -156,6 +190,11 @@ public class CustomEntityRegistry extends DefaultedMappedRegistry<EntityType<?>>
         return wrapped.getValue(key);
     }
 
+    @Override
+    public Codec<Holder<EntityType<?>>> holderByNameCodec() {
+        return wrapped.holderByNameCodec();
+    }
+
     @Override
     public boolean isEmpty() {
         return wrapped.isEmpty();
@@ -166,16 +205,41 @@ public class CustomEntityRegistry extends DefaultedMappedRegistry<EntityType<?>>
         return wrapped.iterator();
     }
 
+    @Override
+    public ResourceKey<? extends Registry<EntityType<?>>> key() {
+        return wrapped.key();
+    }
+
     @Override
     public Set<ResourceLocation> keySet() {
         return wrapped.keySet();
     }
 
+    @Override
+    public Stream<Reference<EntityType<?>>> listElements() {
+        return wrapped.listElements();
+    }
+
+    @Override
+    public Stream<Named<EntityType<?>>> listTags() {
+        return wrapped.listTags();
+    }
+
+    @Override
+    public PendingTags<EntityType<?>> prepareTagReload(LoadResult<EntityType<?>> var0) {
+        return wrapped.prepareTagReload(var0);
+    }
+
     public void put(int entityId, ResourceLocation key, EntityType entityClass) {
         entities.put(key, entityClass.builtInRegistryHolder());
         entityIds.put(entityClass, entityId);
     }
 
+    @Override
+    public Optional<RegistrationInfo> registrationInfo(ResourceKey<EntityType<?>> key) {
+        return wrapped.registrationInfo(key);
+    }
+
     @Override
     public Set<ResourceKey<EntityType<?>>> registryKeySet() {
         return wrapped.registryKeySet();
@@ -191,6 +255,16 @@ public class CustomEntityRegistry extends DefaultedMappedRegistry<EntityType<?>>
         return wrapped.size();
     }
 
+    @Override
+    public String toString() {
+        return wrapped.toString();
+    }
+
+    @Override
+    public Holder<EntityType<?>> wrapAsHolder(EntityType<?> type) {
+        return wrapped.wrapAsHolder(type);
+    }
+
     private static final MethodHandle IREGISTRY_LIFECYCLE = NMS.getFirstGetter(MappedRegistry.class, Lifecycle.class);
     private static final MethodHandle IREGISTRY_RESOURCE_KEY = NMS.getFirstGetter(MappedRegistry.class,
             ResourceKey.class);
diff --git a/v1_21_R2/src/main/java/net/citizensnpcs/nms/v1_21_R2/util/NMSImpl.java b/v1_21_R2/src/main/java/net/citizensnpcs/nms/v1_21_R2/util/NMSImpl.java
index 1486d5f25..6fe7d3c4a 100644
--- a/v1_21_R2/src/main/java/net/citizensnpcs/nms/v1_21_R2/util/NMSImpl.java
+++ b/v1_21_R2/src/main/java/net/citizensnpcs/nms/v1_21_R2/util/NMSImpl.java
@@ -779,10 +779,9 @@ public class NMSImpl implements NMSBridge {
             return DEFAULT_SPEED;
         LivingEntity handle = getHandle((org.bukkit.entity.LivingEntity) npc.getEntity());
         if (handle == null) {
+            return DEFAULT_SPEED;
         }
-        return DEFAULT_SPEED;
-        // return (float)
-        // handle.getAttribute(Attributes.d).getValue();
+        return (float) handle.getAttributeBaseValue(Attributes.MOVEMENT_SPEED);
     }
 
     @Override
diff --git a/v1_8_R3/src/main/java/net/citizensnpcs/nms/v1_8_R3/util/NMSImpl.java b/v1_8_R3/src/main/java/net/citizensnpcs/nms/v1_8_R3/util/NMSImpl.java
index 993e1932c..8446aaeb7 100644
--- a/v1_8_R3/src/main/java/net/citizensnpcs/nms/v1_8_R3/util/NMSImpl.java
+++ b/v1_8_R3/src/main/java/net/citizensnpcs/nms/v1_8_R3/util/NMSImpl.java
@@ -533,10 +533,9 @@ public class NMSImpl implements NMSBridge {
             return DEFAULT_SPEED;
         EntityLiving handle = NMSImpl.getHandle((LivingEntity) npc.getEntity());
         if (handle == null) {
+            return DEFAULT_SPEED;
         }
-        return DEFAULT_SPEED;
-        // return (float)
-        // handle.getAttributeInstance(GenericAttributes.d).getValue();
+        return (float) handle.getAttributeInstance(GenericAttributes.MOVEMENT_SPEED).getValue();
     }
 
     @Override