bitset: add BIT_FFS_AT() for finding the first bit set greater than a start bit

Reviewed by:	kib
Approved by:	scottl (implicit)
MFC after:	1 week
Sponsored by:	Ampere Computing, Inc.
Differential Revision:	https://reviews.freebsd.org/D26128

2 files changed, 33 insertions, 6 deletions

diff --git a/share/man/man9/bitset.9 b/share/man/man9/bitset.9
index 4aabb81e7605..eaa8d32a846d 100644
--- a/share/man/man9/bitset.9
+++ b/share/man/man9/bitset.9
@@ -24,7 +24,7 @@
 .\"
 .\" $FreeBSD$
 .\"
-.Dd December 12, 2019
+.Dd August 25, 2020
 .Dt BITSET 9
 .Os
 .Sh NAME
@@ -43,6 +43,7 @@
 .Nm BIT_EMPTY ,
 .Nm BIT_ISFULLSET ,
 .Nm BIT_FFS ,
+.Nm BIT_FFS_AT ,
 .Nm BIT_FLS ,
 .Nm BIT_COUNT ,
 .Nm BIT_SUBSET ,
@@ -86,6 +87,8 @@
 .Ft int
 .Fn BIT_FFS "const SETSIZE" "struct STRUCTNAME *bitset"
 .Ft int
+.Fn BIT_FFS_AT "const SETSIZE" "struct STRUCTNAME *bitset" "int start"
+.Ft int
 .Fn BIT_FLS "const SETSIZE" "struct STRUCTNAME *bitset"
 .Ft int
 .Fn BIT_COUNT "const SETSIZE" "struct STRUCTNAME *bitset"
@@ -285,6 +288,18 @@ index parameter to any other
 macro, you must subtract one from the result.
 .Pp
 The
+.Fn BIT_FFS_AT
+macro returns the 1-index of the first (lowest) set bit in
+.Fa bitset ,
+which is greater than the given 1-indexed
+.Fa start ,
+or zero if no bits in
+.Fa bitset
+greater than
+.Fa start
+are set.
+.Pp
+The
 .Fn BIT_FLS
 macro returns the 1-index of the last (highest) set bit in
 .Fa bitset ,
@@ -518,7 +533,8 @@ argument to all of these macros must match the value given to
 .Fn BITSET_DEFINE .
 .Pp
 Unlike every other reference to individual set members, which are zero-indexed,
-.Fn BIT_FFS
+.Fn BIT_FFS ,
+.Fn BIT_FFS_AT
 and
 .Fn BIT_FLS
 return a one-indexed result (or zero if the set is empty).
diff --git a/sys/sys/bitset.h b/sys/sys/bitset.h
index 9f9ec9ea791e..578cfc621bfb 100644
--- a/sys/sys/bitset.h
+++ b/sys/sys/bitset.h
@@ -207,21 +207,32 @@
 		    (f)->__bits[__i]);					\
 } while (0)
 
-#define	BIT_FFS(_s, p) __extension__ ({					\
+/*
+ * Note that `start` and the returned value from BIT_FFS_AT are
+ * 1-based bit indices.
+ */
+#define	BIT_FFS_AT(_s, p, start) __extension__ ({			\
 	__size_t __i;							\
+	long __mask;							\
 	int __bit;							\
 									\
+	__mask = ~0UL << ((start) % _BITSET_BITS);			\
 	__bit = 0;							\
-	for (__i = 0; __i < __bitset_words((_s)); __i++) {		\
-		if ((p)->__bits[__i] != 0) {				\
-			__bit = ffsl((p)->__bits[__i]);			\
+	for (__i = __bitset_word((_s), (start));			\
+	    __i < __bitset_words((_s));					\
+	    __i++) {							\
+		if (((p)->__bits[__i] & __mask) != 0) {			\
+			__bit = ffsl((p)->__bits[__i] & __mask);	\
 			__bit += __i * _BITSET_BITS;			\
 			break;						\
 		}							\
+		__mask = ~0UL;						\
 	}								\
 	__bit;								\
 })
 
+#define	BIT_FFS(_s, p) BIT_FFS_AT((_s), (p), 0)
+
 #define	BIT_FLS(_s, p) __extension__ ({					\
 	__size_t __i;							\
 	int __bit;							\