part of an AI optimisation pass

this shaves an average of about 10ns off of the required memory checks for wild pokemobs, from 85ns to 75ns average check time.

src/main/java/pokecube/core/ai/logic/LogicMiscUpdate.java

@@ -6,8 +6,7 @@
 import java.util.Random;
 import java.util.UUID;
 
-import com.google.common.collect.Maps;
-
+import it.unimi.dsi.fastutil.objects.Object2FloatOpenHashMap;
 import net.minecraft.core.BlockPos;
 import net.minecraft.core.particles.ParticleTypes;
 import net.minecraft.network.chat.Component;
@@ -334,7 +333,7 @@ public void tick(final Level world)
         this.prevID = uuid;
 
         // Here we apply worn/held equipment modifiers
-        final Map<Stats, Float> vals = Maps.newHashMap();
+        final Map<Stats, Float> vals = new Object2FloatOpenHashMap<IPokemob.Stats>();
         for (final EquipmentSlot type : EquipmentSlot.values())
             LogicMiscUpdate.getStatModifiers(type, this.entity.getItemBySlot(type), vals);
         if (this.mods == null) this.mods = this.pokemob.getModifiers().getModifiers(StatModifiers.ARMOUR);

src/main/java/pokecube/core/ai/tasks/TaskBase.java

@@ -120,9 +120,6 @@ public static boolean canMove(final IPokemob pokemob)
 
     int priority = 0;
 
-    boolean tempRun = false;
-    boolean tempCont = false;
-
     public TaskBase(final IPokemob pokemob)
     {
         this(pokemob, ImmutableMap.of());

src/main/java/pokecube/core/ai/tasks/combat/CombatTask.java

@@ -1,5 +1,6 @@
 package pokecube.core.ai.tasks.combat;
 
+import java.util.Comparator;
 import java.util.Map;
 
 import com.google.common.collect.Maps;
@@ -32,4 +33,13 @@ public CombatTask(final IPokemob pokemob, final Map<MemoryModuleType<?>, MemoryS
         super(pokemob, RootTask.merge(CombatTask.MEMS, mems));
     }
 
+    @Override
+    protected Comparator<MemoryModuleType<?>> getCheckOrder()
+    {
+        return (a, b) -> {
+            int a1 = a == MemoryModules.ATTACKTARGET.get() ? 1 : 0;
+            int b1 = b == MemoryModules.ATTACKTARGET.get() ? 1 : 0;
+            return Integer.compare(a1, b1);
+        };
+    }
 }

src/main/java/pokecube/core/ai/tasks/idle/IdleWalkTask.java

@@ -90,6 +90,12 @@ public IdleWalkTask(final IPokemob pokemob)
         this.entry = pokemob.getPokedexEntry();
     }
 
+    @Override
+    protected boolean simpleRun()
+    {
+        return true;
+    }
+
     /** Floating things try to stay their preferedHeight from the ground. */
     private void doFloatingIdle()
     {

src/main/java/pokecube/core/ai/tasks/misc/FollowOwnerTask.java

@@ -6,6 +6,7 @@
 import com.google.common.collect.Maps;
 
 import net.minecraft.world.entity.LivingEntity;
+import net.minecraft.world.entity.Mob;
 import net.minecraft.world.entity.ai.attributes.AttributeInstance;
 import net.minecraft.world.entity.ai.attributes.AttributeModifier;
 import net.minecraft.world.entity.ai.attributes.Attributes;
@@ -120,21 +121,28 @@ else if (!this.petPathfinder.isDone()
         this.setWalkTo(target);
     }
 
+    @Override
+    protected boolean shouldNotRun(Mob mobIn)
+    {
+        return this.pokemob.getOwner() == null;
+    }
+
     @Override
     public boolean shouldRun()
     {
         if (!this.pokemob.isRoutineEnabled(AIRoutine.FOLLOW)) return false;
         if (!TaskBase.canMove(this.pokemob)) return false;
+        if (this.pokemob.getGeneralState(GeneralStates.STAYING)) return false;
         final LivingEntity LivingEntity = this.pokemob.getOwner();
-        this.petPathfinder = this.entity.getNavigation();
         // Nothing to follow
         if (LivingEntity == null) return false;
-        else if (this.pokemob.getGeneralState(GeneralStates.STAYING)) return false;
-        else if (this.pathing && this.entity.distanceToSqr(LivingEntity) > this.maxDist * this.maxDist) return true;
-        else if (this.entity.distanceToSqr(LivingEntity) < this.minDist * this.minDist) return false;
-        else if (new Vector3().set(LivingEntity).distToSq(this.ownerPos) < this.minDist * this.minDist) return false;
+        double dr2 = this.entity.distanceToSqr(LivingEntity);
+        if (this.pathing && dr2 > this.maxDist * this.maxDist) return true;
+        if (dr2 < this.minDist * this.minDist) return false;
+        if (new Vector3().set(LivingEntity).distToSq(this.ownerPos) < this.minDist * this.minDist) return false;
+        this.petPathfinder = this.entity.getNavigation();
         // Follow owner.
-        else return true;
+        return true;
     }
 
 }

src/main/java/pokecube/core/ai/tasks/misc/LookAtMob.java

@@ -45,6 +45,12 @@ public LookAtMob(final Predicate<LivingEntity> matcher, final float distance)
         this.distance_squared = distance * distance;
     }
 
+    @Override
+    protected boolean simpleRun()
+    {
+        return true;
+    }
+
     @Override
     protected boolean canTimeOut()
     {

src/main/java/pokecube/core/ai/tasks/misc/RunAway.java

@@ -25,6 +25,12 @@ public RunAway(final MemoryModuleType<? extends Entity> badMemory, final float s
         this.runSpeed = speed;
     }
 
+    @Override
+    protected boolean simpleRun()
+    {
+        return true;
+    }
+
     @Override
     protected boolean checkExtraStartConditions(final ServerLevel worldIn, final PathfinderMob owner)
     {

src/main/java/pokecube/mixin/brain/MixinMobEntity.java

@@ -20,6 +20,7 @@
 import net.minecraft.world.entity.ai.memory.MemoryStatus;
 import net.minecraft.world.entity.schedule.Activity;
 import net.minecraft.world.level.Level;
+import pokecube.core.ai.brain.BrainUtils;
 import pokecube.core.ai.brain.MemoryModules;
 import thut.api.entity.ai.BrainUtil;
 import thut.api.entity.ai.RootTask;
@@ -84,6 +85,7 @@ protected void onPostUpdateAITasks(final CallbackInfo cbi)
         {
             this.ticked_default_ai = this.brain.getMemory(MemoryModules.DUMMY.get()).get();
             this.checked_for_ai = true;
+            BrainUtils.removeMatchingTasks(this.brain, s -> s instanceof DummySetTask);
         }
         if (!this.ticked_default_ai)
         {

src/main/java/thut/api/entity/ai/RootTask.java

@@ -1,13 +1,15 @@
 package thut.api.entity.ai;
 
+import java.util.Comparator;
 import java.util.List;
 import java.util.Map;
-import java.util.Random;
 
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Lists;
 import com.google.common.collect.Maps;
 
+import it.unimi.dsi.fastutil.objects.Object2IntArrayMap;
+import it.unimi.dsi.fastutil.objects.Object2LongArrayMap;
 import net.minecraft.core.BlockPos;
 import net.minecraft.world.entity.Entity;
 import net.minecraft.world.entity.LivingEntity;
@@ -31,33 +33,38 @@ public static Map<MemoryModuleType<?>, MemoryStatus> merge(final Map<MemoryModul
         return ImmutableMap.copyOf(ret);
     }
 
+    public static record MemoryRequriment(MemoryModuleType<?> memory, MemoryStatus status)
+    {
+    }
+
     public static boolean doLoadThrottling = false;
 
     public static int runRate = 10;
 
+    public static Comparator<MemoryModuleType<?>> _NO_ORDER = (a, b) -> 0;
+
     protected final Map<MemoryModuleType<?>, MemoryStatus> neededMems;
 
-    private final MemoryModuleType<?>[] neededModules;
-    private final MemoryStatus[] neededStatus;
+    private final MemoryRequriment[] memoryRequirements;
 
     protected E entity;
 
-    private Random _run_rng;
-
     protected boolean runWhileDead = false;
+    // This is a minimum run rate when load balancing is enabled, or for tasks
+    // that autothrottle, 1 means every tick.
+    protected int _run_rate = 10;
+
+    protected boolean tempRun = false;
+    protected boolean tempCont = false;
 
     public RootTask(final E entity, final Map<MemoryModuleType<?>, MemoryStatus> neededMems, final int duration,
             final int maxDuration)
     {
         super(neededMems, duration, maxDuration);
         this.entity = entity;
         this.neededMems = neededMems;
-        final List<MemoryModuleType<?>> neededModules = Lists.newArrayList();
-        final List<MemoryStatus> neededStatus = Lists.newArrayList();
-        neededModules.addAll(neededMems.keySet());
-        for (final MemoryModuleType<?> mod : neededModules) neededStatus.add(neededMems.get(mod));
-        this.neededModules = neededModules.toArray(new MemoryModuleType<?>[neededModules.size()]);
-        this.neededStatus = neededStatus.toArray(new MemoryStatus[neededModules.size()]);
+        memoryRequirements = new MemoryRequriment[neededMems.size()];
+        initMems(memoryRequirements, neededMems);
     }
 
     public RootTask(final E entity, final Map<MemoryModuleType<?>, MemoryStatus> neededMems, final int duration)
@@ -85,9 +92,28 @@ public RootTask(final Map<MemoryModuleType<?>, MemoryStatus> neededMems)
         this(null, neededMems, 60);
     }
 
-    protected boolean canTimeOut()
+    protected Comparator<MemoryModuleType<?>> getCheckOrder()
     {
-        return false;
+        return _NO_ORDER;
+    }
+
+    protected void initMems(MemoryRequriment[] reqs, final Map<MemoryModuleType<?>, MemoryStatus> neededMems)
+    {
+        final List<MemoryModuleType<?>> neededModules = Lists.newArrayList();
+        neededModules.addAll(neededMems.keySet());
+        neededModules.sort(getCheckOrder());
+        for (int i = 0; i < neededModules.size(); i++)
+        {
+            var mem = neededModules.get(i);
+            var status = neededMems.get(mem);
+            reqs[i] = new MemoryRequriment(mem, status);
+        }
+    }
+
+    protected boolean runTick(final E mobIn)
+    {
+        int rate = Math.max(runRate, _run_rate);
+        return mobIn.tickCount % rate == mobIn.id % rate;
     }
 
     protected void setWalkTo(final Vector3 pos, final double speed, final int dist)
@@ -130,35 +156,132 @@ protected void setWalkTo(WalkTarget target)
     protected final boolean isPaused(final E mobIn)
     {
         if (!this.loadThrottle() || !RootTask.doLoadThrottling) return false;
-        if (this._run_rng == null) this._run_rng = new Random(entity.getUUID().hashCode());
-        final int tick = _run_rng.nextInt(RootTask.runRate);
-        return mobIn.tickCount % RootTask.runRate != tick;
+        return !runTick(mobIn);
+    }
+
+    @Override
+    protected boolean timedOut(final long gameTime)
+    {
+        if (!this.canTimeOut()) return false;
+        return super.timedOut(gameTime);
     }
 
     public boolean loadThrottle()
     {
         return false;
     }
 
-    @Override
-    protected boolean timedOut(final long gameTime)
+    protected boolean canTimeOut()
     {
-        if (!this.canTimeOut()) return false;
-        return super.timedOut(gameTime);
+        return false;
+    }
+
+    protected boolean simpleRun()
+    {
+        return false;
+    }
+
+    protected boolean shouldNotRun(final E mobIn)
+    {
+        return false;
+    }
+
+    static long start = System.nanoTime();
+    static long n = 0;
+    static long dt = 0;
+    static Object2LongArrayMap<Class<?>> taskTimes = new Object2LongArrayMap<>();
+    static Object2IntArrayMap<Class<?>> taskNs = new Object2IntArrayMap<>();
+
+    static void timerStart()
+    {
+        start = System.nanoTime();
+    }
+
+    static void timerEnd(Class<?> involved)
+    {
+        long _dt = System.nanoTime() - start;
+        dt += _dt;
+        taskTimes.compute(involved, (key, value) -> {
+            if (value == null) value = _dt;
+            else value += _dt;
+            return value;
+        });
+        taskNs.compute(involved, (key, value) -> {
+            if (value == null) value = 1;
+            else value += 1;
+            return value;
+        });
+        n++;
+        if (n >= 1000000)
+        {
+            double avg = dt / ((double) n);
+            System.out.println("Average time: " + (avg / 1000d) + "us");
+            System.out.println("class\ttime per\ttime total");
+            taskTimes.forEach((clazz, val) -> {
+                double avg2 = val / ((double) taskNs.getInt(clazz));
+                String key = "%s\t%.2f\t%.2f";
+                System.out.println(key.formatted(clazz, (avg2 / 1000d), (val / 1000d)));
+            });
+            taskTimes.clear();
+            taskNs.clear();
+            n = 0;
+            dt = 0;
+        }
     }
 
     @Override
     public boolean hasRequiredMemories(final E mobIn)
     {
-        this.entity = mobIn;
-        // If we are paused, return early here.
-        if (this.isPaused(mobIn)) return false;
-        final Brain<?> brain = mobIn.getBrain();
-        for (int i = 0; i < this.neededStatus.length; i++)
-            if (!brain.checkMemory(this.neededModules[i], this.neededStatus[i])) return false;
-        // Dead mobs don't have AI!
-        if (!this.runWhileDead && !mobIn.isAlive()) return false;
-        // Otherwise continue;
-        return true;
+//        timerStart();
+        // Default to true, everything below will set false.
+        boolean ret = true;
+        check:
+        {
+            this.entity = mobIn;
+
+            // if we are a "simple" check, we only run every so often, as we
+            // will run everything immediately.
+            if (this.simpleRun() && !this.runTick(mobIn))
+            {
+                ret = false;
+                break check;
+            }
+
+            // This means it has a simpler check for if not to run, so we do
+            // that instead of below.
+            if (this.shouldNotRun(mobIn))
+            {
+                ret = false;
+                break check;
+            }
+
+            // If we are paused, return early here.
+            if (this.isPaused(mobIn))
+            {
+                ret = this.tempCont;
+                break check;
+            }
+
+            final Brain<?> brain = mobIn.getBrain();
+            // Check memories, this loop takes 50% of the time to run as the
+            // vanilla way, so we do it like this.
+            for (var mem : this.memoryRequirements)
+            {
+                if (!brain.checkMemory(mem.memory(), mem.status()))
+                {
+                    ret = false;
+                    break check;
+                }
+            }
+            // Dead mobs don't have AI!
+            if (!this.runWhileDead && !mobIn.isAlive())
+            {
+                ret = false;
+                break check;
+            }
+        }
+        this.tempCont = ret;
+//        timerEnd(this.getClass());
+        return ret;
     }
 }