Avoid storing every single visited pixel, store only edges instead

src/PixiEditor.ChangeableDocument/Changes/Selection/MagicWand/MagicWandHelper.cs

@@ -9,8 +9,6 @@ using PixiEditor.DrawingApi.Core.Surface.Vector;
 namespace PixiEditor.ChangeableDocument.Changes.Selection.MagicWand;
 internal class MagicWandHelper
 {
-    private const byte Visited = 2;
-
     private static readonly VecI Up = new VecI(0, -1);
     private static readonly VecI Down = new VecI(0, 1);
     private static readonly VecI Left = new VecI(-1, 0);
@@ -18,6 +16,88 @@ internal class MagicWandHelper
 
     private static MagicWandVisualizer visualizer = new MagicWandVisualizer(Path.Combine("Debugging", "MagicWand"));
 
+    private class UnvisitedStack
+    {
+        private int chunkSize;
+        private readonly VecI imageSizeInChunks;
+        private Stack<(VecI chunkPos, VecI posOnChunk)> likelyUnvisited = new();
+        private HashSet<VecI> certainlyVisited = new();
+
+        public UnvisitedStack(int chunkSize, VecI imageSizeInChunks)
+        {
+            this.chunkSize = chunkSize;
+            this.imageSizeInChunks = imageSizeInChunks;
+        }
+
+        public void PushAll(VecI chunkPos)
+        {
+            VecI chunkOffset = chunkPos * chunkSize;
+            for (int i = 0; i < chunkSize; i++)
+            {
+                if (chunkPos.Y > 0)
+                    likelyUnvisited.Push((new(chunkPos.X, chunkPos.Y - 1), new(i, chunkSize - 1)));
+                certainlyVisited.Add(chunkOffset + new VecI(i, 0));
+
+                if (chunkPos.Y < imageSizeInChunks.Y - 1)
+                    likelyUnvisited.Push((new(chunkPos.X, chunkPos.Y + 1), new(i, 0)));
+                certainlyVisited.Add(chunkOffset + new VecI(i, chunkSize - 1));
+
+                if (chunkPos.X > 0)
+                    likelyUnvisited.Push((new(chunkPos.X - 1, chunkPos.Y), new(chunkSize - 1, i)));
+                certainlyVisited.Add(chunkOffset + new VecI(0, i));
+
+                if (chunkPos.X < imageSizeInChunks.X - 1)
+                    likelyUnvisited.Push((new(chunkPos.X + 1, chunkPos.Y), new(0, i)));
+                certainlyVisited.Add(chunkOffset + new VecI(chunkSize - 1, i));
+            }
+        }
+
+        public void Push(VecI chunkPos, VecI posOnChunk)
+        {
+            likelyUnvisited.Push((chunkPos, posOnChunk));
+        }
+
+        public void Push(VecI chunkPos, bool[] visitedArray)
+        {
+            VecI chunkOffset = chunkPos * chunkSize;
+            for (int i = 0; i < chunkSize; i++)
+            {
+                if (chunkPos.Y > 0 && visitedArray[i]) //Top
+                    likelyUnvisited.Push((new(chunkPos.X, chunkPos.Y - 1), new(i, chunkSize - 1)));
+                if (visitedArray[i])
+                    certainlyVisited.Add(chunkOffset + new VecI(i, 0));
+
+                if (chunkPos.Y < imageSizeInChunks.Y - 1 && visitedArray[chunkSize * (chunkSize - 1) + i]) // Bottom
+                    likelyUnvisited.Push((new(chunkPos.X, chunkPos.Y + 1), new(i, 0)));
+                if (visitedArray[chunkSize * (chunkSize - 1) + i])
+                    certainlyVisited.Add(chunkOffset + new VecI(i, chunkSize - 1));
+
+                if (chunkPos.X > 0 && visitedArray[i * chunkSize]) // Left
+                    likelyUnvisited.Push((new(chunkPos.X - 1, chunkPos.Y), new(chunkSize - 1, i)));
+                if (visitedArray[i * chunkSize])
+                    certainlyVisited.Add(chunkOffset + new VecI(0, i));
+
+                if (chunkPos.X < imageSizeInChunks.X - 1 && visitedArray[i * chunkSize + (chunkSize - 1)]) // Right
+                    likelyUnvisited.Push((new(chunkPos.X + 1, chunkPos.Y), new(0, i)));
+                if (visitedArray[i * chunkSize + (chunkSize - 1)])
+                    certainlyVisited.Add(chunkOffset + new VecI(chunkSize - 1, i));
+            }
+        }
+
+        public (VecI chunkPos, VecI posOnChunk)? PopUnvisited()
+        {
+            while (likelyUnvisited.Count > 0)
+            {
+                var (chunkPos, posOnChunk) = likelyUnvisited.Pop();
+                VecI global = chunkPos * chunkSize + posOnChunk;
+                if (certainlyVisited.Contains(global))
+                    continue;
+                return (chunkPos, posOnChunk);
+            }
+            return null;
+        }
+    }
+
     public static VectorPath DoMagicWandFloodFill(VecI startingPos, HashSet<Guid> membersToFloodFill,
         IReadOnlyDocument document)
     {
@@ -41,16 +121,19 @@ internal class MagicWandHelper
         ColorBounds colorRange = new(colorToReplace);
 
         HashSet<VecI> processedEmptyChunks = new();
-        HashSet<VecI> processedPositions = new();
-        Stack<(VecI chunkPos, VecI posOnChunk)> positionsToFloodFill = new();
-        positionsToFloodFill.Push((initChunkPos, initPosOnChunk));
 
-        Lines lines = new();
+        UnvisitedStack positionsToFloodFill = new(chunkSize, imageSizeInChunks);
 
+        Lines lines = new();
         VectorPath selection = new();
-        while (positionsToFloodFill.Count > 0)
+
+        positionsToFloodFill.Push(initChunkPos, initPosOnChunk);
+        while (true)
         {
-            var (chunkPos, posOnChunk) = positionsToFloodFill.Pop();
+            (VecI initChunkPos, VecI initPosOnChunk)? popped = positionsToFloodFill.PopUnvisited();
+            if (popped is null)
+                break;
+            var (chunkPos, posOnChunk) = popped.Value;
             var referenceChunk = cache.GetChunk(chunkPos);
 
             // don't call floodfill if the chunk is empty
@@ -58,18 +141,8 @@ internal class MagicWandHelper
             {
                 if (colorToReplace.A == 0 && !processedEmptyChunks.Contains(chunkPos))
                 {
-                    AddLinesForEmptyChunk(lines, chunkPos, document.Size, imageSizeInChunks, chunkSize);
-                    for (int i = 0; i < chunkSize; i++)
-                    {
-                        if (chunkPos.Y > 0)
-                            positionsToFloodFill.Push((new(chunkPos.X, chunkPos.Y - 1), new(i, chunkSize - 1)));
-                        if (chunkPos.Y < imageSizeInChunks.Y - 1)
-                            positionsToFloodFill.Push((new(chunkPos.X, chunkPos.Y + 1), new(i, 0)));
-                        if (chunkPos.X > 0)
-                            positionsToFloodFill.Push((new(chunkPos.X - 1, chunkPos.Y), new(chunkSize - 1, i)));
-                        if (chunkPos.X < imageSizeInChunks.X - 1)
-                            positionsToFloodFill.Push((new(chunkPos.X + 1, chunkPos.Y), new(0, i)));
-                    }
+                    AddLinesForEmptyChunk(lines, chunkPos, document.Size, chunkSize);
+                    positionsToFloodFill.PushAll(chunkPos);
                     processedEmptyChunks.Add(chunkPos);
                 }
                 continue;
@@ -79,10 +152,6 @@ internal class MagicWandHelper
             var reallyReferenceChunk = referenceChunk.AsT0;
 
             VecI globalPos = chunkPos * chunkSize + posOnChunk;
-
-            if (processedPositions.Contains(globalPos))
-                continue;
-
             visualizer.CurrentContext = $"FloodFill_{chunkPos}";
             var maybeArray = AddLinesForChunkViaFloodFill(
                 reallyReferenceChunk,
@@ -90,21 +159,11 @@ internal class MagicWandHelper
                 chunkPos * chunkSize,
                 document.Size,
                 posOnChunk,
-                colorRange, lines, processedPositions);
+                colorRange, lines);
 
             if (maybeArray is null)
                 continue;
-            for (int i = 0; i < chunkSize; i++)
-            {
-                if (chunkPos.Y > 0 && maybeArray[i] == Visited) //Top
-                    positionsToFloodFill.Push((new(chunkPos.X, chunkPos.Y - 1), new(i, chunkSize - 1)));
-                if (chunkPos.Y < imageSizeInChunks.Y - 1 && maybeArray[chunkSize * (chunkSize - 1) + i] == Visited) // Bottom
-                    positionsToFloodFill.Push((new(chunkPos.X, chunkPos.Y + 1), new(i, 0)));
-                if (chunkPos.X > 0 && maybeArray[i * chunkSize] == Visited) // Left
-                    positionsToFloodFill.Push((new(chunkPos.X - 1, chunkPos.Y), new(chunkSize - 1, i)));
-                if (chunkPos.X < imageSizeInChunks.X - 1 && maybeArray[i * chunkSize + (chunkSize - 1)] == Visited) // Right
-                    positionsToFloodFill.Push((new(chunkPos.X + 1, chunkPos.Y), new(0, i)));
-            }
+            positionsToFloodFill.Push(chunkPos, maybeArray);
         }
 
         if (lines.Count > 0)
@@ -117,8 +176,7 @@ internal class MagicWandHelper
         return selection;
     }
 
-    private static void AddLinesForEmptyChunk(Lines lines, VecI chunkPos, VecI imageSize,
-        VecI imageSizeInChunks, int chunkSize)
+    private static void AddLinesForEmptyChunk(Lines lines, VecI chunkPos, VecI imageSize, int chunkSize)
     {
         visualizer.CurrentContext = "EmptyChunk";
 
@@ -144,7 +202,7 @@ internal class MagicWandHelper
     {
         Line previous = default;
         Line? current = firstLine;
-        while (current != null)
+        while (current is not null)
         {
             (previous, current) = ((Line)current, allLines.RemoveLineAt(current.Value.End, current.Value.NormalizedDirection));
         }
@@ -157,7 +215,7 @@ internal class MagicWandHelper
         path.MoveTo(startingLine.Start);
 
         Line? current = startingLine;
-        while (current != null)
+        while (current is not null)
         {
             VecI straightPathEnd = GoStraight(allLines, (Line)current);
             path.LineTo(straightPathEnd);
@@ -170,7 +228,7 @@ internal class MagicWandHelper
         VectorPath selection = new();
 
         Line? current = lines.PopLine();
-        while (current != null)
+        while (current is not null)
         {
             FollowPath(lines, (Line)current, selection);
             current = lines.PopLine();
@@ -179,20 +237,20 @@ internal class MagicWandHelper
         return selection;
     }
 
-    private static unsafe byte[]? AddLinesForChunkViaFloodFill(
+    private static unsafe bool[]? AddLinesForChunkViaFloodFill(
         Chunk referenceChunk,
         int chunkSize,
         VecI chunkOffset,
         VecI documentSize,
         VecI pos,
-        ColorBounds bounds, Lines lines, HashSet<VecI> processedPositions)
+        ColorBounds bounds, Lines lines)
     {
         if (!bounds.IsWithinBounds(referenceChunk.Surface.GetSRGBPixel(pos)))
         {
             return null;
         }
 
-        byte[] pixelStates = new byte[chunkSize * chunkSize];
+        bool[] pixelVisitedStates = new bool[chunkSize * chunkSize];
 
         using var refPixmap = referenceChunk.Surface.DrawingSurface.PeekPixels();
         Half* refArray = (Half*)refPixmap.GetPixels();
@@ -205,36 +263,31 @@ internal class MagicWandHelper
             VecI curPos = toVisit.Pop();
 
             int pixelOffset = curPos.X + curPos.Y * chunkSize;
+            if (pixelVisitedStates[pixelOffset])
+                continue;
+
             VecI globalPos = curPos + chunkOffset;
             Half* refPixel = refArray + pixelOffset * 4;
 
             if (!bounds.IsWithinBounds(refPixel))
-            {
-                processedPositions.Add(globalPos);
                 continue;
-            }
 
-            pixelStates[pixelOffset] = Visited;
+            pixelVisitedStates[pixelOffset] = true;
 
             visualizer.CurrentContext = "AddFillContourLines";
-            AddFillContourLines(chunkSize, chunkOffset, bounds, lines, curPos, pixelStates, pixelOffset, refPixel, toVisit, globalPos, documentSize, processedPositions);
-
-            processedPositions.Add(globalPos);
+            AddFillContourLines(chunkSize, chunkOffset, bounds, lines, curPos, pixelVisitedStates, pixelOffset, refPixel, toVisit, globalPos, documentSize);
         }
 
-        return pixelStates;
+        return pixelVisitedStates;
     }
 
     private static unsafe void AddFillContourLines(int chunkSize, VecI chunkOffset, ColorBounds bounds, Lines lines,
-        VecI curPos, byte[] pixelStates, int pixelOffset, Half* refPixel, Stack<VecI> toVisit, VecI globalPos,
-        VecI documentSize, HashSet<VecI> processedPositions)
+        VecI curPos, bool[] pixelStates, int pixelOffset, Half* refPixel, Stack<VecI> toVisit, VecI globalPos,
+        VecI documentSize)
     {
-
-        if (processedPositions.Contains(globalPos)) return;
-
         // Left pixel
         bool leftEdgePresent = curPos.X == 0 || globalPos.X == 0 || !bounds.IsWithinBounds(refPixel - 4);
-        if (!leftEdgePresent && pixelStates[pixelOffset - 1] != Visited)
+        if (!leftEdgePresent && !pixelStates[pixelOffset - 1])
         {
             toVisit.Push(new(curPos.X - 1, curPos.Y));
         }
@@ -247,7 +300,7 @@ internal class MagicWandHelper
 
         // Right pixel
         bool rightEdgePresent = globalPos.X == documentSize.X - 1 || curPos.X == chunkSize - 1 || !bounds.IsWithinBounds(refPixel + 4);
-        if (!rightEdgePresent && pixelStates[pixelOffset + 1] != Visited)
+        if (!rightEdgePresent && !pixelStates[pixelOffset + 1])
         {
             toVisit.Push(new(curPos.X + 1, curPos.Y));
         }
@@ -260,7 +313,7 @@ internal class MagicWandHelper
 
         // Top pixel
         bool topEdgePresent = curPos.Y == 0 || globalPos.Y == 0 || !bounds.IsWithinBounds(refPixel - 4 * chunkSize);
-        if (!topEdgePresent && pixelStates[pixelOffset - chunkSize] != Visited)
+        if (!topEdgePresent && !pixelStates[pixelOffset - chunkSize])
         {
             toVisit.Push(new(curPos.X, curPos.Y - 1));
         }
@@ -273,7 +326,7 @@ internal class MagicWandHelper
 
         //Bottom pixel
         bool bottomEdgePresent = globalPos.Y == documentSize.Y - 1 || curPos.Y == chunkSize - 1 || !bounds.IsWithinBounds(refPixel + 4 * chunkSize);
-        if (!bottomEdgePresent && pixelStates[pixelOffset + chunkSize] != Visited)
+        if (!bottomEdgePresent && !pixelStates[pixelOffset + chunkSize])
         {
             toVisit.Push(new(curPos.X, curPos.Y + 1));
         }