[PATCH 01/12] bitmap: add find_nth_bit_from()

[Yury Norov <yury.norov@gmail.com>]

Similarly to find_next_bit(), add a _from flavor for find_nth_bit(). It's
useful when traversing bitmaps in a loop because it allows to not count
bits from the beginning every time.

The test is added as well.

Signed-off-by: Yury Norov <yury.norov@gmail.com>
---
 include/linux/find.h | 37 +++++++++++++++++++++++++++++++++++++
 lib/find_bit.c       | 29 +++++++++++++++++++++++++++++
 lib/test_bitmap.c    | 10 +++++++++-
 3 files changed, 75 insertions(+), 1 deletion(-)

diff --git a/include/linux/find.h b/include/linux/find.h
index 5e4f39ef2e72..f7fb88405201 100644
--- a/include/linux/find.h
+++ b/include/linux/find.h
@@ -27,6 +27,8 @@ unsigned long __find_nth_andnot_bit(const unsigned long *addr1, const unsigned l
 unsigned long __find_nth_and_andnot_bit(const unsigned long *addr1, const unsigned long *addr2,
 					const unsigned long *addr3, unsigned long size,
 					unsigned long n);
+unsigned long __find_nth_bit_from(const unsigned long *addr, unsigned long size,
+					unsigned long start, unsigned long nth);
 extern unsigned long _find_first_and_bit(const unsigned long *addr1,
 					 const unsigned long *addr2, unsigned long size);
 extern unsigned long _find_first_zero_bit(const unsigned long *addr, unsigned long size);
@@ -237,6 +239,41 @@ unsigned long find_nth_bit(const unsigned long *addr, unsigned long size, unsign
 	return __find_nth_bit(addr, size, n);
 }
 
+/**
+ * find_nth_bit_from - find N'th set bit in a memory region starting at @off
+ * @addr: The address to start the search at
+ * @size: The maximum number of bits to search
+ * @off: The offset to start search
+ * @n: The number of set bit, which position is needed, counting from 0
+ *
+ * The following is semantically equivalent:
+ *	 idx = find_nth_bit_from(addr, size, off, 0);
+ *	 idx = find_next_bit(addr, size, off);
+ *
+ * Return: the bit number of the N'th set bit.
+ * If no such, returns @size.
+ */
+static __always_inline
+unsigned long find_nth_bit_from(const unsigned long *addr, unsigned long size,
+				unsigned long start, unsigned long n)
+{
+	if (n >= size - start)
+		return size;
+
+	if (small_const_nbits(size - start) && size / BITS_PER_LONG == start / BITS_PER_LONG) {
+		unsigned long val, idx = start / BITS_PER_LONG;
+
+		val =  addr[idx] & GENMASK(size - 1, start);
+		if (val == 0)
+			return size;
+
+		val = idx * BITS_PER_LONG + fns(val, n);
+		return val < size ? val : size;
+	}
+
+	return __find_nth_bit_from(addr, size, start, n);
+}
+
 /**
  * find_nth_and_bit - find N'th set bit in 2 memory regions
  * @addr1: The 1st address to start the search at
diff --git a/lib/find_bit.c b/lib/find_bit.c
index 32f99e9a670e..49959b51df8c 100644
--- a/lib/find_bit.c
+++ b/lib/find_bit.c
@@ -164,6 +164,35 @@ unsigned long __find_nth_and_andnot_bit(const unsigned long *addr1,
 }
 EXPORT_SYMBOL(__find_nth_and_andnot_bit);
 
+#define FIND_NTH_BIT_FROM(FETCH, size, start, nth)				\
+({										\
+	unsigned long mask, idx, tmp, w, sz = (size), __start = (start), n = (nth);\
+										\
+	if (unlikely(__start >= sz || n > sz - __start))			\
+		goto out;							\
+										\
+	mask = BITMAP_FIRST_WORD_MASK(__start);					\
+	idx = __start / BITS_PER_LONG;						\
+										\
+	for (tmp = (FETCH) & mask; (w = hweight_long(tmp)) <= n; tmp = (FETCH)) {\
+		if ((idx + 1) * BITS_PER_LONG >= sz)				\
+			goto out;						\
+		idx++;								\
+		n -= w;								\
+	}									\
+										\
+	sz = min(idx * BITS_PER_LONG + fns(tmp, n), sz);			\
+out:										\
+	sz;									\
+})
+
+unsigned long __find_nth_bit_from(const unsigned long *addr, unsigned long size,
+					unsigned long start, unsigned long nth)
+{
+	return FIND_NTH_BIT_FROM(addr[idx], size, start, nth);
+}
+EXPORT_SYMBOL(__find_nth_bit_from);
+
 #ifndef find_next_and_bit
 unsigned long _find_next_and_bit(const unsigned long *addr1, const unsigned long *addr2,
 					unsigned long nbits, unsigned long start)
diff --git a/lib/test_bitmap.c b/lib/test_bitmap.c
index def7d2f9bd14..3248ed58a817 100644
--- a/lib/test_bitmap.c
+++ b/lib/test_bitmap.c
@@ -223,7 +223,7 @@ static void __init test_zero_clear(void)
 
 static void __init test_find_nth_bit(void)
 {
-	unsigned long b, bit, cnt = 0;
+	unsigned long i, b, bit, cnt = 0;
 	DECLARE_BITMAP(bmap, 64 * 3);
 
 	bitmap_zero(bmap, 64 * 3);
@@ -260,6 +260,14 @@ static void __init test_find_nth_bit(void)
 		b = find_nth_bit(exp1, EXP1_IN_BITS, cnt++);
 		expect_eq_uint(b, bit);
 	}
+
+	for (i = 0; i < EXP1_IN_BITS; i++) {
+		bit = i; cnt = 0;
+		for_each_set_bit_from(bit, exp1, EXP1_IN_BITS) {
+			b = find_nth_bit_from(exp1, EXP1_IN_BITS, i, cnt++);
+			expect_eq_uint(b, bit);
+		}
+	}
 }
 
 static void __init test_fill_set(void)
-- 
2.39.2