Improve BallApplicator performance

Previous logic would check a new LegalityAnalysis for each ball attempted, and would additionally allocate a new PKM for testing the balls.

With the refactoring to BallVerifier, we can just pass in the template and ball and get a truth, skipping the entire re-check. There might be a deficiency when the current ball invalidates the encounter (GO) but I don't think that's worth considering at this time.

PKHeX.Core/Editing/Applicators/BallApplicator.cs

@@ -1,5 +1,4 @@
 using System;
-using System.Collections.Generic;
 using static PKHeX.Core.Ball;
 
 namespace PKHeX.Core;
@@ -9,194 +8,127 @@ namespace PKHeX.Core;
 /// </summary>
 public static class BallApplicator
 {
+    private static readonly Ball[] BallList = Enum.GetValues<Ball>();
+    public const byte MaxBallSpanAlloc = (byte)LAOrigin + 1;
+
+    /// <remarks>
+    /// Requires checking the <see cref="LegalityAnalysis"/>.
+    /// </remarks>
+    /// <inheritdoc cref="GetLegalBalls(Span{Ball}, PKM, IEncounterTemplate)"/>
+    public static int GetLegalBalls(Span<Ball> result, PKM pk) => GetLegalBalls(result, pk, new LegalityAnalysis(pk));
+
+    /// <inheritdoc cref="GetLegalBalls(Span{Ball}, PKM, IEncounterTemplate)"/>
+    public static int GetLegalBalls(Span<Ball> result, PKM pk, LegalityAnalysis la) => GetLegalBalls(result, pk, la.EncounterOriginal);
+
     /// <summary>
     /// Gets all balls that are legal for the input <see cref="PKM"/>.
     /// </summary>
-    /// <remarks>
+    /// <param name="result">Result storage.</param>
-    /// Requires checking the <see cref="LegalityAnalysis"/> for every <see cref="Ball"/> that is tried.
-    /// </remarks>
     /// <param name="pk">Pokémon to retrieve a list of valid balls for.</param>
-    /// <returns>Enumerable list of <see cref="Ball"/> values that the <see cref="PKM"/> is legal with.</returns>
+    /// <param name="enc">Encounter matched to.</param>
-    public static IEnumerable<Ball> GetLegalBalls(PKM pk)
+    /// <returns>Count of <see cref="Ball"/> values that the <see cref="PKM"/> is legal with.</returns>
+    public static int GetLegalBalls(Span<Ball> result, PKM pk, IEncounterTemplate enc)
     {
-        var clone = pk.Clone();
+        if (enc.Species is (ushort)Species.Nincada && pk.Species is (ushort)Species.Shedinja)
+            return GetLegalBallsEvolvedShedinja(result, pk, enc);
+        return LoadLegalBalls(result, pk, enc);
+    }
+
+    private static ReadOnlySpan<Ball> ShedinjaEvolve4 => [Sport, Poke];
+
+    private static int GetLegalBallsEvolvedShedinja(Span<Ball> result, PKM pk, IEncounterTemplate enc)
+    {
+        switch (enc)
+        {
+            case EncounterSlot4 when IsNincadaEvolveInOrigin(pk, enc):
+                ShedinjaEvolve4.CopyTo(result);
+                return ShedinjaEvolve4.Length;
+            case EncounterSlot3 when IsNincadaEvolveInOrigin(pk, enc):
+                return LoadLegalBalls(result, pk, enc);
+        }
+        result[0] = Poke;
+        return 1;
+    }
+
+    private static bool IsNincadaEvolveInOrigin(PKM pk, IEncounterTemplate enc)
+    {
+        // Rough check to see if Nincada evolved in the origin context (Gen3/4).
+        // Does not do PID/IV checks to know the original met level.
+        var current = pk.CurrentLevel;
+        var met = pk.MetLevel;
+        if (pk.Format == enc.Generation)
+            return current > met;
+        return enc.LevelMin != met && current > enc.LevelMin;
+    }
+
+    private static int LoadLegalBalls(Span<Ball> result, PKM pk, IEncounterTemplate enc)
+    {
+        int ctr = 0;
         foreach (var b in BallList)
         {
-            var ball = (byte)b;
+            if (BallVerifier.VerifyBall(enc, b, pk).IsValid())
-            clone.Ball = ball;
+                result[ctr++] = b;
-            if (clone.Ball != ball)
-                continue; // Some setters guard against out of bounds values.
-            if (new LegalityAnalysis(clone).Valid)
-                yield return b;
         }
+        return ctr;
     }
 
     /// <summary>
     /// Applies a random legal ball value if any exist.
     /// </summary>
     /// <remarks>
-    /// Requires checking the <see cref="LegalityAnalysis"/> for every <see cref="Ball"/> that is tried.
+    /// Requires checking the <see cref="LegalityAnalysis"/>.
     /// </remarks>
     /// <param name="pk">Pokémon to modify.</param>
     public static byte ApplyBallLegalRandom(PKM pk)
     {
         Span<Ball> balls = stackalloc Ball[MaxBallSpanAlloc];
-        var count = GetBallListFromColor(pk, balls);
+        var count = GetLegalBalls(balls, pk);
         balls = balls[..count];
         Util.Rand.Shuffle(balls);
-        return ApplyFirstLegalBall(pk, balls);
+        return ApplyFirstLegalBall(pk, balls, []);
     }
 
+    public static byte ApplyBallLegalByColor(PKM pk) => ApplyBallLegalByColor(pk, PersonalColorUtil.GetColor(pk));
+    public static byte ApplyBallLegalByColor(PKM pk, PersonalColor color) => ApplyBallLegalByColor(pk, new LegalityAnalysis(pk), color);
+    public static byte ApplyBallLegalByColor(PKM pk, LegalityAnalysis la, PersonalColor color) => ApplyBallLegalByColor(pk, la.EncounterOriginal, color);
+
     /// <summary>
     /// Applies a legal ball value if any exist, ordered by color.
     /// </summary>
-    /// <remarks>
-    /// Requires checking the <see cref="LegalityAnalysis"/> for every <see cref="Ball"/> that is tried.
-    /// </remarks>
     /// <param name="pk">Pokémon to modify.</param>
-    public static byte ApplyBallLegalByColor(PKM pk)
+    /// <param name="enc">Encounter matched to.</param>
+    /// <param name="color">Color preference to order by.</param>
+    private static byte ApplyBallLegalByColor(PKM pk, IEncounterTemplate enc, PersonalColor color)
     {
         Span<Ball> balls = stackalloc Ball[MaxBallSpanAlloc];
-        GetBallListFromColor(pk, balls);
+        var count = GetLegalBalls(balls, pk, enc);
-        return ApplyFirstLegalBall(pk, balls);
+        balls = balls[..count];
+        var prefer = PersonalColorUtil.GetPreferredByColor(enc, color);
+        return ApplyFirstLegalBall(pk, balls, prefer);
     }
 
-    /// <summary>
+    private static byte ApplyFirstLegalBall(PKM pk, Span<Ball> legal, ReadOnlySpan<Ball> prefer)
-    /// Applies a random ball value in a cyclical manner.
-    /// </summary>
-    /// <param name="pk">Pokémon to modify.</param>
-    public static byte ApplyBallNext(PKM pk)
     {
-        Span<Ball> balls = stackalloc Ball[MaxBallSpanAlloc];
+        foreach (var ball in prefer)
-        GetBallList(pk.Ball, balls);
-        var next = balls[0];
-        return pk.Ball = (byte)next;
-    }
-
-    private static byte ApplyFirstLegalBall(PKM pk, ReadOnlySpan<Ball> balls)
-    {
-        var initial = pk.Ball;
-        foreach (var b in balls)
         {
-            var test = (byte)b;
+            if (Contains(legal, ball))
-            pk.Ball = test;
+                return pk.Ball = (byte)ball;
-            if (new LegalityAnalysis(pk).Valid)
-                return test;
         }
-        return initial; // fail, revert
+        foreach (var ball in legal)
-    }
-
-    private static int GetBallList(byte ball, Span<Ball> result)
-    {
-        var balls = BallList;
-        var currentBall = (Ball)ball;
-        return GetCircularOnce(balls, currentBall, result);
-    }
-
-    private static int GetBallListFromColor(PKM pk, Span<Ball> result)
-    {
-        // Gen1/2 don't store color in personal info
-        var pi = pk.Format >= 3 ? pk.PersonalInfo : PersonalTable.USUM.GetFormEntry(pk.Species, 0);
-        var color = (PersonalColor)pi.Color;
-        var balls = BallColors[(int)color];
-        var currentBall = (Ball)pk.Ball;
-        return GetCircularOnce(balls, currentBall, result);
-    }
-
-    private static int GetCircularOnce<T>(T[] items, T current, Span<T> result)
-    {
-        var currentIndex = Array.IndexOf(items, current);
-        if (currentIndex < 0)
-            currentIndex = items.Length - 2;
-        return GetCircularOnce(items, currentIndex, result);
-    }
-
-    private static int GetCircularOnce<T>(ReadOnlySpan<T> items, int startIndex, Span<T> result)
-    {
-        var tail = items[(startIndex + 1)..];
-        tail.CopyTo(result);
-        items[..startIndex].CopyTo(result[tail.Length..]);
-        return items.Length;
-    }
-
-    private static readonly Ball[] BallList = Enum.GetValues<Ball>();
-    private static int MaxBallSpanAlloc => BallList.Length;
-
-    static BallApplicator()
-    {
-        ReadOnlySpan<Ball> exclude = [None, Poke];
-        ReadOnlySpan<Ball> end = [Poke];
-        Span<Ball> all = stackalloc Ball[BallList.Length - exclude.Length];
-        all = all[..FillExcept(all, exclude, BallList)];
-
-        var colors = Enum.GetValues<PersonalColor>();
-        foreach (var color in colors)
         {
-            int c = (int)color;
+            if (!Contains(prefer, ball))
-            // Replace the array reference with a new array that appends non-matching values, followed by the end values.
+                return pk.Ball = (byte)ball;
-            var defined = BallColors[c];
-            Span<Ball> match = (BallColors[c] = new Ball[all.Length + end.Length]);
-            defined.CopyTo(match);
-            FillExcept(match[defined.Length..], defined, all);
-            end.CopyTo(match[^end.Length..]);
         }
+        return pk.Ball; // fail
 
-        static int FillExcept(Span<Ball> result, ReadOnlySpan<Ball> exclude, ReadOnlySpan<Ball> all)
+        static bool Contains(ReadOnlySpan<Ball> balls, Ball ball)
         {
-            int ctr = 0;
+            foreach (var b in balls)
-            foreach (var b in all)
             {
-                if (Contains(exclude, b))
+                if (b == ball)
-                    continue;
+                    return true;
-                result[ctr++] = b;
-            }
-            return ctr;
-
-            static bool Contains(ReadOnlySpan<Ball> arr, Ball b)
-            {
-                foreach (var a in arr)
-                {
-                    if (a == b)
-                        return true;
-                }
-                return false;
             }
+            return false;
         }
     }
-
-    /// <summary>
-    /// Priority Match ball IDs that match the color ID in descending order
-    /// </summary>
-    private static readonly Ball[][] BallColors =
-    [
-        /* Red */    [Cherish, Repeat, Fast, Heal, Great, Dream, Lure],
-        /* Blue */   [Dive, Net, Great, Beast, Lure],
-        /* Yellow */ [Level, Ultra, Repeat, Quick, Moon],
-        /* Green */  [Safari, Friend, Nest, Dusk],
-        /* Black */  [Luxury, Heavy, Ultra, Moon, Net, Beast],
-
-        /* Brown */  [Level, Heavy],
-        /* Purple */ [Master, Love, Dream, Heal],
-        /* Gray */   [Heavy, Premier, Luxury],
-        /* White */  [Premier, Timer, Luxury, Ultra],
-        /* Pink */   [Love, Dream, Heal],
-    ];
-
-    /// <summary>
-    /// Personal Data color IDs
-    /// </summary>
-    private enum PersonalColor : byte
-    {
-        Red,
-        Blue,
-        Yellow,
-        Green,
-        Black,
-
-        Brown,
-        Purple,
-        Gray,
-        White,
-        Pink,
-    }
 }

PKHeX.Core/Editing/Applicators/PersonalColorUtil.cs

@@ -0,0 +1,65 @@
+using System;
+
+namespace PKHeX.Core;
+
+public static class PersonalColorUtil
+{
+    public static PersonalColor GetColor(PKM pk)
+    {
+        // Gen1/2 don't store color in personal info
+        if (pk.Format < 3)
+            return (PersonalColor)PersonalTable.USUM[pk.Species, 0].Color;
+        return (PersonalColor)pk.PersonalInfo.Color;
+    }
+
+    public static PersonalColor GetColor(IEncounterTemplate enc)
+    {
+        // Gen1/2 don't store color in personal info
+        if (enc.Generation < 3)
+            return (PersonalColor)PersonalTable.USUM[enc.Species, 0].Color;
+
+        var pt = GameData.GetPersonal(enc.Version);
+        var pi = pt[enc.Species, enc.Form];
+        return (PersonalColor)pi.Color;
+    }
+
+    public static ReadOnlySpan<Ball> GetPreferredByColor(IEncounterTemplate enc) => GetPreferredByColor(enc, GetColor(enc));
+
+    public static ReadOnlySpan<Ball> GetPreferredByColor<T>(T enc, PersonalColor color) where T : IVersion
+    {
+        if (enc.Version is GameVersion.PLA)
+            return GetPreferredByColorLA(color);
+        return GetPreferredByColor(color);
+    }
+
+    /// <summary>
+    /// Priority Match ball IDs that match the color ID
+    /// </summary>
+    public static ReadOnlySpan<Ball> GetPreferredByColor(PersonalColor color) => color switch
+    {
+        PersonalColor.Red => [Ball.Repeat, Ball.Fast, Ball.Heal, Ball.Great, Ball.Dream, Ball.Lure],
+        PersonalColor.Blue => [Ball.Dive, Ball.Net, Ball.Great, Ball.Lure, Ball.Beast],
+        PersonalColor.Yellow => [Ball.Level, Ball.Ultra, Ball.Repeat, Ball.Quick, Ball.Moon],
+        PersonalColor.Green => [Ball.Safari, Ball.Friend, Ball.Nest, Ball.Dusk],
+        PersonalColor.Black => [Ball.Luxury, Ball.Heavy, Ball.Ultra, Ball.Moon, Ball.Net, Ball.Beast],
+        PersonalColor.Brown => [Ball.Level, Ball.Heavy],
+        PersonalColor.Purple => [Ball.Master, Ball.Love, Ball.Heal, Ball.Dream],
+        PersonalColor.Gray => [Ball.Heavy, Ball.Premier, Ball.Luxury],
+        PersonalColor.White => [Ball.Premier, Ball.Timer, Ball.Luxury, Ball.Ultra],
+        _ => [Ball.Love, Ball.Heal, Ball.Dream],
+    };
+
+    public static ReadOnlySpan<Ball> GetPreferredByColorLA(PersonalColor color) => color switch
+    {
+        PersonalColor.Red => [Ball.LAPoke],
+        PersonalColor.Blue => [Ball.LAFeather, Ball.LAGreat, Ball.LAJet],
+        PersonalColor.Yellow => [Ball.LAUltra],
+        PersonalColor.Green => [Ball.LAPoke],
+        PersonalColor.Black => [Ball.LAGigaton, Ball.LALeaden, Ball.LAHeavy, Ball.LAUltra],
+        PersonalColor.Brown => [Ball.LAPoke],
+        PersonalColor.Purple => [Ball.LAPoke],
+        PersonalColor.Gray => [Ball.LAGigaton, Ball.LALeaden, Ball.LAHeavy],
+        PersonalColor.White => [Ball.LAWing, Ball.LAJet],
+        _ => [Ball.LAPoke],
+    };
+}

PKHeX.Core/PersonalInfo/Enums/EggGroup.cs

@@ -0,0 +1,24 @@
+namespace PKHeX.Core;
+
+/// <summary>
+/// Personal Data egg groups for breeding compatibility
+/// </summary>
+public enum EggGroup : byte
+{
+    None = 0,
+    Monster = 1,
+    Water1 = 2,
+    Bug = 3,
+    Flying = 4,
+    Field = 5,
+    Fairy = 6,
+    Grass = 7,
+    HumanLike = 8,
+    Water3 = 9,
+    Mineral = 10,
+    Amorphous = 11,
+    Water2 = 12,
+    Ditto = 13,
+    Dragon = 14,
+    Undiscovered = 15,
+}

PKHeX.Core/PersonalInfo/Enums/PersonalColor.cs

@@ -0,0 +1,19 @@
+namespace PKHeX.Core;
+
+/// <summary>
+/// Personal Data color IDs
+/// </summary>
+public enum PersonalColor : byte
+{
+    Red = 0,
+    Blue = 1,
+    Yellow = 2,
+    Green = 3,
+    Black = 4,
+
+    Brown = 5,
+    Purple = 6,
+    Gray = 7,
+    White = 8,
+    Pink = 9,
+}

PKHeX.WinForms/Controls/PKM Editor/BallBrowser.cs

@@ -17,20 +17,21 @@ public partial class BallBrowser : Form
 
     public void LoadBalls(PKM pk)
     {
-        var legal = BallApplicator.GetLegalBalls(pk);
+        Span<Ball> valid = stackalloc Ball[BallApplicator.MaxBallSpanAlloc];
-        LoadBalls(legal, pk.MaxBallID + 1);
+        var legal = BallApplicator.GetLegalBalls(valid, pk);
+        LoadBalls(valid[..legal], pk.MaxBallID);
     }
 
-    private void LoadBalls(IEnumerable<Ball> legal, int max)
+    private void LoadBalls(ReadOnlySpan<Ball> legal, int max)
     {
-        Span<bool> flags = stackalloc bool[max];
+        Span<bool> flags = stackalloc bool[BallApplicator.MaxBallSpanAlloc];
         foreach (var ball in legal)
             flags[(int)ball] = true;
 
         int countLegal = 0;
         List<PictureBox> controls = [];
         var names = GameInfo.BallDataSource;
-        for (byte ballID = 1; ballID < flags.Length; ballID++)
+        for (byte ballID = 1; ballID <= max; ballID++)
         {
             var name = GetBallName(ballID, names);
             var pb = GetBallView(ballID, name, flags[ballID]);