bitset: Reimplement BIT_FOREACH_IS(SET|CLR)

Eliminate the nested loops and re-implement following a suggestion from
rlibby.

Add some simple regression tests.

Reviewed by:	rlibby, kib
MFC after:	2 weeks
Sponsored by:	The FreeBSD Foundation
Differential Revision:	https://reviews.freebsd.org/D32472

4 files changed, 123 insertions, 7 deletions

diff --git a/share/man/man9/bitset.9 b/share/man/man9/bitset.9
index 5788f09f2465..1a5ec05b01c6 100644
--- a/share/man/man9/bitset.9
+++ b/share/man/man9/bitset.9
@@ -357,6 +357,10 @@ Similarly,
 .Fn BIT_FOREACH_ISCLR
 iterates over all clear bits in
 .Fa bitset .
+In the loop body, the currently indexed bit may be set or cleared.
+However, setting or clearing bits other than the currently indexed
+bit does not guarantee that they will or will not be returned in
+subsequent iterations of the same loop.
 .Pp
 The
 .Fn BIT_COUNT
diff --git a/sys/sys/bitset.h b/sys/sys/bitset.h
index ab6eaf85b8df..1c154d5601ab 100644
--- a/sys/sys/bitset.h
+++ b/sys/sys/bitset.h
@@ -271,12 +271,31 @@
 	__count;							\
 })
 
-/* Non-destructively loop over all set or clear bits in the set. */
-#define	_BIT_FOREACH(_s, i, p, op)					\
-	for (__size_t __i = 0; __i < __bitset_words(_s); __i++)		\
-		for (long __j = op((p)->__bits[__i]), __b = ffsl(__j);	\
-		    (i = (__b - 1) + __i * _BITSET_BITS), __j != 0;	\
-		    __j &= ~(1l << i), __b = ffsl(__j))
+#define	_BIT_FOREACH_ADVANCE(_s, i, p, op) __extension__ ({		\
+	int __found;							\
+	for (;;) {							\
+		if (__bits != 0) {					\
+			int __bit = ffsl(__bits) - 1;			\
+			__bits &= ~(1ul << __bit);			\
+			(i) = __i * _BITSET_BITS + __bit;		\
+			__found = 1;					\
+			break;						\
+		}							\
+		if (++__i == __bitset_words(_s)) {			\
+			__found = 0;					\
+			break;						\
+		}							\
+		__bits = op((p)->__bits[__i]);				\
+	}								\
+	__found != 0;							\
+})
+
+/*
+ * Non-destructively loop over all set or clear bits in the set.
+ */
+#define _BIT_FOREACH(_s, i, p, op)					\
+	for (long __i = -1, __bits = 0;					\
+	    _BIT_FOREACH_ADVANCE(_s, i, p, op); )
 
 #define	BIT_FOREACH_ISSET(_s, i, p)	_BIT_FOREACH(_s, i, p, )
 #define	BIT_FOREACH_ISCLR(_s, i, p)	_BIT_FOREACH(_s, i, p, ~)
diff --git a/tests/sys/sys/Makefile b/tests/sys/sys/Makefile
index 95739c127632..f6c45971d93c 100644
--- a/tests/sys/sys/Makefile
+++ b/tests/sys/sys/Makefile
@@ -4,7 +4,12 @@
 
 TESTSDIR=	${TESTSBASE}/sys/sys
 
-ATF_TESTS_C=	arb_test bitstring_test qmath_test rb_test splay_test
+ATF_TESTS_C=	arb_test \
+		bitset_test \
+		bitstring_test \
+		qmath_test \
+		rb_test \
+		splay_test
 
 .if ${COMPILER_TYPE} == "gcc"
 CFLAGS.bitstring_test=	-fno-strict-overflow
diff --git a/tests/sys/sys/bitset_test.c b/tests/sys/sys/bitset_test.c
new file mode 100644
index 000000000000..781b523dae97
--- /dev/null
+++ b/tests/sys/sys/bitset_test.c
@@ -0,0 +1,88 @@
+/*-
+ * SPDX-License-Identifier: BSD-2-Clause-FreeBSD
+ *
+ * Copyright (c) 2021 The FreeBSD Foundation
+ *
+ * This software was developed by Mark Johnston under sponsorship from
+ * the FreeBSD Foundation.
+ */
+
+#include <sys/types.h>
+#include <sys/_bitset.h>
+#include <sys/bitset.h>
+#include <stdio.h>
+#include <stdlib.h>
+
+#include <atf-c.h>
+
+BITSET_DEFINE(bs256, 256);
+
+ATF_TC_WITHOUT_HEAD(bit_foreach);
+ATF_TC_BODY(bit_foreach, tc)
+{
+	struct bs256 bs0, bs1, bsrand;
+	int setc, clrc, i;
+
+#define	_BIT_FOREACH_COUNT(s, bs) do {					\
+	int prev = -1;							\
+	setc = clrc = 0;						\
+	BIT_FOREACH_ISSET((s), i, (bs)) {				\
+		ATF_REQUIRE_MSG(prev < i, "incorrect bit ordering");	\
+		ATF_REQUIRE_MSG(BIT_ISSET((s), i, (bs)),		\
+		    "bit %d is not set", i);				\
+		setc++;							\
+		prev = i;						\
+	}								\
+	prev = -1;							\
+	BIT_FOREACH_ISCLR((s), i, (bs)) {				\
+		ATF_REQUIRE_MSG(prev < i, "incorrect bit ordering");	\
+		ATF_REQUIRE_MSG(!BIT_ISSET((s), i, (bs)),		\
+		    "bit %d is set", i);				\
+		clrc++;							\
+		prev = i;						\
+	}								\
+} while (0)
+
+	/*
+	 * Create several bitsets, and for each one count the number
+	 * of set and clear bits and make sure they match what we expect.
+	 */
+
+	BIT_FILL(256, &bs1);
+	_BIT_FOREACH_COUNT(256, &bs1);
+	ATF_REQUIRE_MSG(setc == 256, "incorrect set count %d", setc);
+	ATF_REQUIRE_MSG(clrc == 0, "incorrect clear count %d", clrc);
+
+	BIT_ZERO(256, &bs0);
+	_BIT_FOREACH_COUNT(256, &bs0);
+	ATF_REQUIRE_MSG(setc == 0, "incorrect set count %d", setc);
+	ATF_REQUIRE_MSG(clrc == 256, "incorrect clear count %d", clrc);
+
+	BIT_ZERO(256, &bsrand);
+	for (i = 0; i < 256; i++)
+		if (random() % 2 != 0)
+			BIT_SET(256, i, &bsrand);
+	_BIT_FOREACH_COUNT(256, &bsrand);
+	ATF_REQUIRE_MSG(setc + clrc == 256, "incorrect counts %d, %d",
+	    setc, clrc);
+
+	/*
+	 * Try to verify that we can safely clear bits in the set while
+	 * iterating.
+	 */
+	BIT_FOREACH_ISSET(256, i, &bsrand) {
+		ATF_REQUIRE(setc-- > 0);
+		BIT_CLR(256, i, &bsrand);
+	}
+	_BIT_FOREACH_COUNT(256, &bsrand);
+	ATF_REQUIRE_MSG(setc == 0, "incorrect set count %d", setc);
+	ATF_REQUIRE_MSG(clrc == 256, "incorrect clear count %d", clrc);
+
+#undef _BIT_FOREACH_COUNT
+}
+
+ATF_TP_ADD_TCS(tp)
+{
+	ATF_TP_ADD_TC(tp, bit_foreach);
+	return (atf_no_error());
+}