Leaf/leaf-server/minecraft-patches/features/0159-Replace-brain-with-optimized-collection.patch

From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001
From: HaHaWTH <102713261+HaHaWTH@users.noreply.github.com>
Date: Sat, 26 Oct 2024 00:06:04 +0800
Subject: [PATCH] Replace brain with optimized collection

Co-authored-by: Taiyou06 <kaandindar21@gmail.com>

diff --git a/net/minecraft/world/entity/ai/Brain.java b/net/minecraft/world/entity/ai/Brain.java
index 29fdf94db0308031edfe7915fc587a2aa5a1a18a..1d533dc2db5df6f31f307bc94f9f8e5784b0ab23 100644
--- a/net/minecraft/world/entity/ai/Brain.java
+++ b/net/minecraft/world/entity/ai/Brain.java
@@ -45,14 +45,20 @@ public class Brain<E extends LivingEntity> {
     static final Logger LOGGER = LogUtils.getLogger();
     private final Supplier<Codec<Brain<E>>> codec;
     private static final int SCHEDULE_UPDATE_DELAY = 20;
-    private final Map<MemoryModuleType<?>, Optional<? extends ExpirableValue<?>>> memories = Maps.newHashMap();
-    private final Map<SensorType<? extends Sensor<? super E>>, Sensor<? super E>> sensors = Maps.newLinkedHashMap();
-    private final Map<Integer, Map<Activity, Set<BehaviorControl<? super E>>>> availableBehaviorsByPriority = Maps.newTreeMap();
+    // Leaf start - Replace brain maps with optimized collection
+    private final Map<MemoryModuleType<?>, Optional<? extends ExpirableValue<?>>> memories = new it.unimi.dsi.fastutil.objects.Reference2ReferenceOpenHashMap<>();
+    private final Map<SensorType<? extends Sensor<? super E>>, Sensor<? super E>> sensors = new it.unimi.dsi.fastutil.objects.Reference2ReferenceArrayMap<>();
+    private final Map<Integer, Map<Activity, Set<BehaviorControl<? super E>>>> availableBehaviorsByPriority = new it.unimi.dsi.fastutil.objects.Object2ObjectRBTreeMap<>();
+    public static final Map[] EMPTY_ARRAY = {};
+    private Map<Activity, Set<BehaviorControl<? super E>>>[] availableBehaviorsByPriorityArray = EMPTY_ARRAY;
+    // Leaf end - Replace brain maps with optimized collection
     private Schedule schedule = Schedule.EMPTY;
-    private final Map<Activity, Set<Pair<MemoryModuleType<?>, MemoryStatus>>> activityRequirements = Maps.newHashMap();
-    private final Map<Activity, Set<MemoryModuleType<?>>> activityMemoriesToEraseWhenStopped = Maps.newHashMap();
-    private Set<Activity> coreActivities = Sets.newHashSet();
-    private final Set<Activity> activeActivities = Sets.newHashSet();
+    // Leaf start - Replace brain maps with optimized collection
+    private final Map<Activity, Set<Pair<MemoryModuleType<?>, MemoryStatus>>> activityRequirements = new org.dreeam.leaf.util.map.ActivityArrayMap<>();
+    private final Map<Activity, Set<MemoryModuleType<?>>> activityMemoriesToEraseWhenStopped = new org.dreeam.leaf.util.map.ActivityArrayMap<>();
+    private Set<Activity> coreActivities = new org.dreeam.leaf.util.map.ActivityBitSet();
+    private final Set<Activity> activeActivities = new org.dreeam.leaf.util.map.ActivityBitSet();
+    // Leaf end - Replace brain maps with optimized collection
     private Activity defaultActivity = Activity.IDLE;
     private long lastScheduleUpdate = -9999L;
 
@@ -162,6 +168,8 @@ public class Brain<E extends LivingEntity> {
         for (Brain.MemoryValue<?> memoryValue : memoryValues) {
             memoryValue.setMemoryInternal(this);
         }
+
+        ((it.unimi.dsi.fastutil.objects.Reference2ReferenceOpenHashMap<MemoryModuleType<?>, Optional<? extends ExpirableValue<?>>>) this.memories).trim(); // Leaf
     }
 
     public <T> DataResult<T> serializeStart(DynamicOps<T> ops) {
@@ -178,6 +186,7 @@ public class Brain<E extends LivingEntity> {
 
     public void clearMemories() {
         this.memories.keySet().forEach(memoryModuleType -> this.memories.put((MemoryModuleType<?>)memoryModuleType, Optional.empty()));
+        this.availableBehaviorsByPriorityArray = EMPTY_ARRAY; // Leaf
     }
 
     public <U> void eraseMemory(MemoryModuleType<U> type) {
@@ -362,7 +371,7 @@ public class Brain<E extends LivingEntity> {
     }
 
     public void addActivity(Activity activity, ImmutableList<? extends Pair<Integer, ? extends BehaviorControl<? super E>>> tasks) {
-        this.addActivityAndRemoveMemoriesWhenStopped(activity, tasks, ImmutableSet.of(), Sets.newHashSet());
+        this.addActivityAndRemoveMemoriesWhenStopped(activity, tasks, ImmutableSet.of(), new it.unimi.dsi.fastutil.objects.ReferenceOpenHashSet<>());
     }
 
     public void addActivityWithConditions(
@@ -395,8 +404,8 @@ public class Brain<E extends LivingEntity> {
 
         for (Pair<Integer, ? extends BehaviorControl<? super E>> pair : tasks) {
             this.availableBehaviorsByPriority
-                .computeIfAbsent(pair.getFirst(), integer -> Maps.newHashMap())
-                .computeIfAbsent(activity, activity1 -> Sets.newLinkedHashSet())
+                .computeIfAbsent(pair.getFirst(), integer -> new org.dreeam.leaf.util.map.ActivityArrayMap<>()) // Leaf - Replace brain activity maps with optimized collection
+                .computeIfAbsent(activity, activity1 -> new org.dreeam.leaf.util.map.BehaviorControlArraySet<>()) // Leaf - Replace brain activity maps with optimized collection
                 .add((BehaviorControl<? super E>)pair.getSecond());
         }
     }
@@ -404,6 +413,7 @@ public class Brain<E extends LivingEntity> {
     @VisibleForTesting
     public void removeAllBehaviors() {
         this.availableBehaviorsByPriority.clear();
+        this.availableBehaviorsByPriorityArray = EMPTY_ARRAY; // Leaf
     }
 
     public boolean isActive(Activity activity) {
@@ -437,14 +447,28 @@ public class Brain<E extends LivingEntity> {
     }
 
     private void forgetOutdatedMemories() {
+        boolean flag = false;
+        for (var entry : this.memories.values()) {
+            if (entry.isPresent()) {
+                ExpirableValue<?> expirableValue = entry.get();
+                if (expirableValue.hasExpired()) {
+                    flag = true;
+                }
+                expirableValue.tick();
+            }
+        }
+        if (flag) {
+            eraseOutdatedMemories();
+        }
+    }
+
+    private void eraseOutdatedMemories() {
         for (Entry<MemoryModuleType<?>, Optional<? extends ExpirableValue<?>>> entry : this.memories.entrySet()) {
             if (entry.getValue().isPresent()) {
-                ExpirableValue<?> expirableValue = (ExpirableValue<?>)entry.getValue().get();
+                ExpirableValue<?> expirableValue = entry.getValue().get();
                 if (expirableValue.hasExpired()) {
                     this.eraseMemory(entry.getKey());
                 }
-
-                expirableValue.tick();
             }
         }
     }
@@ -452,34 +476,92 @@ public class Brain<E extends LivingEntity> {
     public void stopAll(ServerLevel level, E owner) {
         long gameTime = owner.level().getGameTime();
 
-        for (BehaviorControl<? super E> behaviorControl : this.getRunningBehaviors()) {
-            behaviorControl.doStop(level, owner, gameTime);
+        // Leaf start
+        if (this.availableBehaviorsByPriorityArray.length != this.availableBehaviorsByPriority.size()) {
+            this.availableBehaviorsByPriorityArray = this.availableBehaviorsByPriority.values().toArray(EMPTY_ARRAY);
         }
+        for (Map<Activity, Set<BehaviorControl<? super E>>> map : availableBehaviorsByPriorityArray) {
+            var map1 = (org.dreeam.leaf.util.map.ActivityArrayMap<Set<BehaviorControl<? super E>>>) map;
+            for (int index = 0; index <= org.dreeam.leaf.util.map.ActivityArrayMap.BITS; index++) {
+                if ((map1.bitset & (1 << index)) != 0) {
+                    var behaviorControls = (org.dreeam.leaf.util.map.BehaviorControlArraySet<? super E>) map1.a[index];
+                    var behaviorControlsRaw = behaviorControls.raw();
+                    var behaviorControlsSize = behaviorControls.size();
+                    for (int i = 0; i < behaviorControlsSize; i++) {
+                        BehaviorControl<? super E> behaviorControl = (BehaviorControl<? super E>) behaviorControlsRaw[i];
+                        if (behaviorControl.getStatus() == Behavior.Status.RUNNING) {
+                            behaviorControl.doStop(level, owner, gameTime);
+                            behaviorControls.dec();
+                        }
+                    }
+                }
+            }
+        }
+        // Leaf end
     }
 
     private void startEachNonRunningBehavior(ServerLevel level, E entity) {
         long gameTime = level.getGameTime();
 
-        for (Map<Activity, Set<BehaviorControl<? super E>>> map : this.availableBehaviorsByPriority.values()) {
-            for (Entry<Activity, Set<BehaviorControl<? super E>>> entry : map.entrySet()) {
-                Activity activity = entry.getKey();
-                if (this.activeActivities.contains(activity)) {
-                    for (BehaviorControl<? super E> behaviorControl : entry.getValue()) {
+        // Leaf start
+        if (this.availableBehaviorsByPriorityArray.length != this.availableBehaviorsByPriority.size()) {
+            this.availableBehaviorsByPriorityArray = this.availableBehaviorsByPriority.values().toArray(EMPTY_ARRAY);
+        }
+        var aact = (org.dreeam.leaf.util.map.ActivityBitSet) this.activeActivities;
+        for (int index = 0; index <= org.dreeam.leaf.util.map.ActivityArrayMap.BITS; index++) {
+            if ((aact.bitset & (1 << index)) != 0) {
+                for (Map<Activity, Set<BehaviorControl<? super E>>> map : availableBehaviorsByPriorityArray) {
+                    var map1 = ((org.dreeam.leaf.util.map.ActivityArrayMap<Set<BehaviorControl<? super E>>>) map);
+                    var ele = map1.a[index];
+                    if (ele == null) {
+                        continue;
+                    }
+                    var behaviorControls = (org.dreeam.leaf.util.map.BehaviorControlArraySet<? super E>) ele;
+                    var behaviorControlsRaw = behaviorControls.raw();
+                    var behaviorControlsSize = behaviorControls.size();
+                    for (int i = 0; i < behaviorControlsSize; i++) {
+                        BehaviorControl<? super E> behaviorControl = (BehaviorControl<? super E>) behaviorControlsRaw[i];
                         if (behaviorControl.getStatus() == Behavior.Status.STOPPED) {
-                            behaviorControl.tryStart(level, entity, gameTime);
+                            if (behaviorControl.tryStart(level, entity, gameTime)) {
+                                behaviorControls.inc();
+                            }
                         }
                     }
                 }
             }
         }
+        // Leaf end
     }
 
     private void tickEachRunningBehavior(ServerLevel level, E entity) {
         long gameTime = level.getGameTime();
 
-        for (BehaviorControl<? super E> behaviorControl : this.getRunningBehaviors()) {
-            behaviorControl.tickOrStop(level, entity, gameTime);
+        // Leaf start
+        if (this.availableBehaviorsByPriorityArray.length != this.availableBehaviorsByPriority.size()) {
+            this.availableBehaviorsByPriorityArray = this.availableBehaviorsByPriority.values().toArray(new Map[0]);
+        }
+        for (Map<Activity, Set<BehaviorControl<? super E>>> map : availableBehaviorsByPriorityArray) {
+            var map1 = (org.dreeam.leaf.util.map.ActivityArrayMap<Set<BehaviorControl<? super E>>>) map;
+            for (int index = 0; index <= org.dreeam.leaf.util.map.ActivityArrayMap.BITS; index++) {
+                if ((map1.bitset & (1 << index)) != 0) {
+                    var behaviorControls = (org.dreeam.leaf.util.map.BehaviorControlArraySet<? super E>) map1.a[index];
+                    if (behaviorControls.running()) {
+                        var behaviorControlsRaw = behaviorControls.raw();
+                        var behaviorControlsSize = behaviorControls.size();
+                        for (int i = 0; i < behaviorControlsSize; i++) {
+                            BehaviorControl<? super E> behaviorControl = (BehaviorControl<? super E>) behaviorControlsRaw[i];
+                            if (behaviorControl.getStatus() == Behavior.Status.RUNNING) {
+                                behaviorControl.tickOrStop(level, entity, gameTime);
+                                if (behaviorControl.getStatus() == Behavior.Status.STOPPED) {
+                                    behaviorControls.dec();
+                                }
+                            }
+                        }
+                    }
+                }
+            }
         }
+        // Leaf end
     }
 
     private boolean activityRequirementsAreMet(Activity activity) {
diff --git a/net/minecraft/world/entity/schedule/Activity.java b/net/minecraft/world/entity/schedule/Activity.java
index 5a143bb6fabba3dc4e2272afb0be636d5722ea22..7753b2b2ab00c08dc15444a1caa08faac1d82f0c 100644
--- a/net/minecraft/world/entity/schedule/Activity.java
+++ b/net/minecraft/world/entity/schedule/Activity.java
@@ -32,10 +32,12 @@ public class Activity {
     public static final Activity DIG = register("dig");
     private final String name;
     private final int hashCode;
+    public final int id; // Leaf
 
-    private Activity(String name) {
+    private Activity(String name, int id) { // Leaf
         this.name = name;
         this.hashCode = name.hashCode();
+        this.id = id; // Leaf
     }
 
     public String getName() {
@@ -43,7 +45,7 @@ public class Activity {
     }
 
     private static Activity register(String key) {
-        return Registry.register(BuiltInRegistries.ACTIVITY, key, new Activity(key));
+        return Registry.register(BuiltInRegistries.ACTIVITY, key, new Activity(key, BuiltInRegistries.ACTIVITY.size())); // Leaf
     }
 
     @Override