flood fill

buildconfig/stubs/pygame/draw.pyi

@@ -1,6 +1,6 @@
 from pygame.rect import Rect
 from pygame.surface import Surface
-from typing import overload
+from typing import Union, overload
 
 from ._common import ColorValue, Coordinate, RectValue, Sequence
 
@@ -89,3 +89,8 @@ def aalines(
     closed: bool,
     points: Sequence[Coordinate],
 ) -> Rect: ...
+def flood_fill(
+    surface: Surface,
+    color: Union[ColorValue, Surface],
+    start_point: Coordinate
+) -> Rect: ...

docs/reST/ref/draw.rst

@@ -597,6 +597,33 @@ object around the draw calls (see :func:`pygame.Surface.lock` and
    .. versionchanged:: 2.5.0 ``blend`` argument readded for backcompat, but will always raise a deprecation exception when used
 
    .. ## pygame.draw.aalines ##
+
+.. function:: flood_fill
+
+   | :sl:`fill in a connected area of same-color pixels`
+   | :sg:`flood_fill(surface, color, starting_point) -> Rect`
+   | :sg:`flood_fill(surface, pattern_surface, starting_point) -> Rect`
+
+   Replace the color of a cluster of connected same-color pixels, beginning
+   from the starting point, with a repeating pattern or solid single color
+
+   :param Surface surface: surface to draw on
+   :param color: color to draw with, the alpha value is optional if using a
+      tuple ``(RGB[A])``
+   :type color: Color or string (for :doc:`color_list`) or int or tuple(int, int, int, [int])
+   :param pattern_surface: pattern to fill with, as a surface
+   :param starting_point: starting point as a sequence of 2 ints/floats,
+      e.g. ``(x, y)``
+   :type starting_point: tuple(int or float, int or float) or
+      list(int or float, int or float) or Vector2(int or float, int or float)
+
+   :returns: a rect bounding the changed pixels, if nothing is drawn the
+      bounding rect's position will be the position of the starting point
+      and its width and height will be 0
+   :rtype: Rect
+
+   .. versionadded:: 2.5.0
+   .. ## pygame.draw.flood_fill ##
 
 .. ## pygame.draw ##
 

src_c/doc/draw_doc.h

@@ -10,3 +10,4 @@
 #define DOC_DRAW_LINES "lines(surface, color, closed, points) -> Rect\nlines(surface, color, closed, points, width=1) -> Rect\ndraw multiple contiguous straight line segments"
 #define DOC_DRAW_AALINE "aaline(surface, color, start_pos, end_pos) -> Rect\ndraw a straight antialiased line"
 #define DOC_DRAW_AALINES "aalines(surface, color, closed, points) -> Rect\ndraw multiple contiguous straight antialiased line segments"
+#define DOC_DRAW_FLOODFILL "flood_fill(surface, color, starting_point) -> Rect\nflood_fill(surface, pattern_surface, starting_point) -> Rect\nfill in a connected area of same-color pixels"

src_c/draw.c

@@ -93,6 +93,13 @@ draw_round_rect(SDL_Surface *surf, int x1, int y1, int x2, int y2, int radius,
                 int width, Uint32 color, int top_left, int top_right,
                 int bottom_left, int bottom_right, int *drawn_area);
 
+static int
+flood_fill_inner(SDL_Surface *surf, int x1, int y1, Uint32 new_color,
+                 SDL_Surface *pattern, int *drawn_area);
+
+static void
+unsafe_set_at(SDL_Surface *surf, int x, int y, Uint32 color);
+
 // validation of a draw color
 #define CHECK_LOAD_COLOR(colorobj)                               \
     if (!pg_MappedColorFromObj((colorobj), surf->format, &color, \
@@ -1105,6 +1112,93 @@ rect(PyObject *self, PyObject *args, PyObject *kwargs)
         return pgRect_New4(rect->x, rect->y, 0, 0);
 }
 
+static PyObject *
+flood_fill(PyObject *self, PyObject *arg, PyObject *kwargs)
+{
+    pgSurfaceObject *surfobj;
+    pgSurfaceObject *pat_surfobj;
+    PyObject *colorobj, *start;
+    SDL_Surface *surf = NULL;
+    int startx, starty;
+    Uint32 color;
+    SDL_Surface *pattern = NULL;
+    SDL_bool did_lock = SDL_FALSE;
+    int flood_fill_result;
+
+    int drawn_area[4] = {INT_MAX, INT_MAX, INT_MIN,
+                         INT_MIN}; /* Used to store bounding box values */
+    static char *keywords[] = {"surface", "color", "start_pos", NULL};
+
+    if (!PyArg_ParseTupleAndKeywords(arg, kwargs, "O!OO", keywords,
+                                     &pgSurface_Type, &surfobj, &colorobj,
+                                     &start)) {
+        return NULL; /* Exception already set. */
+    }
+
+    surf = pgSurface_AsSurface(surfobj);
+    SURF_INIT_CHECK(surf)
+
+    if (PG_SURF_BytesPerPixel(surf) <= 0 || PG_SURF_BytesPerPixel(surf) > 4) {
+        return PyErr_Format(PyExc_ValueError,
+                            "unsupported surface bit depth (%d) for drawing",
+                            PG_SURF_BytesPerPixel(surf));
+    }
+
+    if (pgSurface_Check(colorobj)) {
+        pat_surfobj = ((pgSurfaceObject *)colorobj);
+
+        pattern = SDL_ConvertSurface(pat_surfobj->surf, surf->format, 0);
+
+        if (pattern == NULL) {
+            return RAISE(PyExc_RuntimeError, "error converting pattern surf");
+        }
+
+        SDL_SetSurfaceRLE(pattern, SDL_FALSE);
+
+        color = 0;
+    }
+    else {
+        CHECK_LOAD_COLOR(colorobj);
+    }
+
+    if (!pg_TwoIntsFromObj(start, &startx, &starty)) {
+        return RAISE(PyExc_TypeError, "invalid start_pos argument");
+    }
+
+    if (SDL_MUSTLOCK(surf)) {
+        did_lock = SDL_TRUE;
+        if (!pgSurface_Lock(surfobj)) {
+            return RAISE(PyExc_RuntimeError, "error locking surface");
+        }
+    }
+
+    flood_fill_result =
+        flood_fill_inner(surf, startx, starty, color, pattern, drawn_area);
+
+    if (pattern != NULL) {
+        SDL_FreeSurface(pattern);
+    }
+
+    if (did_lock) {
+        if (!pgSurface_Unlock(surfobj)) {
+            return RAISE(PyExc_RuntimeError, "error unlocking surface");
+        }
+    }
+
+    if (flood_fill_result == -1) {
+        return PyErr_NoMemory();
+    }
+
+    /* Compute return rect. */
+    if (drawn_area[0] != INT_MAX && drawn_area[1] != INT_MAX &&
+        drawn_area[2] != INT_MIN && drawn_area[3] != INT_MIN)
+        return pgRect_New4(drawn_area[0], drawn_area[1],
+                           drawn_area[2] - drawn_area[0] + 1,
+                           drawn_area[3] - drawn_area[1] + 1);
+    else
+        return pgRect_New4(startx, starty, 0, 0);
+}
+
 /* Functions used in drawing algorithms */
 
 static void
@@ -1115,6 +1209,33 @@ swap(float *a, float *b)
     *b = temp;
 }
 
+#define WORD_BITS (8 * sizeof(unsigned int))
+
+struct point2d {
+    Uint32 x;
+    Uint32 y;
+};
+
+static inline void
+_bitarray_set(unsigned int *bitarray, size_t idx, SDL_bool value)
+{
+    if (value) {
+        bitarray[idx / WORD_BITS] |= (1 << (idx % WORD_BITS));
+    }
+    else {
+        bitarray[idx / WORD_BITS] &= (~(1) << (idx % WORD_BITS));
+    }
+}
+
+static inline SDL_bool
+_bitarray_get(unsigned int *bitarray, size_t idx)
+{
+    if (bitarray[idx / WORD_BITS] & (1 << (idx % WORD_BITS)))
+        return SDL_TRUE;
+    else
+        return SDL_FALSE;
+}
+
 static int
 compare_int(const void *a, const void *b)
 {
@@ -1775,6 +1896,183 @@ draw_line(SDL_Surface *surf, int x1, int y1, int x2, int y2, Uint32 color,
     set_and_check_rect(surf, x2, y2, color, drawn_area);
 }
 
+#define SURF_GET_AT(p_color, p_surf, p_x, p_y, p_pixels, p_format, p_pix)    \
+    switch (PG_FORMAT_BytesPerPixel(p_format)) {                             \
+        case 1:                                                              \
+            p_color = (Uint32) *                                             \
+                      ((Uint8 *)(p_pixels) + (p_y) * p_surf->pitch + (p_x)); \
+            break;                                                           \
+        case 2:                                                              \
+            p_color =                                                        \
+                (Uint32) *                                                   \
+                ((Uint16 *)((p_pixels) + (p_y) * p_surf->pitch) + (p_x));    \
+            break;                                                           \
+        case 3:                                                              \
+            p_pix =                                                          \
+                ((Uint8 *)(p_pixels + (p_y) * p_surf->pitch) + (p_x) * 3);   \
+            p_color = (SDL_BYTEORDER == SDL_LIL_ENDIAN)                      \
+                          ? (p_pix[0]) + (p_pix[1] << 8) + (p_pix[2] << 16)  \
+                          : (p_pix[2]) + (p_pix[1] << 8) + (p_pix[0] << 16); \
+            break;                                                           \
+        default: /* case 4: */                                               \
+            p_color =                                                        \
+                *((Uint32 *)(p_pixels + (p_y) * p_surf->pitch) + (p_x));     \
+            break;                                                           \
+    }
+
+static int
+flood_fill_inner(SDL_Surface *surf, int x1, int y1, Uint32 new_color,
+                 SDL_Surface *pattern, int *drawn_area)
+{
+    // breadth first flood fill, like graph search
+    SDL_Rect cliprect;
+    size_t mask_idx;
+
+    SDL_GetClipRect(surf, &cliprect);
+    size_t frontier_bufsize = 8, frontier_size = 1, next_frontier_size = 0;
+
+    // Instead of a queue, we use two arrays and swap them between steps.
+    // This makes implementation easier, especially memory management.
+    struct point2d *frontier =
+        malloc(frontier_bufsize * sizeof(struct point2d));
+    if (frontier == NULL) {
+        return -1;
+    }
+
+    struct point2d *frontier_next =
+        malloc(frontier_bufsize * sizeof(struct point2d));
+
+    if (frontier_next == NULL) {
+        free(frontier);
+        return -1;
+    }
+
+    // 2D bitmask for queued nodes
+    // we could check drawn color, but that doesnt work for patterns
+    size_t mask_size = cliprect.w * cliprect.h;
+    unsigned int *mask = calloc((mask_size) / 8 + 1, sizeof(unsigned int));
+
+    if (mask == NULL) {
+        free(frontier);
+        free(frontier_next);
+        return -1;
+    }
+    Uint32 old_color = 0;
+    Uint8 *pix;
+
+    // Von Neumann neighbourhood
+    int VN_X[] = {0, 0, 1, -1};
+    int VN_Y[] = {1, -1, 0, 0};
+
+    if (!(x1 >= cliprect.x && x1 < (cliprect.x + cliprect.w) &&
+          y1 >= cliprect.y && y1 < (cliprect.y + cliprect.h))) {
+        // not an error, but nothing to do here
+        goto flood_fill_finished;
+    }
+
+    SURF_GET_AT(old_color, surf, x1, y1, (Uint8 *)surf->pixels, surf->format,
+                pix);
+
+    if (pattern == NULL && old_color == new_color) {
+        // not an error, but nothing to do here
+        goto flood_fill_finished;
+    }
+
+    frontier[0].x = x1;
+    frontier[0].y = y1;
+
+    // mark starting point already queued
+    mask_idx = (y1 - cliprect.y) * cliprect.w + (x1 - cliprect.x);
+    _bitarray_set(mask, mask_idx, SDL_TRUE);
+
+    while (frontier_size != 0) {
+        next_frontier_size = 0;
+
+        for (size_t i = 0; i < frontier_size; i++) {
+            unsigned int x = frontier[i].x;
+            unsigned int y = frontier[i].y;
+
+            Uint32 current_color = 0;
+
+            SURF_GET_AT(current_color, surf, x, y, (Uint8 *)surf->pixels,
+                        surf->format, pix);
+
+            if (current_color != old_color) {
+                continue;
+            }
+
+            if (pattern != NULL) {
+                SURF_GET_AT(new_color, pattern, x % pattern->w, y % pattern->h,
+                            (Uint8 *)pattern->pixels, pattern->format, pix);
+            }
+
+            // clipping and color mapping have already happened here
+            unsafe_set_at(surf, x, y, new_color);
+            add_pixel_to_drawn_list(x, y, drawn_area);
+
+            for (int n = 0; n < 4; n++) {
+                long nx = x + VN_X[n];
+                long ny = y + VN_Y[n];
+
+                if (!(nx >= cliprect.x && nx < cliprect.x + cliprect.w &&
+                      ny >= cliprect.y && ny < cliprect.y + cliprect.h)) {
+                    continue;
+                }
+
+                mask_idx = (ny - cliprect.y) * cliprect.w + (nx - cliprect.x);
+                if (_bitarray_get(mask, mask_idx))
+                    continue;
+
+                // only queue node once
+                _bitarray_set(mask, mask_idx, SDL_TRUE);
+
+                if (next_frontier_size == frontier_bufsize) {
+                    // grow frontier arrays
+                    struct point2d *old_buf = frontier_next;
+
+                    frontier_bufsize *= 4;
+
+                    frontier_next =
+                        realloc(frontier_next,
+                                frontier_bufsize * sizeof(struct point2d));
+                    if (frontier_next == NULL) {
+                        free(mask);
+                        free(frontier);
+                        free(old_buf);
+                        return -1;
+                    }
+
+                    old_buf = frontier;
+                    frontier = realloc(
+                        frontier, frontier_bufsize * sizeof(struct point2d));
+                    if (frontier == NULL) {
+                        free(old_buf);
+                        free(mask);
+                        free(frontier_next);
+                        return -1;
+                    }
+                }
+
+                frontier_next[next_frontier_size].x = nx;
+                frontier_next[next_frontier_size].y = ny;
+                next_frontier_size++;
+            }
+        }
+        // swap buffers
+        struct point2d *temp_buf;
+        temp_buf = frontier;
+        frontier = frontier_next;
+        frontier_next = temp_buf;
+
+        frontier_size = next_frontier_size;
+    }
+
+flood_fill_finished:
+    free(frontier);
+    free(mask);
+    free(frontier_next);
+    return 0;
+}
 static int
 check_pixel_in_arc(int x, int y, double min_dotproduct, double invsqr_radius1,
                    double invsqr_radius2, double invsqr_inner_radius1,
@@ -3115,6 +3413,8 @@ static PyMethodDef _draw_methods[] = {
      DOC_DRAW_LINES},
     {"ellipse", (PyCFunction)ellipse, METH_VARARGS | METH_KEYWORDS,
      DOC_DRAW_ELLIPSE},
+    {"flood_fill", (PyCFunction)flood_fill, METH_VARARGS | METH_KEYWORDS,
+     DOC_DRAW_FLOODFILL},
     {"arc", (PyCFunction)arc, METH_VARARGS | METH_KEYWORDS, DOC_DRAW_ARC},
     {"circle", (PyCFunction)circle, METH_VARARGS | METH_KEYWORDS,
      DOC_DRAW_CIRCLE},