From 7db6e06354328ea7c6164723f504e8ba58d0c4a4 Mon Sep 17 00:00:00 2001
From: LoGin <longjin@DragonOS.org>
Date: Tue, 30 Apr 2024 18:45:01 +0800
Subject: [PATCH] feat(bitmap): Add bit and for AllocBitMap (#793)

---
 kernel/crates/bitmap/src/alloc_bitmap.rs   | 14 ++++++++++++++
 kernel/crates/bitmap/tests/alloc-bitmap.rs | 20 ++++++++++++++++++++
 kernel/src/libs/cpumask.rs                 | 11 +++++++++++
 3 files changed, 45 insertions(+)

diff --git a/kernel/crates/bitmap/src/alloc_bitmap.rs b/kernel/crates/bitmap/src/alloc_bitmap.rs
index d4c6c510..7f06350e 100644
--- a/kernel/crates/bitmap/src/alloc_bitmap.rs
+++ b/kernel/crates/bitmap/src/alloc_bitmap.rs
@@ -1,3 +1,5 @@
+use core::ops::BitAnd;
+
 use alloc::vec::Vec;
 
 use crate::{bitmap_core::BitMapCore, traits::BitMapOps};
@@ -108,3 +110,15 @@ impl BitMapOps<usize> for AllocBitmap {
         self.core.set_all(self.elements, &mut self.data, value);
     }
 }
+
+impl BitAnd for AllocBitmap {
+    type Output = Self;
+
+    fn bitand(self, rhs: Self) -> Self::Output {
+        let mut result = AllocBitmap::new(self.elements);
+        for i in 0..rhs.data.len() {
+            result.data[i] = self.data[i] & rhs.data[i];
+        }
+        result
+    }
+}
diff --git a/kernel/crates/bitmap/tests/alloc-bitmap.rs b/kernel/crates/bitmap/tests/alloc-bitmap.rs
index bd01b964..80d14913 100644
--- a/kernel/crates/bitmap/tests/alloc-bitmap.rs
+++ b/kernel/crates/bitmap/tests/alloc-bitmap.rs
@@ -643,3 +643,23 @@ fn test_alloc_bitmap_full_128() {
     assert_eq!(bitmap.is_full(), false);
     assert_eq!(bitmap.is_empty(), true);
 }
+
+#[test]
+fn test_alloc_bitmap_bitand_128() {
+    let mut bitmap = AllocBitmap::new(128);
+    bitmap.set_all(true);
+
+    let mut bitmap2 = AllocBitmap::new(128);
+
+    bitmap2.set(0, true);
+    bitmap2.set(1, true);
+    bitmap2.set(67, true);
+
+    let bitmap3 = bitmap & bitmap2;
+
+    assert_eq!(bitmap3.len(), 128);
+    assert_eq!(bitmap3.size(), 16);
+    assert_eq!(bitmap3.first_index(), Some(0));
+    assert_eq!(bitmap3.first_false_index(), Some(2));
+    assert_eq!(bitmap3.last_index(), Some(67));
+}
diff --git a/kernel/src/libs/cpumask.rs b/kernel/src/libs/cpumask.rs
index b75ca303..7cbac899 100644
--- a/kernel/src/libs/cpumask.rs
+++ b/kernel/src/libs/cpumask.rs
@@ -1,3 +1,5 @@
+use core::ops::BitAnd;
+
 use bitmap::{traits::BitMapOps, AllocBitmap};
 
 use crate::{mm::percpu::PerCpu, smp::cpu::ProcessorId};
@@ -86,6 +88,15 @@ impl CpuMask {
     }
 }
 
+impl BitAnd for CpuMask {
+    type Output = Self;
+
+    fn bitand(self, rhs: Self) -> Self::Output {
+        let bmp = self.bmp & rhs.bmp;
+        Self { bmp }
+    }
+}
+
 pub struct CpuMaskIter<'a> {
     mask: &'a CpuMask,
     index: Option<ProcessorId>,