Implement bitmap_ffu(), which finds the first unset bit.

6 files changed, 136 insertions, 25 deletions

diff --git a/include/jemalloc/internal/bitmap_externs.h b/include/jemalloc/internal/bitmap_externs.h
index 4df63eb..034a4e6 100644
--- a/include/jemalloc/internal/bitmap_externs.h
+++ b/include/jemalloc/internal/bitmap_externs.h
@@ -2,7 +2,7 @@
 #define JEMALLOC_INTERNAL_BITMAP_EXTERNS_H
 
 void	bitmap_info_init(bitmap_info_t *binfo, size_t nbits);
-void	bitmap_init(bitmap_t *bitmap, const bitmap_info_t *binfo);
+void	bitmap_init(bitmap_t *bitmap, const bitmap_info_t *binfo, bool fill);
 size_t	bitmap_size(const bitmap_info_t *binfo);
 
 #endif /* JEMALLOC_INTERNAL_BITMAP_EXTERNS_H */
diff --git a/include/jemalloc/internal/bitmap_inlines.h b/include/jemalloc/internal/bitmap_inlines.h
index df582bb..07166ba 100644
--- a/include/jemalloc/internal/bitmap_inlines.h
+++ b/include/jemalloc/internal/bitmap_inlines.h
@@ -2,11 +2,13 @@
 #define JEMALLOC_INTERNAL_BITMAP_INLINES_H
 
 #ifndef JEMALLOC_ENABLE_INLINE
-bool	bitmap_full(bitmap_t *bitmap, const bitmap_info_t *binfo);
-bool	bitmap_get(bitmap_t *bitmap, const bitmap_info_t *binfo, size_t bit);
-void	bitmap_set(bitmap_t *bitmap, const bitmap_info_t *binfo, size_t bit);
-size_t	bitmap_sfu(bitmap_t *bitmap, const bitmap_info_t *binfo);
-void	bitmap_unset(bitmap_t *bitmap, const bitmap_info_t *binfo, size_t bit);
+bool bitmap_full(bitmap_t *bitmap, const bitmap_info_t *binfo);
+bool bitmap_get(bitmap_t *bitmap, const bitmap_info_t *binfo, size_t bit);
+void bitmap_set(bitmap_t *bitmap, const bitmap_info_t *binfo, size_t bit);
+size_t bitmap_ffu(const bitmap_t *bitmap, const bitmap_info_t *binfo,
+    size_t min_bit);
+size_t bitmap_sfu(bitmap_t *bitmap, const bitmap_info_t *binfo);
+void bitmap_unset(bitmap_t *bitmap, const bitmap_info_t *binfo, size_t bit);
 #endif
 
 #if (defined(JEMALLOC_ENABLE_INLINE) || defined(JEMALLOC_BITMAP_C_))
@@ -75,6 +77,64 @@ bitmap_set(bitmap_t *bitmap, const bitmap_info_t *binfo, size_t bit) {
 #endif
 }
 
+/* ffu: find first unset >= bit. */
+JEMALLOC_INLINE size_t
+bitmap_ffu(const bitmap_t *bitmap, const bitmap_info_t *binfo, size_t min_bit) {
+	assert(min_bit < binfo->nbits);
+
+#ifdef BITMAP_USE_TREE
+	unsigned level = binfo->nlevels - 1;
+	size_t lg_bits_per_group = (LG_BITMAP_GROUP_NBITS * (level+1));
+	size_t bits_per_group = 1LU << lg_bits_per_group;
+	size_t bits_per_group_mask = bits_per_group - 1;
+	unsigned group_nmask = (min_bit & bits_per_group_mask) >> (level *
+	    LG_BITMAP_GROUP_NBITS);
+	bitmap_t group_mask = ~((1LU << group_nmask) - 1);
+	bitmap_t group = bitmap[binfo->levels[level].group_offset] & group_mask;
+	if (group == 0LU) {
+		return binfo->nbits;
+	}
+	size_t bit = ffs_lu(group) - 1;
+
+	while (level > 0) {
+		level--;
+
+		lg_bits_per_group = (LG_BITMAP_GROUP_NBITS * (level+1));
+		bits_per_group = 1LU << lg_bits_per_group;
+		bits_per_group_mask = bits_per_group - 1;
+
+		group = bitmap[binfo->levels[level].group_offset + bit];
+		size_t cur_base = bit << lg_bits_per_group;
+		if (cur_base < min_bit) {
+			group_nmask = (min_bit & bits_per_group_mask) >> (level
+			    * LG_BITMAP_GROUP_NBITS);
+			group_mask = ~((1LU << group_nmask) - 1);
+			group &= group_mask;
+		}
+		if (group == 0LU) {
+			return binfo->nbits;
+		}
+		bit = (bit << LG_BITMAP_GROUP_NBITS) + (ffs_lu(group) - 1);
+	}
+	assert(bit < binfo->nbits);
+	return bit;
+#else
+	size_t i = min_bit >> LG_BITMAP_GROUP_NBITS;
+	bitmap_t g = bitmap[i] & ~((1LU << (min_bit & BITMAP_GROUP_NBITS_MASK))
+	    - 1);
+	size_t bit;
+	do {
+		bit = ffs_lu(g);
+		if (bit != 0) {
+			return (i << LG_BITMAP_GROUP_NBITS) + (bit - 1);
+		}
+		i++;
+		g = bitmap[i];
+	} while (i < binfo->ngroups);
+	return binfo->nbits;
+#endif
+}
+
 /* sfu: set first unset. */
 JEMALLOC_INLINE size_t
 bitmap_sfu(bitmap_t *bitmap, const bitmap_info_t *binfo) {
diff --git a/include/jemalloc/internal/private_symbols.txt b/include/jemalloc/internal/private_symbols.txt
index 1af1f91..d22cd87 100644
--- a/include/jemalloc/internal/private_symbols.txt
+++ b/include/jemalloc/internal/private_symbols.txt
@@ -94,6 +94,7 @@ base_postfork_child
 base_postfork_parent
 base_prefork
 base_stats_get
+bitmap_ffu
 bitmap_full
 bitmap_get
 bitmap_info_init
diff --git a/src/arena.c b/src/arena.c
index 519111e..b0913c3 100644
--- a/src/arena.c
+++ b/src/arena.c
@@ -1210,7 +1210,7 @@ arena_slab_alloc(tsdn_t *tsdn, arena_t *arena, szind_t binind,
 	/* Initialize slab internals. */
 	arena_slab_data_t *slab_data = extent_slab_data_get(slab);
 	slab_data->nfree = bin_info->nregs;
-	bitmap_init(slab_data->bitmap, &bin_info->bitmap_info);
+	bitmap_init(slab_data->bitmap, &bin_info->bitmap_info, false);
 
 	arena_nactive_add(arena, extent_size_get(slab) >> LG_PAGE);
 
diff --git a/src/bitmap.c b/src/bitmap.c
index 17efb73..81d2a6d 100644
--- a/src/bitmap.c
+++ b/src/bitmap.c
@@ -39,16 +39,26 @@ bitmap_info_ngroups(const bitmap_info_t *binfo) {
 }
 
 void
-bitmap_init(bitmap_t *bitmap, const bitmap_info_t *binfo) {
+bitmap_init(bitmap_t *bitmap, const bitmap_info_t *binfo, bool fill) {
 	size_t extra;
 	unsigned i;
 
 	/*
 	 * Bits are actually inverted with regard to the external bitmap
-	 * interface, so the bitmap starts out with all 1 bits, except for
-	 * trailing unused bits (if any).  Note that each group uses bit 0 to
-	 * correspond to the first logical bit in the group, so extra bits
-	 * are the most significant bits of the last group.
+	 * interface.
+	 */
+
+	if (fill) {
+		/* The "filled" bitmap starts out with all 0 bits. */
+		memset(bitmap, 0, bitmap_size(binfo));
+		return;
+	}
+
+	/*
+	 * The "empty" bitmap starts out with all 1 bits, except for trailing
+	 * unused bits (if any).  Note that each group uses bit 0 to correspond
+	 * to the first logical bit in the group, so extra bits are the most
+	 * significant bits of the last group.
 	 */
 	memset(bitmap, 0xffU, bitmap_size(binfo));
 	extra = (BITMAP_GROUP_NBITS - (binfo->nbits & BITMAP_GROUP_NBITS_MASK))
@@ -84,9 +94,14 @@ bitmap_info_ngroups(const bitmap_info_t *binfo) {
 }
 
 void
-bitmap_init(bitmap_t *bitmap, const bitmap_info_t *binfo) {
+bitmap_init(bitmap_t *bitmap, const bitmap_info_t *binfo, bool fill) {
 	size_t extra;
 
+	if (fill) {
+		memset(bitmap, 0, bitmap_size(binfo));
+		return;
+	}
+
 	memset(bitmap, 0xffU, bitmap_size(binfo));
 	extra = (BITMAP_GROUP_NBITS - (binfo->nbits & BITMAP_GROUP_NBITS_MASK))
 	    & BITMAP_GROUP_NBITS_MASK;
diff --git a/test/unit/bitmap.c b/test/unit/bitmap.c
index ca65760..92a07de 100644
--- a/test/unit/bitmap.c
+++ b/test/unit/bitmap.c
@@ -171,12 +171,18 @@ test_bitmap_init_body(const bitmap_info_t *binfo, size_t nbits) {
 	size_t i;
 	bitmap_t *bitmap = (bitmap_t *)malloc(bitmap_size(binfo));
 	assert_ptr_not_null(bitmap, "Unexpected malloc() failure");
-	bitmap_init(bitmap, binfo);
 
+	bitmap_init(bitmap, binfo, false);
 	for (i = 0; i < nbits; i++) {
 		assert_false(bitmap_get(bitmap, binfo, i),
 		    "Bit should be unset");
 	}
+
+	bitmap_init(bitmap, binfo, true);
+	for (i = 0; i < nbits; i++) {
+		assert_true(bitmap_get(bitmap, binfo, i), "Bit should be set");
+	}
+
 	free(bitmap);
 }
 
@@ -202,7 +208,7 @@ test_bitmap_set_body(const bitmap_info_t *binfo, size_t nbits) {
 	size_t i;
 	bitmap_t *bitmap = (bitmap_t *)malloc(bitmap_size(binfo));
 	assert_ptr_not_null(bitmap, "Unexpected malloc() failure");
-	bitmap_init(bitmap, binfo);
+	bitmap_init(bitmap, binfo, false);
 
 	for (i = 0; i < nbits; i++) {
 		bitmap_set(bitmap, binfo, i);
@@ -233,7 +239,7 @@ test_bitmap_unset_body(const bitmap_info_t *binfo, size_t nbits) {
 	size_t i;
 	bitmap_t *bitmap = (bitmap_t *)malloc(bitmap_size(binfo));
 	assert_ptr_not_null(bitmap, "Unexpected malloc() failure");
-	bitmap_init(bitmap, binfo);
+	bitmap_init(bitmap, binfo, false);
 
 	for (i = 0; i < nbits; i++) {
 		bitmap_set(bitmap, binfo, i);
@@ -268,14 +274,22 @@ TEST_END
 
 static void
 test_bitmap_sfu_body(const bitmap_info_t *binfo, size_t nbits) {
-	size_t i;
 	bitmap_t *bitmap = (bitmap_t *)malloc(bitmap_size(binfo));
 	assert_ptr_not_null(bitmap, "Unexpected malloc() failure");
-	bitmap_init(bitmap, binfo);
+	bitmap_init(bitmap, binfo, false);
 
 	/* Iteratively set bits starting at the beginning. */
-	for (i = 0; i < nbits; i++) {
-		assert_zd_eq(bitmap_sfu(bitmap, binfo), i,
+	for (size_t i = 0; i < nbits; i++) {
+		assert_zu_eq(bitmap_ffu(bitmap, binfo, 0), i,
+		    "First unset bit should be just after previous first unset "
+		    "bit");
+		assert_zu_eq(bitmap_ffu(bitmap, binfo, (i > 0) ? i-1 : i), i,
+		    "First unset bit should be just after previous first unset "
+		    "bit");
+		assert_zu_eq(bitmap_ffu(bitmap, binfo, i), i,
+		    "First unset bit should be just after previous first unset "
+		    "bit");
+		assert_zu_eq(bitmap_sfu(bitmap, binfo), i,
 		    "First unset bit should be just after previous first unset "
 		    "bit");
 	}
@@ -285,9 +299,15 @@ test_bitmap_sfu_body(const bitmap_info_t *binfo, size_t nbits) {
 	 * Iteratively unset bits starting at the end, and verify that
 	 * bitmap_sfu() reaches the unset bits.
 	 */
-	for (i = nbits - 1; i < nbits; i--) { /* (nbits..0] */
+	for (size_t i = nbits - 1; i < nbits; i--) { /* (nbits..0] */
 		bitmap_unset(bitmap, binfo, i);
-		assert_zd_eq(bitmap_sfu(bitmap, binfo), i,
+		assert_zu_eq(bitmap_ffu(bitmap, binfo, 0), i,
+		    "First unset bit should the bit previously unset");
+		assert_zu_eq(bitmap_ffu(bitmap, binfo, (i > 0) ? i-1 : i), i,
+		    "First unset bit should the bit previously unset");
+		assert_zu_eq(bitmap_ffu(bitmap, binfo, i), i,
+		    "First unset bit should the bit previously unset");
+		assert_zu_eq(bitmap_sfu(bitmap, binfo), i,
 		    "First unset bit should the bit previously unset");
 		bitmap_unset(bitmap, binfo, i);
 	}
@@ -297,14 +317,29 @@ test_bitmap_sfu_body(const bitmap_info_t *binfo, size_t nbits) {
 	 * Iteratively set bits starting at the beginning, and verify that
 	 * bitmap_sfu() looks past them.
 	 */
-	for (i = 1; i < nbits; i++) {
+	for (size_t i = 1; i < nbits; i++) {
 		bitmap_set(bitmap, binfo, i - 1);
-		assert_zd_eq(bitmap_sfu(bitmap, binfo), i,
+		assert_zu_eq(bitmap_ffu(bitmap, binfo, 0), i,
+		    "First unset bit should be just after the bit previously "
+		    "set");
+		assert_zu_eq(bitmap_ffu(bitmap, binfo, (i > 0) ? i-1 : i), i,
+		    "First unset bit should be just after the bit previously "
+		    "set");
+		assert_zu_eq(bitmap_ffu(bitmap, binfo, i), i,
+		    "First unset bit should be just after the bit previously "
+		    "set");
+		assert_zu_eq(bitmap_sfu(bitmap, binfo), i,
 		    "First unset bit should be just after the bit previously "
 		    "set");
 		bitmap_unset(bitmap, binfo, i);
 	}
-	assert_zd_eq(bitmap_sfu(bitmap, binfo), nbits - 1,
+	assert_zu_eq(bitmap_ffu(bitmap, binfo, 0), nbits - 1,
+	    "First unset bit should be the last bit");
+	assert_zu_eq(bitmap_ffu(bitmap, binfo, (nbits > 1) ? nbits-2 : nbits-1),
+	    nbits - 1, "First unset bit should be the last bit");
+	assert_zu_eq(bitmap_ffu(bitmap, binfo, nbits - 1), nbits - 1,
+	    "First unset bit should be the last bit");
+	assert_zu_eq(bitmap_sfu(bitmap, binfo), nbits - 1,
 	    "First unset bit should be the last bit");
 	assert_true(bitmap_full(bitmap, binfo), "All bits should be set");
 	free(bitmap);