[PATCH 1/9] bitmap: add bitmap_and_and() and bitmap_and_andnot()

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

Add bitmap_and_and() and bitmap_and_andnot() to combine three bitmaps
in a single pass. Both helpers return whether the resulting bitmap is
non-empty.

Introduce BITMAP_OP() to share the word iteration with bitmap_and(),
and add tests for small constants, multiword bitmaps, aliases, tail
masking, empty results and zero-sized bitmaps.

Signed-off-by: Yury Norov <ynorov@nvidia.com>
---
 include/linux/bitmap.h | 32 +++++++++++++++++
 lib/bitmap.c           | 46 +++++++++++++++++++------
 lib/test_bitmap.c      | 78 ++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 146 insertions(+), 10 deletions(-)

diff --git a/include/linux/bitmap.h b/include/linux/bitmap.h
index 7df1573a409c..d524ef6b0300 100644
--- a/include/linux/bitmap.h
+++ b/include/linux/bitmap.h
@@ -44,6 +44,10 @@ struct device;
  *  bitmap_fill(dst, nbits)                     *dst = ~0UL
  *  bitmap_copy(dst, src, nbits)                *dst = *src
  *  bitmap_and(dst, src1, src2, nbits)          *dst = *src1 & *src2
+ *  bitmap_and_and(dst, src1, src2, src3, nbits)
+ *                                               *dst = *src1 & *src2 & *src3
+ *  bitmap_and_andnot(dst, src1, src2, src3, nbits)
+ *                                               *dst = *src1 & *src2 & ~(*src3)
  *  bitmap_or(dst, src1, src2, nbits)           *dst = *src1 | *src2
  *  bitmap_weighted_or(dst, src1, src2, nbits)	*dst = *src1 | *src2. Returns Hamming Weight of dst
  *  bitmap_weighted_xor(dst, src1, src2, nbits)	*dst = *src1 ^ *src2. Returns Hamming Weight of dst
@@ -166,6 +170,12 @@ void bitmap_cut(unsigned long *dst, const unsigned long *src,
 		unsigned int first, unsigned int cut, unsigned int nbits);
 bool __bitmap_and(unsigned long *dst, const unsigned long *bitmap1,
 		 const unsigned long *bitmap2, unsigned int nbits);
+bool __bitmap_and_and(unsigned long *dst, const unsigned long *bitmap1,
+		      const unsigned long *bitmap2,
+		      const unsigned long *bitmap3, unsigned int nbits);
+bool __bitmap_and_andnot(unsigned long *dst, const unsigned long *bitmap1,
+			 const unsigned long *bitmap2,
+			 const unsigned long *bitmap3, unsigned int nbits);
 void __bitmap_or(unsigned long *dst, const unsigned long *bitmap1,
 		 const unsigned long *bitmap2, unsigned int nbits);
 unsigned int __bitmap_weighted_or(unsigned long *dst, const unsigned long *bitmap1,
@@ -337,6 +347,28 @@ bool bitmap_and(unsigned long *dst, const unsigned long *src1,
 	return __bitmap_and(dst, src1, src2, nbits);
 }
 
+static __always_inline
+bool bitmap_and_and(unsigned long *dst, const unsigned long *src1,
+		    const unsigned long *src2, const unsigned long *src3,
+		    unsigned int nbits)
+{
+	if (small_const_nbits(nbits))
+		return (*dst = *src1 & *src2 & *src3 &
+			       BITMAP_LAST_WORD_MASK(nbits)) != 0;
+	return __bitmap_and_and(dst, src1, src2, src3, nbits);
+}
+
+static __always_inline
+bool bitmap_and_andnot(unsigned long *dst, const unsigned long *src1,
+		       const unsigned long *src2, const unsigned long *src3,
+		       unsigned int nbits)
+{
+	if (small_const_nbits(nbits))
+		return (*dst = *src1 & *src2 & ~(*src3) &
+			       BITMAP_LAST_WORD_MASK(nbits)) != 0;
+	return __bitmap_and_andnot(dst, src1, src2, src3, nbits);
+}
+
 static __always_inline
 void bitmap_or(unsigned long *dst, const unsigned long *src1,
 	       const unsigned long *src2, unsigned int nbits)
diff --git a/lib/bitmap.c b/lib/bitmap.c
index ed685127a107..85ce3cbaa9ab 100644
--- a/lib/bitmap.c
+++ b/lib/bitmap.c
@@ -34,6 +34,25 @@
  * for the best explanations of this ordering.
  */
 
+/*
+ * Common helper for bitmap operations.
+ * @FETCH: The expression that fetches and combines each word of the bitmaps
+ * @bits: The bitmap size in bits
+ */
+#define BITMAP_OP(FETCH, bits)							\
+({										\
+	unsigned long idx, val, sz = (bits), result = 0;			\
+										\
+	for (idx = 0; idx * BITS_PER_LONG < sz; idx++) {			\
+		val = (FETCH);							\
+		if (sz - idx * BITS_PER_LONG < BITS_PER_LONG)		\
+			val &= BITMAP_LAST_WORD_MASK(sz);			\
+		result |= (dst[idx] = val);					\
+	}									\
+										\
+	result != 0;								\
+})
+
 bool __bitmap_equal(const unsigned long *bitmap1,
 		    const unsigned long *bitmap2, unsigned int bits)
 {
@@ -230,19 +249,26 @@ EXPORT_SYMBOL(bitmap_cut);
 bool __bitmap_and(unsigned long *dst, const unsigned long *bitmap1,
 				const unsigned long *bitmap2, unsigned int bits)
 {
-	unsigned int k;
-	unsigned int lim = bits/BITS_PER_LONG;
-	unsigned long result = 0;
-
-	for (k = 0; k < lim; k++)
-		result |= (dst[k] = bitmap1[k] & bitmap2[k]);
-	if (bits % BITS_PER_LONG)
-		result |= (dst[k] = bitmap1[k] & bitmap2[k] &
-			   BITMAP_LAST_WORD_MASK(bits));
-	return result != 0;
+	return BITMAP_OP(bitmap1[idx] & bitmap2[idx], bits);
 }
 EXPORT_SYMBOL(__bitmap_and);
 
+bool __bitmap_and_and(unsigned long *dst, const unsigned long *bitmap1,
+		      const unsigned long *bitmap2,
+		      const unsigned long *bitmap3, unsigned int bits)
+{
+	return BITMAP_OP(bitmap1[idx] & bitmap2[idx] & bitmap3[idx], bits);
+}
+EXPORT_SYMBOL(__bitmap_and_and);
+
+bool __bitmap_and_andnot(unsigned long *dst, const unsigned long *bitmap1,
+			 const unsigned long *bitmap2,
+			 const unsigned long *bitmap3, unsigned int bits)
+{
+	return BITMAP_OP(bitmap1[idx] & bitmap2[idx] & ~bitmap3[idx], bits);
+}
+EXPORT_SYMBOL(__bitmap_and_andnot);
+
 void __bitmap_or(unsigned long *dst, const unsigned long *bitmap1,
 				const unsigned long *bitmap2, unsigned int bits)
 {
diff --git a/lib/test_bitmap.c b/lib/test_bitmap.c
index 56bd23059b26..fbc46ed960a4 100644
--- a/lib/test_bitmap.c
+++ b/lib/test_bitmap.c
@@ -193,6 +193,81 @@ static void __init test_zero_clear(void)
 	expect_eq_pbl("", bmap, 1024);
 }
 
+static void __init test_bitmap_and(void)
+{
+	enum { nbits = BITS_PER_LONG + 13 };
+	DECLARE_BITMAP(src1, nbits);
+	DECLARE_BITMAP(src2, nbits);
+	DECLARE_BITMAP(src3, nbits);
+	DECLARE_BITMAP(dst, nbits);
+	DECLARE_BITMAP(expected, nbits);
+	unsigned long small_src1 = ~0UL;
+	unsigned long small_src2 = ~0UL;
+	unsigned long small_src3 = BIT(2);
+	unsigned long small_dst;
+	bool ret;
+
+	ret = bitmap_and_and(&small_dst, &small_src1, &small_src2,
+			     &small_src3, 4);
+	expect_eq_ulong(true, ret);
+	expect_eq_ulong(BIT(2), small_dst);
+
+	ret = bitmap_and_andnot(&small_dst, &small_src1, &small_src2,
+				&small_src3, 4);
+	expect_eq_ulong(true, ret);
+	expect_eq_ulong(GENMASK(3, 0) & ~BIT(2), small_dst);
+
+	bitmap_zero(src1, nbits);
+	bitmap_zero(src2, nbits);
+	bitmap_zero(src3, nbits);
+	bitmap_zero(expected, nbits);
+	__set_bit(1, src1);
+	__set_bit(2, src1);
+	__set_bit(BITS_PER_LONG + 1, src1);
+	__set_bit(BITS_PER_LONG + 12, src1);
+	__set_bit(BITS_PER_LONG + 13, src1);
+	__set_bit(1, src2);
+	__set_bit(BITS_PER_LONG + 1, src2);
+	__set_bit(BITS_PER_LONG + 12, src2);
+	__set_bit(BITS_PER_LONG + 13, src2);
+	__set_bit(BITS_PER_LONG + 1, src3);
+	__set_bit(1, expected);
+	__set_bit(BITS_PER_LONG + 1, expected);
+	__set_bit(BITS_PER_LONG + 12, expected);
+
+	ret = bitmap_and(dst, src1, src2, nbits);
+	expect_eq_ulong(true, ret);
+	expect_eq_bitmap(expected, dst, nbits);
+	expect_eq_ulong(BIT(1) | BIT(12), dst[1]);
+
+	bitmap_zero(expected, nbits);
+	__set_bit(BITS_PER_LONG + 1, expected);
+	ret = bitmap_and_and(dst, src1, src2, src3, nbits);
+	expect_eq_ulong(true, ret);
+	expect_eq_bitmap(expected, dst, nbits);
+	expect_eq_ulong(BIT(1), dst[1]);
+
+	bitmap_zero(expected, nbits);
+	__set_bit(1, expected);
+	__set_bit(2, expected);
+	__set_bit(BITS_PER_LONG + 12, expected);
+	ret = bitmap_andnot(dst, src1, src3, nbits);
+	expect_eq_ulong(true, ret);
+	expect_eq_bitmap(expected, dst, nbits);
+	expect_eq_ulong(BIT(12), dst[1]);
+
+	__clear_bit(2, expected);
+	ret = bitmap_and_andnot(src2, src1, src2, src3, nbits);
+	expect_eq_ulong(true, ret);
+	expect_eq_bitmap(expected, src2, nbits);
+	expect_eq_ulong(BIT(12), src2[1]);
+
+	bitmap_fill(src3, nbits);
+	ret = bitmap_and_andnot(src2, src1, src2, src3, nbits);
+	expect_eq_ulong(false, ret);
+	expect_eq_pbl("", src2, nbits);
+}
+
 static void __init test_find_nth_bit(void)
 {
 	unsigned long b, bit, cnt = 0;
@@ -1533,6 +1608,8 @@ static void __init test_zero_nbits(void)
 	bitmap_zero(NULL, 0);
 
 	ret = bitmap_and(NULL, NULL, NULL, 0);
+	ret = bitmap_and_and(NULL, NULL, NULL, NULL, 0);
+	ret = bitmap_and_andnot(NULL, NULL, NULL, NULL, 0);
 	ret = bitmap_empty(NULL, 0);
 	ret = bitmap_equal(NULL, NULL, 0);
 	ret = bitmap_full(NULL, 0);
@@ -1566,6 +1643,7 @@ static void __init test_zero_nbits(void)
 static void __init selftest(void)
 {
 	test_zero_clear();
+	test_bitmap_and();
 	test_fill_set();
 	test_copy();
 	test_bitmap_region();
-- 
2.53.0