util/idalloc: optimize foreach by tracking the greatest non-zero element

also add util_idalloc_foreach_no_zero_safe.

Reviewed-by: Adam Jackson <ajax@redhat.com>
Reviewed-by: Pierre-Eric Pelloux-Prayer <pierre-eric.pelloux-prayer@amd.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/27494>

src/util/u_idalloc.c

@@ -74,6 +74,7 @@ util_idalloc_alloc(struct util_idalloc *buf)
       unsigned bit = ffs(~buf->data[i]) - 1;
       buf->data[i] |= 1u << bit;
       buf->lowest_free_idx = i;
+      buf->num_set_elements = MAX2(buf->num_set_elements, i + 1);
       return i * 32 + bit;
    }
 
@@ -82,6 +83,7 @@ util_idalloc_alloc(struct util_idalloc *buf)
 
    buf->lowest_free_idx = num_elements;
    buf->data[num_elements] |= 1;
+   buf->num_set_elements = MAX2(buf->num_set_elements, num_elements + 1);
    return num_elements * 32;
 }
 
@@ -134,6 +136,8 @@ ret:
    if (buf->lowest_free_idx == base)
       buf->lowest_free_idx = base + num / 32;
 
+   buf->num_set_elements = MAX2(buf->num_set_elements, base + num_alloc);
+
    /* Validate this algorithm. */
    for (unsigned i = 0; i < num; i++)
       assert(util_idalloc_exists(buf, base * 32 + i));
@@ -151,14 +155,23 @@ util_idalloc_free(struct util_idalloc *buf, unsigned id)
 
    buf->lowest_free_idx = MIN2(idx, buf->lowest_free_idx);
    buf->data[idx] &= ~(1 << (id % 32));
+
+   /* Decrease num_used to the last used element + 1. */
+   if (buf->num_set_elements == idx + 1) {
+      while (buf->num_set_elements > 0 && !buf->data[buf->num_set_elements - 1])
+         buf->num_set_elements--;
+   }
 }
 
 void
 util_idalloc_reserve(struct util_idalloc *buf, unsigned id)
 {
-   if (id / 32 >= buf->num_elements)
-      util_idalloc_resize(buf, (id / 32 + 1) * 2);
-   buf->data[id / 32] |= 1u << (id % 32);
+   unsigned idx = id / 32;
+
+   if (idx >= buf->num_elements)
+      util_idalloc_resize(buf, (idx + 1) * 2);
+   buf->data[idx] |= 1u << (id % 32);
+   buf->num_set_elements = MAX2(buf->num_set_elements, idx + 1);
 }
 
 void

src/util/u_idalloc.h

@@ -37,6 +37,7 @@
 #include <inttypes.h>
 #include <stdbool.h>
 #include "simple_mtx.h"
+#include "bitscan.h"
 
 #ifdef __cplusplus
 extern "C" {
@@ -45,7 +46,8 @@ extern "C" {
 struct util_idalloc
 {
    uint32_t *data;
-   unsigned num_elements;    /* number of allocated elements of "data" */
+   unsigned num_elements;     /* number of allocated elements of "data" */
+   unsigned num_set_elements; /* the last non-zero element of "data" + 1 */
    unsigned lowest_free_idx;
 };
 
@@ -70,17 +72,25 @@ util_idalloc_reserve(struct util_idalloc *buf, unsigned id);
 static inline bool
 util_idalloc_exists(struct util_idalloc *buf, unsigned id)
 {
-   return id / 32 < buf->num_elements &&
+   return id / 32 < buf->num_set_elements &&
           buf->data[id / 32] & BITFIELD_BIT(id % 32);
 }
 
 #define util_idalloc_foreach(buf, id) \
-   for (uint32_t i = 0, mask = (buf)->num_elements ? (buf)->data[0] : 0, id, \
-                 count = (buf)->num_elements; \
+   for (uint32_t i = 0, mask = (buf)->num_set_elements ? (buf)->data[0] : 0, id, \
+                 count = (buf)->num_used; \
         i < count; mask = ++i < count ? (buf)->data[i] : 0) \
       while (mask) \
          if ((id = i * 32 + u_bit_scan(&mask)), true)
 
+/* This allows freeing IDs while iterating, excluding ID=0. */
+#define util_idalloc_foreach_no_zero_safe(buf, id) \
+   for (uint32_t i = 0, bit, id, count = (buf)->num_set_elements, \
+         mask = count ? (buf)->data[0] & ~0x1 : 0; \
+        i < count; mask = ++i < count ? (buf)->data[i] : 0) \
+      while (mask) \
+         if ((bit = u_bit_scan(&mask), id = i * 32 + bit), \
+             (buf)->data[i] & BITFIELD_BIT(bit))
 
 /* Thread-safe variant. */
 struct util_idalloc_mt {