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] Optimize Brain


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 afbb027021acfbe25d534a84f1750e420bbde6e0..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;
 
@@ -69,17 +73,21 @@ public class Brain<E extends LivingEntity> {
         mutableObject.setValue(
             (new MapCodec<Brain<E>>() {
                     public <T> Stream<T> keys(DynamicOps<T> dynamicOps) {
-                        return memoryModules.stream()
-                            .flatMap(
-                                memoryType -> memoryType.getCodec()
-                                        .map(codec -> BuiltInRegistries.MEMORY_MODULE_TYPE.getKey((MemoryModuleType<?>)memoryType))
-                                        .stream()
-                            )
-                            .map(id -> dynamicOps.createString(id.toString()));
+                        // Leaf start - Optimize Brain
+                        List<String> results = new java.util.ArrayList<>();
+                        for (MemoryModuleType<?> memoryType : memoryModules) {
+                            Optional<?> codecOptional = memoryType.getCodec();
+                            if (codecOptional.isPresent()) {
+                                net.minecraft.resources.ResourceLocation id = BuiltInRegistries.MEMORY_MODULE_TYPE.getKey(memoryType);
+                                results.add(id.toString());
+                            }
+                        }
+                        return results.stream().map(dynamicOps::createString);
+                        // Leaf end - Optimize Brain
                     }
 
                     public <T> DataResult<Brain<E>> decode(DynamicOps<T> dynamicOps, MapLike<T> mapLike) {
-                        MutableObject<DataResult<Builder<Brain.MemoryValue<?>>>> mutableObject = new MutableObject<>(
+                        MutableObject<DataResult<Builder<Brain.MemoryValue<?>>>> mutableObject2 = new MutableObject<>( // Leaf - Decompile fix
                             DataResult.success(ImmutableList.builder())
                         );
                         mapLike.entries()
@@ -91,10 +99,10 @@ public class Brain<E extends LivingEntity> {
                                     DataResult<? extends Brain.MemoryValue<?>> dataResult2 = dataResult.flatMap(
                                         memoryType -> this.captureRead((MemoryModuleType<T>)memoryType, dynamicOps, (T)pair.getSecond())
                                     );
-                                    mutableObject.setValue(mutableObject.getValue().apply2(Builder::add, dataResult2));
+                                    mutableObject2.setValue(mutableObject2.getValue().apply2(Builder::add, dataResult2)); // Leaf - Decompile fix
                                 }
                             );
-                        ImmutableList<Brain.MemoryValue<?>> immutableList = mutableObject.getValue()
+                        ImmutableList<Brain.MemoryValue<?>> immutableList = mutableObject2.getValue() // Leaf - Decompile fix
                             .resultOrPartial(Brain.LOGGER::error)
                             .map(Builder::build)
                             .orElseGet(ImmutableList::of);
@@ -152,7 +160,13 @@ public class Brain<E extends LivingEntity> {
     }
 
     Stream<Brain.MemoryValue<?>> memories() {
-        return this.memories.entrySet().stream().map(entry -> Brain.MemoryValue.createUnchecked(entry.getKey(), entry.getValue()));
+        // Leaf start - Optimize Brain
+        List<Brain.MemoryValue<?>> result = new java.util.ArrayList<>();
+        for (Entry<MemoryModuleType<?>, Optional<? extends ExpirableValue<?>>> entry : this.memories.entrySet()) {
+            result.add(Brain.MemoryValue.createUnchecked(entry.getKey(), entry.getValue()));
+        }
+        return result.stream();
+        // Leaf end - Optimize Brain
     }
 
     public boolean hasMemoryValue(MemoryModuleType<?> type) {
@@ -194,14 +208,14 @@ public class Brain<E extends LivingEntity> {
         if (optional == null) {
             throw new IllegalStateException("Unregistered memory fetched: " + type);
         } else {
-            return optional.map(ExpirableValue::getValue);
+            return (Optional<U>) optional.map(ExpirableValue::getValue); // Leaf - Decompile fix
         }
     }
 
     @Nullable
     public <U> Optional<U> getMemoryInternal(MemoryModuleType<U> type) {
         Optional<? extends ExpirableValue<?>> optional = this.memories.get(type);
-        return optional == null ? null : optional.map(ExpirableValue::getValue);
+        return optional == null ? null : (Optional<U>) optional.map(ExpirableValue::getValue); // Leaf - Decompile fix
     }
 
     public <U> long getTimeUntilExpiry(MemoryModuleType<U> type) {