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 c0a09b615e9b6c4ec72b8b77a78e7da374d4498b..a550e6bd9dffa9a46f0d7967c9d73ff5cbfbaa0a 100644
--- a/net/minecraft/world/entity/ai/attributes/AttributeMap.java
+++ b/net/minecraft/world/entity/ai/attributes/AttributeMap.java
@@ -20,7 +20,7 @@ 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
     // Leaf start - Multithreaded tracker
     private final Set<AttributeInstance> attributesToSync;
     private final Set<AttributeInstance> attributesToUpdate;
@@ -36,7 +36,7 @@ public class AttributeMap {
     // Leaf end - Multithreaded tracker
     // 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) {
@@ -47,7 +47,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) {
@@ -71,13 +71,24 @@ public class AttributeMap {
 
     @Nullable
     public AttributeInstance getInstance(Holder<Attribute> attribute) {
-        if (org.dreeam.leaf.config.modules.async.MultithreadedTracker.enabled) {
-            synchronized (this) {
-                return this.attributes.computeIfAbsent(attribute, holder -> this.supplier.createInstance(this::onAttributeModified, (Holder<Attribute>) holder));
+        // Leaf start - Optimize AttributeMap
+        AttributeInstance v;
+        if ((v = this.attributes.get(attribute)) == null) {
+            AttributeInstance newValue;
+            if (org.dreeam.leaf.config.modules.async.MultithreadedTracker.enabled) {
+                synchronized (this) {
+                    if ((newValue = this.supplier.createInstance(this::onAttributeModified, attribute)) != null) {
+                        this.attributes.put(attribute, newValue);
+                        return newValue;
+                    }
+                }
+            } else if ((newValue = this.supplier.createInstance(this::onAttributeModified, attribute)) != null) {
+                this.attributes.put(attribute, newValue);
+                return newValue;
             }
-        } else {
-            return this.attributes.computeIfAbsent(attribute, holder -> this.supplier.createInstance(this::onAttributeModified, (Holder<Attribute>) holder));
         }
+        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) {