Add BitMap trait for primitive types

Cargo.toml

@@ -11,3 +11,4 @@ macros = { path = "./macros" }
 
 [dev-dependencies]
 trybuild = "1.0.111"
+paste = "1.0.15"

README.md

@@ -1,12 +1,44 @@
-# `bitmap!`
+# bitmap
+
+## `BitMap` trait
+
+This trait defines the following API:
+
+```rust
+pub trait BitMap<T> {
+    /// Gets the bit at position `index` from `&self`.
+    fn get_bit(&self, index: u8) -> T;
+    /// Sets the bit at position `index` in `&self`.
+    fn set_bit(&mut self, index: u8, value: T);
+    /// Gets the bits at positions `indices.start..indices.end` from `&self`.
+    fn get_bits(&self, indices: Range<u8>) -> T;
+    /// Sets the bits at positions `indices.start..indices.end` in `&self`.
+    fn set_bits(&mut self, indices: Range<u8>, value: T);
+}
+```
+
+By using the crate's `traits` prelude, the `BitMap` trait is implemented for `u8`, `u16`, `u32`, `u64`, and `u128`.
+
+```rust
+use bitmap::traits::*;
+
+fn main() {
+    let mut x: u64 = 0;
+
+    x.set_bit(1, 1);
+    assert_eq!(x, 2);
+    x.set_bit(1, 0);
+    assert_eq!(x, 0);
+}
+```
+
+## `bitmap!` Procedural Macro
 
 Generates a packed bitmap newtype struct with field-level bit access.
 
 The macro expands to a newtype struct around a `u8` to `u128`, depending on the total bit width
 of the definition, with automatically generated getters and setters for each field.
 
-## API
-
 ### Usage Example
 
 ```rust
@@ -23,9 +55,9 @@ bitmap!(
 let mut player = Player(0);
 assert_eq!(std::mem::size_of::<Player>(), 1);
 
-player.set_imposter(1);
-player.set_finished_tasks(5);
-player.set_kills(3);
+player.set_imposter(1)
+    .set_finished_tasks(5)
+    .set_kills(3);
 
 assert_eq!(player.imposter(), 1);
 assert_eq!(player.finished_tasks(), 5);

macros/src/generator.rs

@@ -66,8 +66,9 @@ pub fn expand_bitmap(input: BitmapInput) -> syn::Result<TokenStream2> {
             }
 
             #[inline]
-            pub fn #setter_name(&mut self, val: #this_storage_ty) -> &mut Self {
-                self.0 = ((self.0 & !((#mask) << #index)) | (((val as #storage_ty) & #mask) << #index));
+            pub fn #setter_name(&mut self, value: #this_storage_ty) -> &mut Self {
+                assert!(value <= #mask as #this_storage_ty);
+                self.0 = ((self.0 & !((#mask) << #index)) | (((value as #storage_ty) & #mask) << #index));
                 self
             }
         }

src/lib.rs

@@ -1,4 +1,7 @@
 pub use macros::bitmap;
+pub use traits::*;
+
+pub mod traits;
 
 #[test]
 fn one_bit() {

src/traits.rs

@@ -0,0 +1,106 @@
+pub trait BitMap<T> {
+    /// Gets the bit at position `index` from `&self`.
+    fn get_bit(&self, index: u8) -> T;
+    /// Sets the bit at position `index` in `&self`.
+    fn set_bit(&mut self, index: u8, value: T);
+    /// Gets the bits at positions `indices.start..indices.end` from `&self`.
+    fn get_bits(&self, indices: ::core::ops::Range<u8>) -> T;
+    /// Sets the bits at positions `indices.start..indices.end` in `&self`.
+    fn set_bits(&mut self, indices: ::core::ops::Range<u8>, value: T);
+}
+
+macro_rules! impl_bitmap {
+    ($ty:ident) => {
+        impl BitMap<$ty> for $ty {
+            fn get_bit(&self, index: u8) -> $ty {
+                *self >> index & 0b1
+            }
+
+            fn set_bit(&mut self, index: u8, value: $ty) {
+                *self = (*self & !(1 << index)) | ((value & 1) << index);
+            }
+
+            fn set_bits(&mut self, indices: ::core::ops::Range<u8>, value: $ty) {
+                let width = indices.end - indices.start;
+                let bit_count = ::core::mem::size_of::<$ty>() * 8;
+
+                let mask = if width as usize >= bit_count { $ty::MAX } else { (1 << width) - 1 };
+
+                *self = (*self & !(mask << indices.start)) | ((value & mask) << indices.start);
+            }
+
+            fn get_bits(&self, indices: ::core::ops::Range<u8>) -> $ty {
+                let width = indices.end - indices.start;
+                let bit_count = ::core::mem::size_of::<$ty>() * 8;
+
+                let mask = if width as usize >= bit_count { $ty::MAX } else { (1 << width) - 1 };
+
+                (*self >> indices.start) & mask
+            }
+        }
+    };
+}
+
+impl_bitmap!(u8);
+impl_bitmap!(u16);
+impl_bitmap!(u32);
+impl_bitmap!(u64);
+impl_bitmap!(u128);
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    macro_rules! generate_tests {
+        ($ty:ident) => {
+            paste::paste! {
+                #[test]
+                fn [<test_ $ty _single_bit>]() {
+                    let mut x: $ty = 0;
+                    let bit_width = ::core::mem::size_of::<$ty>() * 8;
+
+                    for bit in 0..bit_width {
+                        x.set_bit(bit as u8, 1);
+                        assert_eq!(x.get_bit(bit as u8), 1);
+                        assert_eq!(x, 1 << bit);
+
+                        x.set_bit(bit as u8, 0);
+                        assert_eq!(x.get_bit(bit as u8), 0);
+                        assert_eq!(x, 0);
+                    }
+                }
+
+                #[test]
+                fn [<test_ $ty _bit_range>]() {
+                    let mut x: $ty = 0;
+                    let bit_width = ::core::mem::size_of::<$ty>() * 8;
+
+                    for start in (0..bit_width).step_by(8) {
+                        for width in 1..=8.min(bit_width - start) {
+                            let end = start + width;
+                            if end > bit_width {
+                                break;
+                            }
+                            let max_val = if width >= bit_width {
+                                $ty::MAX
+                            } else {
+                                (1 << width) - 1
+                            };
+
+                            x.set_bits(start as u8..end as u8, max_val);
+                            assert_eq!(x.get_bits(start as u8..end as u8), max_val, "Failed range test: {start}..{end} with value {max_val}");
+
+                            x = 0;
+                        }
+                    }
+                }
+            }
+        };
+    }
+
+    generate_tests!(u8);
+    generate_tests!(u16);
+    generate_tests!(u32);
+    generate_tests!(u64);
+    generate_tests!(u128);
+}

tests/trybuild_tests.rs

@@ -15,3 +15,9 @@ fn invalid_type_zero() {
     let t = trybuild::TestCases::new();
     t.compile_fail("tests/ui/invalid_type_zero.rs");
 }
+
+#[test]
+fn traits() {
+    let t = trybuild::TestCases::new();
+    t.pass("tests/ui/traits.rs");
+}

tests/ui/traits.rs

@@ -0,0 +1,7 @@
+use bitmap::traits::*;
+
+fn main() {
+    let mut x: u128 = 0;
+    x.set_bits(0..2, 0b11);
+    assert_eq!(x, 0b11);
+}