diff --git a/README.md b/README.md index c6e91df..6dcb878 100644 --- a/README.md +++ b/README.md @@ -88,6 +88,32 @@ for h in handles { } ``` +### Batch Allocation + +```rust +use sbitmap::Sbitmap; + +// Create a bitmap with 1024 bits +let sb = Sbitmap::new(1024, None, false); +let mut hint = 0; + +// Allocate 4 consecutive bits atomically +if let Some(start_bit) = sb.get_batch(4, &mut hint) { + // Use bits: start_bit, start_bit+1, start_bit+2, start_bit+3 + println!("Allocated bits {}-{}", start_bit, start_bit + 3); + + // Process consecutive resources... + for i in 0..4 { + println!("Using bit {}", start_bit + i); + } + + // Free all 4 bits atomically when done + sb.put_batch(start_bit, 4, &mut hint); +} +``` + +**Note:** Batch operations require `nr_bits <= bits_per_word()`. All consecutive bits are guaranteed to be within the same word (no spanning across word boundaries). + ## API ### `Sbitmap::new(depth: usize, shift: Option, round_robin: bool) -> Self` @@ -108,6 +134,23 @@ Allocate a free bit. The `hint` parameter is a mutable reference to the caller's Free a previously allocated bit. The `hint` parameter is updated to improve cache locality for subsequent allocations. +### `get_batch(&self, nr_bits: usize, hint: &mut usize) -> Option` + +Allocate `nr_bits` consecutive free bits from the bitmap atomically. This operation provides acquire barrier semantics on success. Only supports `nr_bits <= bits_per_word()` to ensure all bits are within the same word (no spanning across word boundaries). + +Returns `Some(start_bit)` where `start_bit` is the first bit of the allocated consecutive range, or `None` if no consecutive `nr_bits` are available or `nr_bits > bits_per_word()`. + +**Use cases:** +- Allocating contiguous resource ranges (e.g., multiple consecutive I/O tags) +- Batch resource allocation for improved efficiency +- DMA buffer allocation requiring consecutive indices + +### `put_batch(&self, bitnr: usize, nr_bits: usize, hint: &mut usize)` + +Free `nr_bits` consecutive previously allocated bits starting from `bitnr`. This operation provides release barrier semantics, ensuring that all writes to data associated with these bits are visible before the bits are freed. Only supports `nr_bits <= bits_per_word()` to ensure all bits are within the same word. + +The `hint` parameter is updated for better cache locality in subsequent allocations. + ### `test_bit(&self, bitnr: usize) -> bool` Check if a bit is currently allocated. @@ -125,12 +168,15 @@ Get the total number of bits in the bitmap. - **Tag allocation**: I/O tag allocation for block devices - **Resource pools**: Any scenario requiring efficient concurrent resource allocation - **Lock-free data structures**: Building block for concurrent algorithms +- **Batch resource allocation**: Allocating multiple consecutive I/O tags, DMA buffers, or contiguous resource ranges - **NUMA machine**: improvement on NUMA machines is obvious ## Performance Characteristics - **Allocation**: O(n) worst case, O(1) average with hints - **Deallocation**: O(1) +- **Batch allocation**: O(n * nr_bits) worst case, finds consecutive bits within single word +- **Batch deallocation**: O(1), atomic clear of consecutive bits - **Memory overhead**: ~56 bytes per word (64 bits) due to cache-line alignment - **Thread safety**: Lock-free with atomic operations - **Scalability**: Linear scaling with number of CPUs up to bitmap depth @@ -165,6 +211,8 @@ let sb = Sbitmap::new(1024, None, false); // Auto-calculated based on depth - `get()`: Acquire semantics - ensures allocated bit is visible before use - `put()`: Release semantics - ensures all writes complete before bit is freed +- `get_batch()`: Acquire semantics - ensures all allocated bits are visible before use +- `put_batch()`: Release semantics - ensures all writes complete before bits are freed ## Comparison with Alternatives @@ -189,14 +237,17 @@ cargo run --bin bench_compare --release -- --depth 1024 --time 5 # Specify bitmap depth, shift, and duration cargo run --bin bench_compare --release -- --depth 512 --shift 5 --time 10 +# Benchmark batch operations (allocating 4 consecutive bits) +cargo run --bin bench_compare --release -- --depth 128 --batch 4 --time 5 + # Show help cargo run --bin bench_compare --release -- --help ``` This benchmark: - Auto-detects available CPUs and spawns N-1 concurrent tasks -- Measures operations per second (get + put pairs) -- Compares sbitmap vs a baseline lockless implementation +- Measures operations per second (get + put pairs for single-bit mode, get_batch + put_batch pairs for batch mode) +- Compares sbitmap vs a baseline lockless implementation (single-bit mode only) - Defaults: 32 bits, auto-calculated shift, 10 seconds, N-1 tasks (where N is total CPU count) Options: @@ -204,6 +255,7 @@ Options: - `--shift SHIFT` - log2(bits per word), auto-calculated if not specified - `--time TIME` - Benchmark duration in seconds (default: 10) - `--tasks TASKS` - Number of concurrent tasks (default: NUM_CPUS - 1) +- `--batch NR_BITS` - Use get_batch/put_batch with NR_BITS (default: 1, single bit mode) - `--round-robin` - Enable round-robin allocation mode (default: disabled) See [benches/README.md](benches/README.md) for more details. diff --git a/benches/compare.rs b/benches/compare.rs index 7610acc..6c34372 100644 --- a/benches/compare.rs +++ b/benches/compare.rs @@ -76,28 +76,32 @@ impl SimpleBitmap { } } +/// Initialize allocation hint combining stack address + system time for better randomization +/// +/// This creates a pseudo-random starting point for each task by combining: +/// - Stack address (different for each thread) +/// - Current time in nanoseconds +/// This helps spread allocations across the bitmap and reduce contention. +fn init_hint(depth: usize) -> usize { + let stack_var = 0u8; + let addr = &stack_var as *const _ as usize; + let time_ns = SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos() as usize; + ((addr / 64).wrapping_add(time_ns)) % depth +} + /// Run benchmark workload: continuous get() and put() operations -fn run_workload( - bitmap: Arc, - duration: Duration, - ops_counter: Arc, - depth: usize, -) where +fn run_workload(bitmap: Arc, duration: Duration, ops_counter: Arc, depth: usize) +where B: Send + Sync + 'static, B: BitmapOps, { thread::spawn(move || { let start = Instant::now(); let mut local_ops = 0u64; - - // Initialize hint combining stack address + system time for better randomization - let stack_var = 0u8; - let addr = &stack_var as *const _ as usize; - let time_ns = SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap() - .as_nanos() as usize; - let mut hint = ((addr / 64).wrapping_add(time_ns)) % depth; + let mut hint = init_hint(depth); while start.elapsed() < duration { // One operation = get() + put() @@ -111,6 +115,31 @@ fn run_workload( }); } +/// Run benchmark workload with batch operations: continuous get_batch() and put_batch() operations +fn run_batch_workload( + bitmap: Arc, + duration: Duration, + ops_counter: Arc, + depth: usize, + batch_size: usize, +) { + thread::spawn(move || { + let start = Instant::now(); + let mut local_ops = 0u64; + let mut hint = init_hint(depth); + + while start.elapsed() < duration { + // One operation = get_batch() + put_batch() + if let Some(bit) = bitmap.get_batch(batch_size, &mut hint) { + bitmap.put_batch(bit, batch_size, &mut hint); + local_ops += 1; + } + } + + ops_counter.fetch_add(local_ops, Ordering::Relaxed); + }); +} + /// Trait for bitmap operations to allow generic benchmarking trait BitmapOps { fn get(&self, hint: &mut usize) -> Option; @@ -145,9 +174,7 @@ fn detect_numa_nodes() -> usize { let count = entries .filter_map(|e| e.ok()) .filter(|e| { - e.file_name() - .to_string_lossy() - .starts_with("node") + e.file_name().to_string_lossy().starts_with("node") && e.file_name() .to_string_lossy() .chars() @@ -155,7 +182,11 @@ fn detect_numa_nodes() -> usize { .all(|c| c.is_ascii_digit()) }) .count(); - if count > 0 { count } else { 1 } + if count > 0 { + count + } else { + 1 + } } Err(_) => 1, // Default to 1 if we can't detect } @@ -176,6 +207,9 @@ fn print_usage(program: &str) { eprintln!(" --shift SHIFT log2(bits per word) (default: auto-calculated)"); eprintln!(" --time TIME Benchmark duration in seconds (default: 10)"); eprintln!(" --tasks TASKS Number of concurrent tasks (default: NUM_CPUS - 1)"); + eprintln!( + " --batch NR_BITS Use get_batch/put_batch with NR_BITS (default: 1, single bit mode)" + ); eprintln!(" --round-robin Enable round-robin allocation mode (default: disabled)"); eprintln!(" -h, --help Show this help message"); eprintln!(); @@ -183,6 +217,7 @@ fn print_usage(program: &str) { eprintln!(" {} --depth 1024 --time 5", program); eprintln!(" {} --depth 512 --shift 5 --time 10", program); eprintln!(" {} --depth 256 --tasks 8 --round-robin", program); + eprintln!(" {} --depth 128 --batch 4", program); } /// Run benchmark with N tasks @@ -190,11 +225,68 @@ fn benchmark(name: &str, bitmap: Arc, duration: Duration, depth: usize, nu where B: Send + Sync + 'static + BitmapOps, { - println!("\n=== {} Benchmark ===", name); + benchmark_internal( + name, + duration, + depth, + num_tasks, + None, + |bitmap_clone, counter| { + run_workload(bitmap_clone, duration, counter, depth); + }, + bitmap, + ); +} + +/// Run batch benchmark with N tasks (Sbitmap only) +fn batch_benchmark( + name: &str, + bitmap: Arc, + duration: Duration, + depth: usize, + num_tasks: usize, + batch_size: usize, +) { + benchmark_internal( + name, + duration, + depth, + num_tasks, + Some(batch_size), + |bitmap_clone, counter| { + run_batch_workload(bitmap_clone, duration, counter, depth, batch_size); + }, + bitmap, + ); +} + +/// Internal benchmark implementation shared by benchmark() and batch_benchmark() +fn benchmark_internal( + name: &str, + duration: Duration, + depth: usize, + num_tasks: usize, + batch_size: Option, + spawn_workload: F, + bitmap: Arc, +) where + B: Send + Sync + 'static, + F: Fn(Arc, Arc), +{ + // Print header + if batch_size.is_some() { + println!("\n=== {} Benchmark (Batch Mode) ===", name); + } else { + println!("\n=== {} Benchmark ===", name); + } + println!("Configuration:"); println!(" - Duration: {:?}", duration); println!(" - Tasks: {}", num_tasks); println!(" - Bitmap depth: {} bits", depth); + if let Some(batch) = batch_size { + println!(" - Batch size: {} bits", batch); + } // Create counter for each task let mut ops_counters = Vec::new(); @@ -206,7 +298,7 @@ where for i in 0..num_tasks { let bitmap_clone = Arc::clone(&bitmap); let counter = Arc::clone(&ops_counters[i]); - run_workload(bitmap_clone, duration, counter, depth); + spawn_workload(bitmap_clone, counter); } // Wait for duration + a bit more for threads to finish @@ -219,25 +311,35 @@ where for i in 0..num_tasks { let ops = ops_counters[i].load(Ordering::Relaxed); let ops_per_sec = ops as f64 / duration_secs; - println!(" Task {}: {} ops, {} ops/sec ({:.4} Mops/sec)", - i, ops, ops_per_sec as u64, ops_per_sec / 1_000_000.0); + println!( + " Task {}: {} ops, {} ops/sec ({:.4} Mops/sec)", + i, + ops, + ops_per_sec as u64, + ops_per_sec / 1_000_000.0 + ); total_ops += ops; } let total_ops_per_sec = total_ops as f64 / duration_secs; - println!(" Total: {} ops, {} ops/sec ({:.4} Mops/sec)", - total_ops, total_ops_per_sec as u64, total_ops_per_sec / 1_000_000.0); + println!( + " Total: {} ops, {} ops/sec ({:.4} Mops/sec)", + total_ops, + total_ops_per_sec as u64, + total_ops_per_sec / 1_000_000.0 + ); } fn main() { // Parse command line arguments: --depth DEPTH --shift SHIFT --time TIME --tasks TASKS --round-robin let args: Vec = env::args().collect(); - let mut depth = 32usize; // Default depth - let mut shift: Option = None; // Default shift (auto-calculate) - let mut time = 10u64; // Default time in seconds - let mut tasks: Option = None; // Default tasks (auto-calculate: NUM_CPUS - 1) - let mut round_robin = false; // Default round-robin mode (disabled) + let mut depth = 32usize; // Default depth + let mut shift: Option = None; // Default shift (auto-calculate) + let mut time = 10u64; // Default time in seconds + let mut tasks: Option = None; // Default tasks (auto-calculate: NUM_CPUS - 1) + let mut batch_size = 1usize; // Default batch size (1 = single bit mode) + let mut round_robin = false; // Default round-robin mode (disabled) // Simple argument parser let mut i = 1; @@ -302,6 +404,24 @@ fn main() { tasks = Some(tasks_val); i += 2; } + "--batch" => { + if i + 1 >= args.len() { + eprintln!("Error: --batch requires a value"); + print_usage(&args[0]); + std::process::exit(1); + } + batch_size = args[i + 1].parse::().unwrap_or_else(|_| { + eprintln!("Error: Invalid batch value '{}'", args[i + 1]); + print_usage(&args[0]); + std::process::exit(1); + }); + if batch_size == 0 { + eprintln!("Error: batch size must be at least 1"); + print_usage(&args[0]); + std::process::exit(1); + } + i += 2; + } "--round-robin" => { round_robin = true; i += 1; @@ -332,7 +452,7 @@ fn main() { // Determine number of tasks to run let num_cpus = match tasks { - Some(t) => t, // Use user-specified value + Some(t) => t, // Use user-specified value None => { // Default: Use N-1 CPUs if N > 1, else use N if total_cpus > 1 { @@ -347,8 +467,10 @@ fn main() { println!("║ Sbitmap vs Simple Lockless Bitmap Benchmark Comparison ║"); println!("╚═══════════════════════════════════════════════════════════╝"); println!(); - println!("System: {} CPUs detected, {} NUMA nodes, using {} tasks for benchmark", - total_cpus, numa_nodes, num_cpus); + println!( + "System: {} CPUs detected, {} NUMA nodes, using {} tasks for benchmark", + total_cpus, numa_nodes, num_cpus + ); println!("Bitmap depth: {} bits", depth); // Create sbitmap to get actual configuration @@ -360,21 +482,45 @@ fn main() { } else { println!("Shift: auto-calculated (bits per word: {})", bits_per_word); } - println!("Round-robin: {}", if round_robin { "enabled" } else { "disabled" }); + println!( + "Round-robin: {}", + if round_robin { "enabled" } else { "disabled" } + ); + println!( + "Batch size: {} bit{}", + batch_size, + if batch_size == 1 { "" } else { "s" } + ); println!("Duration: {} seconds", duration_secs); println!(); - // Benchmark 1: Sbitmap (cache-line optimized with per-task hints) - benchmark("Sbitmap (Optimized)", sbitmap, duration, depth, num_cpus); + if batch_size > 1 { + // Batch mode: only benchmark Sbitmap with get_batch/put_batch + if batch_size > bits_per_word { + eprintln!( + "Error: batch size ({}) exceeds bits_per_word ({})", + batch_size, bits_per_word + ); + eprintln!("Batch operations require nr_bits <= bits_per_word()"); + std::process::exit(1); + } + batch_benchmark("Sbitmap", sbitmap, duration, depth, num_cpus, batch_size); + } else { + // Single bit mode: benchmark both Sbitmap and SimpleBitmap + // Benchmark 1: Sbitmap (cache-line optimized with per-task hints) + benchmark("Sbitmap (Optimized)", sbitmap, duration, depth, num_cpus); - // Benchmark 2: SimpleBitmap (no cache-line optimization, no hints) - let simple = Arc::new(SimpleBitmap::new(depth)); - benchmark("SimpleBitmap (Baseline)", simple, duration, depth, num_cpus); + // Benchmark 2: SimpleBitmap (no cache-line optimization, no hints) + let simple = Arc::new(SimpleBitmap::new(depth)); + benchmark("SimpleBitmap (Baseline)", simple, duration, depth, num_cpus); + } - println!("\n╔═══════════════════════════════════════════════════════════╗"); - println!("║ Summary ║"); - println!("╚═══════════════════════════════════════════════════════════╝"); - println!(" + if batch_size == 1 { + println!("\n╔═══════════════════════════════════════════════════════════╗"); + println!("║ Summary ║"); + println!("╚═══════════════════════════════════════════════════════════╝"); + println!( + " Tasks: {} concurrent tasks Sbitmap optimizations: @@ -391,5 +537,25 @@ Expected: Sbitmap should show higher ops/sec due to: - Reduced false sharing between CPUs - Better cache locality with caller-provided hints - Less contention on bitmap words -", num_cpus); +", + num_cpus + ); + } else { + println!("\n╔═══════════════════════════════════════════════════════════╗"); + println!("║ Batch Mode Summary ║"); + println!("╚═══════════════════════════════════════════════════════════╝"); + println!( + " +Tasks: {} concurrent tasks +Batch size: {} consecutive bits + +Batch operations: + ✓ Atomic allocation of {} consecutive bits via get_batch() + ✓ Atomic deallocation of {} consecutive bits via put_batch() + ✓ All bits guaranteed within single word (no spanning) + ✓ Lock-free with acquire/release memory ordering +", + num_cpus, batch_size, batch_size, batch_size + ); + } } diff --git a/src/lib.rs b/src/lib.rs index 991aa2e..52867f9 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -66,13 +66,12 @@ impl Sbitmap { pub fn new(depth: usize, shift: Option, round_robin: bool) -> Self { let shift = shift.unwrap_or_else(|| Self::calculate_shift(depth)); let bits_per_word = 1usize << shift; - let map_nr = (depth + bits_per_word - 1) / bits_per_word; // DIV_ROUND_UP + let map_nr = depth.div_ceil(bits_per_word); let map = (0..map_nr).map(|_| SbitmapWord::new()).collect(); log::debug!( - "sbitmap::new: depth={}, shift={}, map_nr={}, bits_per_word={}, round_robin={}", - depth, shift, map_nr, bits_per_word, round_robin + "sbitmap::new: depth={depth}, shift={shift}, map_nr={map_nr}, bits_per_word={bits_per_word}, round_robin={round_robin}" ); Self { @@ -149,6 +148,45 @@ impl Sbitmap { } } + /// Create a bit mask for nr_bits + /// + /// Returns a mask with nr_bits set to 1. + /// Handles the special case when nr_bits == BITS_PER_WORD to avoid shift overflow. + #[inline] + fn make_mask(nr_bits: usize) -> usize { + if nr_bits == BITS_PER_WORD { + usize::MAX + } else { + (1usize << nr_bits).wrapping_sub(1) + } + } + + /// Find nr_bits consecutive zero bits in a word starting from hint + /// + /// Returns the starting position if found, None otherwise. + #[inline] + fn find_next_zero_batch( + word: usize, + depth: usize, + hint: usize, + nr_bits: usize, + ) -> Option { + if depth < nr_bits || hint > depth.saturating_sub(nr_bits) { + return None; + } + + let mask = Self::make_mask(nr_bits); + + for start in hint..=(depth - nr_bits) { + let bits_mask = mask << start; + if (word & bits_mask) == 0 { + return Some(start); + } + } + + None + } + /// Atomically test and set a bit (acquire semantics) /// /// Returns true if the bit was successfully allocated (was 0, now 1) @@ -208,6 +246,58 @@ impl Sbitmap { } } + /// Try to allocate nr_bits consecutive bits from a specific word + fn get_batch_from_word( + &self, + word: &AtomicUsize, + depth: usize, + alloc_hint: usize, + nr_bits: usize, + wrap: bool, + ) -> Option { + if depth < nr_bits { + return None; + } + + let mut hint = alloc_hint; + let wrap = wrap && hint > 0; // don't wrap if starting from 0 + + loop { + // Read current word value + let current = word.load(Ordering::Relaxed); + + // Find nr_bits consecutive zero bits starting from hint + let nr = match Self::find_next_zero_batch(current, depth, hint, nr_bits) { + Some(bit) => bit, + None => { + // If we started with an offset and wrapping is allowed, + // try again from the beginning + if hint > 0 && wrap { + hint = 0; + continue; + } + return None; + } + }; + + // Try to atomically set all nr_bits bits + let mask = Self::make_mask(nr_bits); + let bits_mask = mask << nr; + let old = word.fetch_or(bits_mask, Ordering::Acquire); + + // Check if all bits were zero before we set them + if (old & bits_mask) == 0 { + return Some(nr); + } + + // Some bits were already set, continue searching from next position + hint = nr + 1; + if hint > depth.saturating_sub(nr_bits) { + hint = 0; + } + } + } + /// Find and allocate a bit starting from the given index fn find_bit(&self, start_index: usize, alloc_hint: usize, wrap: bool) -> Option { let mut index = start_index; @@ -232,6 +322,38 @@ impl Sbitmap { None } + /// Find and allocate nr_bits consecutive bits starting from the given index + fn find_batch( + &self, + start_index: usize, + alloc_hint: usize, + nr_bits: usize, + wrap: bool, + ) -> Option { + let mut index = start_index; + let mut hint = alloc_hint; + + for _ in 0..self.map_nr { + let depth = self.map_depth(index); + if depth >= nr_bits { + if let Some(bit) = + self.get_batch_from_word(&self.map[index].word, depth, hint, nr_bits, wrap) + { + return Some((index << self.shift) + bit); + } + } + + // Move to next word + hint = 0; + index += 1; + if index >= self.map_nr { + index = 0; + } + } + + None + } + /// Allocate a free bit from the bitmap /// /// This operation provides acquire barrier semantics on success. @@ -309,6 +431,119 @@ impl Sbitmap { } } + /// Allocate nr_bits consecutive free bits from the bitmap + /// + /// This operation provides acquire barrier semantics on success. + /// Only supports nr_bits <= bits_per_word() to ensure all bits are in the same word. + /// + /// # Arguments + /// * `nr_bits` - Number of consecutive bits to allocate + /// * `hint` - Mutable reference to caller's allocation hint for reducing contention + /// + /// # Returns + /// * `Some(start_bit)` - Successfully allocated starting bit number + /// * `None` - No consecutive nr_bits available or nr_bits > bits_per_word() + pub fn get_batch(&self, nr_bits: usize, hint: &mut usize) -> Option { + // Validate nr_bits + if nr_bits == 0 || nr_bits > self.bits_per_word() { + return None; + } + + // Fall back to single bit allocation for nr_bits == 1 + if nr_bits == 1 { + return self.get(hint); + } + + // Validate and sanitize hint + if *hint >= self.depth { + *hint = 0; + } + + let h = *hint; + let index = self.bit_to_index(h); + + // Calculate bit offset within the word + let alloc_hint = if self.round_robin { + self.bit_to_offset(h) + } else { + 0 + }; + + let allocated = self.find_batch(index, alloc_hint, nr_bits, !self.round_robin); + + // Update hint based on allocation result + match allocated { + None => { + // Map is full, reset hint to 0 + *hint = 0; + } + Some(nr) if nr == h || self.round_robin => { + // Only update if we used the hint or in round-robin mode + let next_hint = nr + nr_bits; + *hint = if next_hint >= self.depth { + 0 + } else { + next_hint + }; + } + _ => { + // Don't update hint if we didn't use it + } + } + + allocated + } + + /// Free nr_bits consecutive previously allocated bits + /// + /// This operation provides release barrier semantics, ensuring that + /// all writes to data associated with these bits are visible before + /// the bits are freed. + /// Only supports nr_bits <= bits_per_word() to ensure all bits are in the same word. + /// + /// # Arguments + /// * `bitnr` - The starting bit number to free (must have been returned by get_batch()) + /// * `nr_bits` - Number of consecutive bits to free + /// * `hint` - Mutable reference to caller's allocation hint for better cache locality + pub fn put_batch(&self, bitnr: usize, nr_bits: usize, hint: &mut usize) { + // Validate nr_bits + if nr_bits == 0 || nr_bits > self.bits_per_word() { + return; + } + + // Fall back to single bit deallocation for nr_bits == 1 + if nr_bits == 1 { + self.put(bitnr, hint); + return; + } + + // Validate range + if bitnr >= self.depth || bitnr + nr_bits > self.depth { + return; // Invalid bit range + } + + let start_index = self.bit_to_index(bitnr); + let end_index = self.bit_to_index(bitnr + nr_bits - 1); + + // Ensure all bits are in the same word + if start_index != end_index { + return; + } + + let offset = self.bit_to_offset(bitnr); + let mask = Self::make_mask(nr_bits); + let clear_mask = !(mask << offset); + + self.map[start_index] + .word + .fetch_and(clear_mask, Ordering::Release); + + // Update hint for better cache locality (non-round-robin mode) + if !self.round_robin && bitnr < self.depth { + *hint = bitnr; + } + } + /// Get the total number of bits in the bitmap pub fn depth(&self) -> usize { self.depth @@ -542,7 +777,11 @@ mod tests { let bit = sb.get(&mut hint).expect("Should allocate bit"); allocated.push(bit); // In round-robin mode, bits should be allocated sequentially - assert_eq!(bit, i, "Round-robin should allocate bit {} but got {}", i, bit); + assert_eq!( + bit, i, + "Round-robin should allocate bit {} but got {}", + i, bit + ); } // Free some bits in the middle @@ -684,4 +923,304 @@ mod tests { // All bits should be free now assert_eq!(sb.weight(), 0); } + + #[test] + fn test_batch_basic() { + let sb = Sbitmap::new(64, None, false); + let mut hint = 0; + + // Allocate 4 consecutive bits + let start = sb.get_batch(4, &mut hint).expect("Should allocate 4 bits"); + assert!(start < 64); + + // Verify all 4 bits are set + for i in 0..4 { + assert!(sb.test_bit(start + i), "Bit {} should be set", start + i); + } + assert_eq!(sb.weight(), 4); + + // Free the 4 bits + sb.put_batch(start, 4, &mut hint); + + // Verify all 4 bits are clear + for i in 0..4 { + assert!(!sb.test_bit(start + i), "Bit {} should be clear", start + i); + } + assert_eq!(sb.weight(), 0); + } + + #[test] + fn test_batch_multiple_allocations() { + let sb = Sbitmap::new(128, None, false); + let mut hint = 0; + let mut batches = Vec::new(); + + // Allocate multiple batches of different sizes + batches.push(sb.get_batch(3, &mut hint).expect("Should allocate 3 bits")); + batches.push(sb.get_batch(5, &mut hint).expect("Should allocate 5 bits")); + batches.push(sb.get_batch(2, &mut hint).expect("Should allocate 2 bits")); + + assert_eq!(sb.weight(), 3 + 5 + 2); + + // Free all batches + sb.put_batch(batches[0], 3, &mut hint); + sb.put_batch(batches[1], 5, &mut hint); + sb.put_batch(batches[2], 2, &mut hint); + + assert_eq!(sb.weight(), 0); + } + + #[test] + fn test_batch_exhaustion() { + // Create a small bitmap where we can easily exhaust consecutive bits + let sb = Sbitmap::new(16, Some(4), false); // 16 bits per word + let mut hint = 0; + + // Allocate bits in a pattern that leaves no room for 4 consecutive bits + // Pattern: allocate 3, skip 1, allocate 3, skip 1, etc. + let bit0 = sb.get_batch(3, &mut hint).expect("Should allocate 3 bits"); + assert_eq!(bit0, 0); + + let _bit4 = sb.get(&mut hint).expect("Should skip to bit 3"); + let _bit5 = sb + .get_batch(3, &mut hint) + .expect("Should allocate bits 4-6"); + + let _bit8 = sb.get(&mut hint).expect("Should skip to bit 7"); + let _bit9 = sb + .get_batch(3, &mut hint) + .expect("Should allocate bits 8-10"); + + let _bit12 = sb.get(&mut hint).expect("Should skip to bit 11"); + let _bit13 = sb + .get_batch(3, &mut hint) + .expect("Should allocate bits 12-14"); + + // Now we have: XXX_XXX_XXX_XXX_ (where X is allocated, _ is free) + // Trying to allocate 4 consecutive bits should fail + assert!( + sb.get_batch(4, &mut hint).is_none(), + "Should not find 4 consecutive bits" + ); + + // But we can still allocate single bits + assert!(sb.get(&mut hint).is_some()); + } + + #[test] + fn test_batch_edge_cases() { + let sb = Sbitmap::new(64, None, false); + let mut hint = 0; + + // Test nr_bits = 0 + assert!(sb.get_batch(0, &mut hint).is_none()); + + // Test nr_bits > bits_per_word + let too_large = sb.bits_per_word() + 1; + assert!(sb.get_batch(too_large, &mut hint).is_none()); + + // Test nr_bits = 1 (should work like regular get) + let bit = sb.get_batch(1, &mut hint).expect("Should allocate 1 bit"); + assert!(sb.test_bit(bit)); + sb.put_batch(bit, 1, &mut hint); + assert!(!sb.test_bit(bit)); + + // Test put_batch with invalid parameters + sb.put_batch(100, 4, &mut hint); // Out of range, should be no-op + assert_eq!(sb.weight(), 0); + + sb.put_batch(62, 4, &mut hint); // Would go past depth (64), should be no-op + assert_eq!(sb.weight(), 0); + + sb.put_batch(10, 0, &mut hint); // nr_bits = 0, should be no-op + assert_eq!(sb.weight(), 0); + } + + #[test] + #[allow(unused_assignments)] + #[cfg(target_pointer_width = "64")] + fn test_batch_word_boundary_64bit() { + // Create bitmap with 64-bit words (only possible on 64-bit systems) + let sb = Sbitmap::new(128, Some(6), false); // 2^6 = 64 bits per word + let mut hint = 0; + + // Allocate bits near the end of the first word (bits 60-63) + for i in 60..64 { + hint = i; + sb.get(&mut hint).expect("Should allocate bit"); + } + + // Try to allocate a batch starting at bit 62 (would span word boundary) + // This should fail because bits 62, 63 are in word 0, and bits 64, 65 are in word 1 + hint = 62; + let batch = sb.get_batch(4, &mut hint); + + // The batch should either: + // 1. Not start at 62 (because it can't span words), or + // 2. Be None if no suitable position found + if let Some(start) = batch { + // If we got a batch, verify all bits are in the same word + let start_word = start / 64; + let end_word = (start + 3) / 64; + assert_eq!(start_word, end_word, "Batch should not span word boundary"); + sb.put_batch(start, 4, &mut hint); + } + + // Verify put_batch rejects spanning word boundary + hint = 0; + sb.put_batch(62, 4, &mut hint); // Should be rejected (spans words) + // Bits 60-63 should still be allocated since put_batch should reject this + assert_eq!(sb.weight(), 4); + } + + #[test] + #[allow(unused_assignments)] + #[cfg(target_pointer_width = "32")] + fn test_batch_word_boundary_32bit() { + // Create bitmap with 32-bit words (for 32-bit systems) + let sb = Sbitmap::new(64, Some(5), false); // 2^5 = 32 bits per word + let mut hint = 0; + + // Allocate bits near the end of the first word (bits 28-31) + for i in 28..32 { + hint = i; + sb.get(&mut hint).expect("Should allocate bit"); + } + + // Try to allocate a batch starting at bit 30 (would span word boundary) + hint = 30; + let batch = sb.get_batch(4, &mut hint); + + // The batch should either: + // 1. Not start at 30 (because it can't span words), or + // 2. Be None if no suitable position found + if let Some(start) = batch { + // If we got a batch, verify all bits are in the same word + let start_word = start / 32; + let end_word = (start + 3) / 32; + assert_eq!(start_word, end_word, "Batch should not span word boundary"); + sb.put_batch(start, 4, &mut hint); + } + + // Verify put_batch rejects spanning word boundary + hint = 0; + sb.put_batch(30, 4, &mut hint); // Should be rejected (spans words) + // Bits 28-31 should still be allocated since put_batch should reject this + assert_eq!(sb.weight(), 4); + } + + #[test] + fn test_batch_concurrent() { + let sb = Arc::new(Sbitmap::new(1024, None, false)); + let mut handles = vec![]; + + // Spawn multiple threads to allocate and free batches + for _ in 0..8 { + let sb_clone = Arc::clone(&sb); + let handle = thread::spawn(move || { + let mut local_batches = Vec::new(); + let mut hint = 0; + + // Allocate some batches of varying sizes + for size in [2, 3, 4, 5, 2, 3].iter() { + if let Some(start) = sb_clone.get_batch(*size, &mut hint) { + local_batches.push((start, *size)); + } + } + + // Verify all allocated bits are set + for (start, size) in &local_batches { + for i in 0..*size { + assert!(sb_clone.test_bit(*start + i)); + } + } + + // Free them + for (start, size) in local_batches { + sb_clone.put_batch(start, size, &mut hint); + } + }); + handles.push(handle); + } + + // Wait for all threads + for handle in handles { + handle.join().unwrap(); + } + + // All bits should be free + assert_eq!(sb.weight(), 0); + } + + #[test] + fn test_batch_fragmentation() { + // Test that batch allocation works correctly with fragmented bitmaps + let sb = Sbitmap::new(64, None, false); + let mut hint = 0; + + // Allocate all bits first + let mut all_bits = Vec::new(); + for _ in 0..64 { + if let Some(bit) = sb.get(&mut hint) { + all_bits.push(bit); + } + } + assert_eq!(sb.weight(), 64); + + // Free every other bit to create a fragmented pattern: _X_X_X_X... + for i in (0..64).step_by(2) { + sb.put(all_bits[i], &mut hint); + } + assert_eq!(sb.weight(), 32); + + // Trying to allocate 2 consecutive bits should fail (all free bits are isolated) + hint = 0; + assert!( + sb.get_batch(2, &mut hint).is_none(), + "Should not find 2 consecutive bits in fragmented bitmap" + ); + + // Free an adjacent bit to create a gap of 2 consecutive free bits + sb.put(all_bits[1], &mut hint); + + // Now we should be able to allocate a batch of 2 + hint = 0; + let batch = sb.get_batch(2, &mut hint); + assert!( + batch.is_some(), + "Should find 2 consecutive bits after creating gap" + ); + + // The batch should be bits 0 and 1 + if let Some(start) = batch { + assert_eq!( + start, 0, + "Should allocate the first available consecutive pair" + ); + } + } + + #[test] + fn test_batch_round_robin() { + // Test batch allocation in round-robin mode + let sb = Sbitmap::new(64, None, true); + let mut hint = 0; + + // In round-robin mode, batches should be allocated sequentially + let batch1 = sb.get_batch(3, &mut hint).expect("Should allocate batch 1"); + assert_eq!(batch1, 0, "First batch should start at 0"); + + let batch2 = sb.get_batch(3, &mut hint).expect("Should allocate batch 2"); + assert_eq!(batch2, 3, "Second batch should start at 3"); + + let batch3 = sb.get_batch(4, &mut hint).expect("Should allocate batch 3"); + assert_eq!(batch3, 6, "Third batch should start at 6"); + + // Free and verify + sb.put_batch(batch1, 3, &mut hint); + sb.put_batch(batch2, 3, &mut hint); + sb.put_batch(batch3, 4, &mut hint); + + assert_eq!(sb.weight(), 0); + } }