From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001
From: hayanesuru <hayanesuru@outlook.jp>
Date: Thu, 15 May 2025 21:11:18 +0900
Subject: [PATCH] Optimize AttributeMap


diff --git a/net/minecraft/world/entity/ai/attributes/Attribute.java b/net/minecraft/world/entity/ai/attributes/Attribute.java
index f8419dde44ebc7324e783f8bee42132d5ec973c3..406767c60ec1a324faaf5d3658b161647497f99b 100644
--- a/net/minecraft/world/entity/ai/attributes/Attribute.java
+++ b/net/minecraft/world/entity/ai/attributes/Attribute.java
@@ -16,10 +16,15 @@ public class Attribute {
     private boolean syncable;
     private final String descriptionId;
     private Attribute.Sentiment sentiment = Attribute.Sentiment.POSITIVE;
+    // Leaf start - Optimize AttributeMap
+    public final int uid;
+    private static final java.util.concurrent.atomic.AtomicInteger SIZE = new java.util.concurrent.atomic.AtomicInteger();
+    // Leaf end - Optimize AttributeMap
 
     protected Attribute(String descriptionId, double defaultValue) {
         this.defaultValue = defaultValue;
         this.descriptionId = descriptionId;
+        this.uid = SIZE.getAndAdd(1); // Leaf - Optimize AttributeMap
     }
 
     public double getDefaultValue() {
diff --git a/net/minecraft/world/entity/ai/attributes/AttributeMap.java b/net/minecraft/world/entity/ai/attributes/AttributeMap.java
index 701025715e0aca3c1f920a66f9b3d03ec08eaf02..1ef934b02bd99c0fe997e9ec163457b39dba55f0 100644
--- a/net/minecraft/world/entity/ai/attributes/AttributeMap.java
+++ b/net/minecraft/world/entity/ai/attributes/AttributeMap.java
@@ -15,12 +15,12 @@ import net.minecraft.resources.ResourceLocation;
 
 public class AttributeMap {
     // Gale start - Lithium - replace AI attributes with optimized collections
-    private final Map<Holder<Attribute>, AttributeInstance> attributes = new it.unimi.dsi.fastutil.objects.Reference2ReferenceOpenHashMap<>(0);
+    private final Map<Holder<Attribute>, AttributeInstance> attributes = new org.dreeam.leaf.util.map.AttributeInstanceArrayMap(); // Leaf - Optimize AttributeMap
     private final Set<AttributeInstance> attributesToSync = new it.unimi.dsi.fastutil.objects.ReferenceOpenHashSet<>(0);
     private final Set<AttributeInstance> attributesToUpdate = new it.unimi.dsi.fastutil.objects.ReferenceOpenHashSet<>(0);
     // Gale end - Lithium - replace AI attributes with optimized collections
     private final AttributeSupplier supplier;
-    private final java.util.function.Function<Holder<Attribute>, AttributeInstance> createInstance; // Gale - Airplane - reduce entity allocations
+    //private final java.util.function.Function<Holder<Attribute>, AttributeInstance> createInstance; // Gale - Airplane - reduce entity allocations // Leaf - Optimize AttributeMap
     private final net.minecraft.world.entity.LivingEntity entity; // Purpur - Ridables
 
     public AttributeMap(AttributeSupplier supplier) {
@@ -31,7 +31,7 @@ public class AttributeMap {
         this.entity = entity;
         // Purpur end - Ridables
         this.supplier = defaultAttributes;
-        this.createInstance = holder -> this.supplier.createInstance(this::onAttributeModified, holder); // Gale - Airplane - reduce entity allocations
+        //this.createInstance = holder -> this.supplier.createInstance(this::onAttributeModified, holder); // Gale - Airplane - reduce entity allocations // Leaf - Optimize AttributeMap
     }
 
     private void onAttributeModified(AttributeInstance instance) {
@@ -55,7 +55,17 @@ public class AttributeMap {
 
     @Nullable
     public AttributeInstance getInstance(Holder<Attribute> attribute) {
-        return this.attributes.computeIfAbsent(attribute, this.createInstance); // Gale - Airplane - reduce entity allocations - cache lambda, as for some reason java allocates it anyways
+        // Leaf start - Optimize AttributeMap
+        AttributeInstance v;
+        if ((v = this.attributes.get(attribute)) == null) {
+            AttributeInstance newValue;
+            if ((newValue = this.supplier.createInstance(this::onAttributeModified, attribute)) != null) {
+                this.attributes.put(attribute, newValue);
+                return newValue;
+            }
+        }
+        return v;
+        // Leaf end - Optimize AttributeMap
     }
 
     public boolean hasAttribute(Holder<Attribute> attribute) {
diff --git a/net/minecraft/world/entity/ai/attributes/AttributeSupplier.java b/net/minecraft/world/entity/ai/attributes/AttributeSupplier.java
index 24710041ccbc70e5506d8d89ae34f0141977f209..09341ef6c651150aba223689badbead490162b2b 100644
--- a/net/minecraft/world/entity/ai/attributes/AttributeSupplier.java
+++ b/net/minecraft/world/entity/ai/attributes/AttributeSupplier.java
@@ -11,7 +11,7 @@ public class AttributeSupplier {
     private final Map<Holder<Attribute>, AttributeInstance> instances;
 
     AttributeSupplier(Map<Holder<Attribute>, AttributeInstance> instances) {
-        this.instances = instances;
+        this.instances = new org.dreeam.leaf.util.map.AttributeInstanceArrayMap(instances); // Leaf - Optimize AttributeMap
     }
 
     public AttributeInstance getAttributeInstance(Holder<Attribute> attribute) {