use std::collections::HashSet;
use bitvec::prelude::*;
pub struct PieceMap {
piece_size: u64,
total_size: u64,
piece_count: usize,
completed: BitVec<u8, Lsb0>,
}
impl PieceMap {
pub fn new(total_size: u64, piece_size: u64) -> Self {
let piece_count = total_size.div_ceil(piece_size) as usize;
Self {
piece_size,
total_size,
piece_count,
completed: bitvec![u8, Lsb0; 0; piece_count],
}
}
pub fn from_bitset(
total_size: u64,
piece_size: u64,
bitset: &[u8],
piece_count: usize,
) -> Self {
let mut completed = BitVec::<u8, Lsb0>::from_slice(bitset);
completed.resize(piece_count, false);
Self {
piece_size,
total_size,
piece_count,
completed,
}
}
pub fn piece_size(&self) -> u64 {
self.piece_size
}
#[allow(dead_code)]
pub fn total_size(&self) -> u64 {
self.total_size
}
pub fn piece_count(&self) -> usize {
self.piece_count
}
pub fn piece_range(&self, piece_id: usize) -> (u64, u64) {
let start = piece_id as u64 * self.piece_size;
let end = ((piece_id as u64 + 1) * self.piece_size).min(self.total_size);
(start, end)
}
#[allow(dead_code)]
pub fn is_complete(&self, piece_id: usize) -> bool {
piece_id < self.piece_count && self.completed[piece_id]
}
pub fn mark_complete(&mut self, piece_id: usize) {
if piece_id < self.piece_count {
self.completed.set(piece_id, true);
}
}
pub fn next_missing_excluding(&self, exclude: &HashSet<usize>) -> Option<usize> {
(0..self.piece_count).find(|&i| !self.completed[i] && !exclude.contains(&i))
}
pub fn all_done(&self) -> bool {
self.completed.all()
}
pub fn completed_count(&self) -> usize {
self.completed.count_ones()
}
pub fn completed_bytes(&self) -> u64 {
let mut bytes = 0u64;
for i in 0..self.piece_count {
if self.completed[i] {
let (start, end) = self.piece_range(i);
bytes += end - start;
}
}
bytes
}
pub fn remaining_count(&self) -> usize {
self.piece_count - self.completed_count()
}
pub fn to_bitset_bytes(&self) -> Vec<u8> {
self.completed.as_raw_slice().to_vec()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_piece_map_basic() {
let mut pm = PieceMap::new(10_000_000, 1_000_000); assert_eq!(pm.piece_count(), 10);
assert!(!pm.all_done());
assert_eq!(pm.completed_bytes(), 0);
pm.mark_complete(0);
pm.mark_complete(5);
assert!(pm.is_complete(0));
assert!(!pm.is_complete(1));
assert!(pm.is_complete(5));
assert_eq!(pm.completed_count(), 2);
assert_eq!(pm.completed_bytes(), 2_000_000);
}
#[test]
fn test_piece_range_last_piece() {
let pm = PieceMap::new(2_500_000, 1_000_000); assert_eq!(pm.piece_count(), 3);
assert_eq!(pm.piece_range(0), (0, 1_000_000));
assert_eq!(pm.piece_range(1), (1_000_000, 2_000_000));
assert_eq!(pm.piece_range(2), (2_000_000, 2_500_000)); }
#[test]
fn test_next_missing_excluding() {
let mut pm = PieceMap::new(5_000_000, 1_000_000);
pm.mark_complete(0);
pm.mark_complete(2);
let mut excl = HashSet::new();
assert_eq!(pm.next_missing_excluding(&excl), Some(1));
excl.insert(1);
assert_eq!(pm.next_missing_excluding(&excl), Some(3));
}
#[test]
fn test_bitset_roundtrip() {
let mut pm = PieceMap::new(5_000_000, 1_000_000);
pm.mark_complete(0);
pm.mark_complete(3);
let bytes = pm.to_bitset_bytes();
let pm2 = PieceMap::from_bitset(5_000_000, 1_000_000, &bytes, 5);
assert!(pm2.is_complete(0));
assert!(!pm2.is_complete(1));
assert!(!pm2.is_complete(2));
assert!(pm2.is_complete(3));
assert!(!pm2.is_complete(4));
}
#[test]
fn test_all_done() {
let mut pm = PieceMap::new(2_000_000, 1_000_000);
assert!(!pm.all_done());
pm.mark_complete(0);
assert!(!pm.all_done());
pm.mark_complete(1);
assert!(pm.all_done());
}
#[test]
fn test_total_size_and_accessors() {
let pm = PieceMap::new(2_500_000, 1_000_000);
assert_eq!(pm.total_size(), 2_500_000);
assert_eq!(pm.piece_size(), 1_000_000);
assert_eq!(pm.piece_count(), 3);
assert_eq!(pm.remaining_count(), 3);
assert_eq!(pm.completed_count(), 0);
}
#[test]
fn test_mark_complete_out_of_bounds() {
let mut pm = PieceMap::new(1_000_000, 1_000_000);
pm.mark_complete(999);
assert!(!pm.is_complete(999));
}
#[test]
fn test_is_complete_out_of_bounds() {
let pm = PieceMap::new(1_000_000, 1_000_000);
assert!(!pm.is_complete(999));
}
#[test]
fn test_completed_bytes_partial() {
let mut pm = PieceMap::new(2_500_000, 1_000_000);
pm.mark_complete(2); assert_eq!(pm.completed_bytes(), 500_000);
}
}