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 maps with optimized collection


diff --git a/src/main/java/net/minecraft/world/entity/ai/Brain.java b/src/main/java/net/minecraft/world/entity/ai/Brain.java
index f6439fb10031a8f17ce3a5bee9f6e89275edebc7..2e0859cd62fd064d43aaf3ac24269d159a6bd49e 100644
--- a/src/main/java/net/minecraft/world/entity/ai/Brain.java
+++ b/src/main/java/net/minecraft/world/entity/ai/Brain.java
@@ -45,14 +45,18 @@ 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 - Optimize Brain
+    private final Map<MemoryModuleType<?>, Optional<? extends ExpirableValue<?>>> memories = new it.unimi.dsi.fastutil.objects.Object2ObjectOpenHashMap<>();
+    private final Map<SensorType<? extends Sensor<? super E>>, Sensor<? super E>> sensors = new it.unimi.dsi.fastutil.objects.Object2ObjectLinkedOpenHashMap<>();
+    private final Map<Integer, Map<Activity, Set<BehaviorControl<? super E>>>> availableBehaviorsByPriority = new it.unimi.dsi.fastutil.objects.Object2ObjectRBTreeMap<>();
+    // Leaf end - Optimize Brain
     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 - Optimize Brain
+    private final Map<Activity, Set<Pair<MemoryModuleType<?>, MemoryStatus>>> activityRequirements = new it.unimi.dsi.fastutil.objects.Object2ObjectOpenHashMap<>();
+    private final Map<Activity, Set<MemoryModuleType<?>>> activityMemoriesToEraseWhenStopped = new it.unimi.dsi.fastutil.objects.Object2ObjectOpenHashMap<>();
+    private Set<Activity> coreActivities = new it.unimi.dsi.fastutil.objects.ObjectOpenHashSet<>();
+    private final Set<Activity> activeActivities = new it.unimi.dsi.fastutil.objects.ObjectOpenHashSet<>();
+    // Leaf end - Optimize Brain
     private Activity defaultActivity = Activity.IDLE;
     private long lastScheduleUpdate = -9999L;