[PATCH 02/12] bitmap: add bitmap_weight_from()

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

Add a _from flavor for bitmap_weight(). It's useful when traversing
bitmaps in a loop to not count bits from the beginning every time.

The test for bitmap_weight{,_from}() is added as well.

Signed-off-by: Yury Norov <yury.norov@gmail.com>
---
 include/linux/bitmap.h | 14 ++++++++++++++
 lib/bitmap.c           | 25 +++++++++++++++++++++++++
 lib/test_bitmap.c      | 18 ++++++++++++++++++
 3 files changed, 57 insertions(+)

diff --git a/include/linux/bitmap.h b/include/linux/bitmap.h
index 692d0a1dbe92..6acbdd2abd0c 100644
--- a/include/linux/bitmap.h
+++ b/include/linux/bitmap.h
@@ -168,6 +168,8 @@ bool __bitmap_subset(const unsigned long *bitmap1,
 unsigned int __bitmap_weight(const unsigned long *bitmap, unsigned int nbits);
 unsigned int __bitmap_weight_and(const unsigned long *bitmap1,
 				 const unsigned long *bitmap2, unsigned int nbits);
+unsigned int __bitmap_weight_from(const unsigned long *bitmap,
+					unsigned int bits, unsigned int start);
 void __bitmap_set(unsigned long *map, unsigned int start, int len);
 void __bitmap_clear(unsigned long *map, unsigned int start, int len);
 
@@ -446,6 +448,18 @@ unsigned long bitmap_weight_and(const unsigned long *src1,
 	return __bitmap_weight_and(src1, src2, nbits);
 }
 
+static __always_inline
+unsigned int bitmap_weight_from(const unsigned long *src, unsigned int nbits, unsigned int start)
+{
+	if (unlikely(start >= nbits))
+		return 0;
+
+	if (small_const_nbits(nbits - start) && nbits / BITS_PER_LONG == start / BITS_PER_LONG)
+		return hweight_long(*src & GENMASK(nbits-1, start));
+
+	return __bitmap_weight_from(src, nbits, start);
+}
+
 static __always_inline void bitmap_set(unsigned long *map, unsigned int start,
 		unsigned int nbits)
 {
diff --git a/lib/bitmap.c b/lib/bitmap.c
index cf77d56cf223..65c64911c92f 100644
--- a/lib/bitmap.c
+++ b/lib/bitmap.c
@@ -345,6 +345,24 @@ EXPORT_SYMBOL(__bitmap_subset);
 	w;									\
 })
 
+#define BITMAP_WEIGHT_FROM(FETCH, bits, start)					\
+({										\
+	unsigned long __bits = (bits), val, idx, w = 0, __start = (start), mask;\
+										\
+	mask = BITMAP_FIRST_WORD_MASK(__start);					\
+	idx = __start / BITS_PER_LONG;						\
+										\
+	for (val = (FETCH) & mask; idx < __bits / BITS_PER_LONG; val = (FETCH))	{\
+		w += hweight_long(val);						\
+		idx++;								\
+	}									\
+										\
+	if (__bits % BITS_PER_LONG)						\
+		w += hweight_long((val) & BITMAP_LAST_WORD_MASK(__bits));	\
+										\
+	w;									\
+})
+
 unsigned int __bitmap_weight(const unsigned long *bitmap, unsigned int bits)
 {
 	return BITMAP_WEIGHT(bitmap[idx], bits);
@@ -358,6 +376,13 @@ unsigned int __bitmap_weight_and(const unsigned long *bitmap1,
 }
 EXPORT_SYMBOL(__bitmap_weight_and);
 
+unsigned int __bitmap_weight_from(const unsigned long *bitmap,
+					unsigned int bits, unsigned int start)
+{
+	return BITMAP_WEIGHT_FROM(bitmap[idx], bits, start);
+}
+EXPORT_SYMBOL(__bitmap_weight_from);
+
 void __bitmap_set(unsigned long *map, unsigned int start, int len)
 {
 	unsigned long *p = map + BIT_WORD(start);
diff --git a/lib/test_bitmap.c b/lib/test_bitmap.c
index 3248ed58a817..a5d823f7589d 100644
--- a/lib/test_bitmap.c
+++ b/lib/test_bitmap.c
@@ -361,6 +361,23 @@ static void __init test_bitmap_region(void)
 	expect_eq_uint(bitmap_weight(bmap, 1000), 0);
 }
 
+static void __init test_weight(void)
+{
+	DECLARE_BITMAP(bmap, 1024);
+	unsigned int idx, w;
+
+	for (idx = 0; idx < 1024; idx++)
+		__assign_bit(idx, bmap, idx);
+
+	w = bitmap_weight(bmap, 1024);
+	for (idx = 0; idx < 1024; idx++) {
+		unsigned int w1 = bitmap_weight(bmap, idx);
+		unsigned int w2 = bitmap_weight_from(bmap, 1024, idx);
+
+		expect_eq_uint(w1 + w2, w);
+	}
+}
+
 #define EXP2_IN_BITS	(sizeof(exp2) * 8)
 
 static void __init test_replace(void)
@@ -1260,6 +1277,7 @@ static void __init selftest(void)
 	test_copy();
 	test_bitmap_region();
 	test_replace();
+	test_weight();
 	test_bitmap_arr32();
 	test_bitmap_arr64();
 	test_bitmap_parse();
-- 
2.39.2