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..1c52a118bad6b4b8abdf2527347ef66888b71f54 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 89f4c5b2d61e27acd48063f9f24ce9ea91898b8b..19f7119256ac50ba8db961f179d3fabb9263944e 100644
--- a/net/minecraft/world/entity/ai/attributes/AttributeMap.java
+++ b/net/minecraft/world/entity/ai/attributes/AttributeMap.java
@@ -20,12 +20,12 @@ import org.slf4j.Logger;
 public class AttributeMap {
     private static final Logger LOGGER = LogUtils.getLogger();
     // 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
     private final net.minecraft.world.entity.LivingEntity entity; // Purpur - Ridables
 
     public AttributeMap(AttributeSupplier supplier) {
@@ -36,7 +36,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
     }
 
     private void onAttributeModified(AttributeInstance instance) {
@@ -60,7 +60,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..608edf272735d356215cf0147e65dcef391e1638 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) {