Disable parent sensor in brained adult animals

lithium-mixin-config.md

@@ -73,6 +73,10 @@ Disabling useless brain sensors to avoid useless sensing calculations.
 (default: `true`)  
 Disable the goat item sensor whose memories are never used.  
 
+### `mixin.ai.useless_sensors.parent_animal_sensor`
+(default: `true`)  
+Disable the parent animal sensor when an animal is not a baby. Would differ from vanilla in the case where an adult animal turns back into a baby animal, as the sensor information is refreshed, leading to a less-outdated value in the first second of turning back into a baby animal. However, there is no way to turn an animal back into a baby without reinitializing the brain, creating entirely new sensors.  
+
 ### `mixin.alloc`
 (default: `true`)  
 Patches that reduce memory allocations  

src/main/java/me/jellysquid/mods/lithium/common/ai/brain/SensorHelper.java

@@ -6,10 +6,14 @@
 import net.minecraft.entity.ai.brain.Brain;
 import net.minecraft.entity.ai.brain.sensor.Sensor;
 import net.minecraft.entity.ai.brain.sensor.SensorType;
+import net.minecraft.server.world.ServerWorld;
 
 public class SensorHelper {
 
     public static void disableSensor(LivingEntity brainedEntity, SensorType<?> sensorType) {
+        if (brainedEntity.getWorld().isClient()) {
+            return;
+        }
         Brain<?> brain = brainedEntity.getBrain();
         Sensor<?> sensor = ((BrainAccessor<?>) brain).getSensors().get(sensorType);
         if (sensor instanceof SensorAccessor sensorAccessor) {
@@ -28,16 +32,29 @@ public static void disableSensor(LivingEntity brainedEntity, SensorType<?> senso
         }
     }
 
-    public static void enableSensor(LivingEntity brainedEntity, SensorType<?> sensorType) {
+    public static <T extends LivingEntity, U extends Sensor<T>> void enableSensor(T brainedEntity, SensorType<U> sensorType) {
+        enableSensor(brainedEntity, sensorType, false);
+    }
+
+    public static <T extends LivingEntity, U extends Sensor<T>> void enableSensor(T brainedEntity, SensorType<U> sensorType, boolean extraTick) {
+        if (brainedEntity.getWorld().isClient()) {
+            return;
+        }
+
         Brain<?> brain = brainedEntity.getBrain();
-        Sensor<?> sensor = ((BrainAccessor<?>) brain).getSensors().get(sensorType);
+        //noinspection unchecked
+        U sensor = (U) ((BrainAccessor<?>) brain).getSensors().get(sensorType);
         if (sensor instanceof SensorAccessor sensorAccessor) {
             long lastSenseTime = sensorAccessor.getLastSenseTime();
             int senseInterval = sensorAccessor.getSenseInterval();
 
             //Recover the random offset of the sensor:
             if (lastSenseTime > senseInterval) {
                 lastSenseTime = lastSenseTime % senseInterval;
+                if (extraTick) {
+                    ((SensorAccessor) sensor).setLastSenseTime(0L);
+                    sensor.tick((ServerWorld) brainedEntity.getWorld(), brainedEntity);
+                }
             }
             sensorAccessor.setLastSenseTime(lastSenseTime);
         }

src/main/java/me/jellysquid/mods/lithium/mixin/ai/useless_sensors/parent_animal_sensor/PassiveEntityMixin.java

@@ -0,0 +1,47 @@
+package me.jellysquid.mods.lithium.mixin.ai.useless_sensors.parent_animal_sensor;
+
+import me.jellysquid.mods.lithium.common.ai.brain.SensorHelper;
+import net.minecraft.entity.EntityType;
+import net.minecraft.entity.LivingEntity;
+import net.minecraft.entity.ai.brain.MemoryModuleType;
+import net.minecraft.entity.ai.brain.sensor.SensorType;
+import net.minecraft.entity.data.TrackedData;
+import net.minecraft.entity.passive.PassiveEntity;
+import net.minecraft.world.World;
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.Shadow;
+import org.spongepowered.asm.mixin.injection.At;
+import org.spongepowered.asm.mixin.injection.Inject;
+import org.spongepowered.asm.mixin.injection.callback.CallbackInfo;
+
+import java.util.Optional;
+
+@Mixin(PassiveEntity.class)
+public abstract class PassiveEntityMixin extends LivingEntity {
+
+    @Shadow
+    public abstract boolean isBaby();
+
+    protected PassiveEntityMixin(EntityType<? extends LivingEntity> entityType, World world) {
+        super(entityType, world);
+    }
+
+    @Inject(
+            method = "onTrackedDataSet",
+            at = @At(value = "INVOKE", target = "Lnet/minecraft/entity/passive/PassiveEntity;calculateDimensions()V"),
+            require = 1, allow = 1
+    )
+    private void handleParentSensor(TrackedData<?> data, CallbackInfo ci) {
+        if (this.getWorld().isClient()) {
+            return;
+        }
+        if (isBaby()) {
+            SensorHelper.enableSensor((PassiveEntity) (Object) this, SensorType.NEAREST_ADULT, true);
+        } else {
+            SensorHelper.disableSensor((PassiveEntity) (Object) this, SensorType.NEAREST_ADULT);
+            if (this.getBrain().hasMemoryModule(MemoryModuleType.NEAREST_VISIBLE_ADULT)) {
+                this.getBrain().remember(MemoryModuleType.NEAREST_VISIBLE_ADULT, Optional.empty());
+            }
+        }
+    }
+}

src/main/java/me/jellysquid/mods/lithium/mixin/ai/useless_sensors/parent_animal_sensor/package-info.java

@@ -0,0 +1,8 @@
+@MixinConfigOption(description = "Disable the parent animal sensor when an animal is not a baby." +
+        " Would differ from vanilla in the case where an adult animal turns back into a baby animal, as the sensor information is refreshed," +
+        " leading to a less-outdated value in the first second of turning back into a baby animal. However, there is no way to turn an animal" +
+        " back into a baby without reinitializing the brain, creating entirely new sensors."
+)
+package me.jellysquid.mods.lithium.mixin.ai.useless_sensors.parent_animal_sensor;
+
+import net.caffeinemc.gradle.MixinConfigOption;

src/main/resources/lithium.mixins.json

@@ -34,6 +34,7 @@
     "ai.useless_sensors.BrainAccessor",
     "ai.useless_sensors.SensorAccessor",
     "ai.useless_sensors.goat_item_sensor.GoatEntityMixin",
+    "ai.useless_sensors.parent_animal_sensor.PassiveEntityMixin",
     "alloc.blockstate.StateMixin",
     "alloc.chunk_random.ServerWorldMixin",
     "alloc.chunk_random.WorldMixin",