disk_core/unix_common/

bitmap.rs

//! Bitmap utilities for Unix-like filesystems.
//!
//! Provides efficient bit-level operations on block and inode bitmaps.
//! Uses little-endian bit order within each byte (bit 0 = LSB), matching
//! the standard Linux bitmap convention used by ext2/3/4, xfs AG bitmaps,
//! UFS cylinder group bitmaps, and Minix zone maps.

/// A read-only view over a bitmap stored as a byte slice.
///
/// `bit_count` specifies the number of valid bits, which may be less than
/// `data.len() * 8` (e.g. when the last byte is only partially used).
pub struct BitmapReader<'a> {
    data: &'a [u8],
    bit_count: u64,
}

impl<'a> BitmapReader<'a> {
    /// Create a new bitmap reader.
    ///
    /// `bit_count` is the number of valid bits. Bits beyond `bit_count`
    /// (in the last partial byte) are ignored by counting and iteration methods.
    pub fn new(data: &'a [u8], bit_count: u64) -> Self {
        Self { data, bit_count }
    }

    /// Returns the total number of valid bits in this bitmap.
    pub fn bit_count(&self) -> u64 {
        self.bit_count
    }

    /// Test whether the bit at `index` is set (1).
    ///
    /// Returns `false` if `index` is out of range.
    #[inline]
    pub fn is_bit_set(&self, index: u64) -> bool {
        if index >= self.bit_count {
            return false;
        }
        let byte_idx = (index / 8) as usize;
        let bit_idx = (index % 8) as u32;
        if byte_idx >= self.data.len() {
            return false;
        }
        self.data[byte_idx] & (1u8 << bit_idx) != 0
    }

    /// Count the number of set (1) bits in the valid range.
    pub fn count_set_bits(&self) -> u64 {
        if self.bit_count == 0 {
            return 0;
        }

        let full_bytes = (self.bit_count / 8) as usize;
        let remaining_bits = (self.bit_count % 8) as u32;

        let mut count: u64 = 0;

        // Count full bytes using popcount
        for &byte in &self.data[..full_bytes] {
            count += byte.count_ones() as u64;
        }

        // Handle the partial last byte
        if remaining_bits > 0 && full_bytes < self.data.len() {
            let mask = (1u8 << remaining_bits) - 1;
            count += (self.data[full_bytes] & mask).count_ones() as u64;
        }

        count
    }

    /// Count the number of clear (0) bits in the valid range.
    pub fn count_clear_bits(&self) -> u64 {
        self.bit_count - self.count_set_bits()
    }

    /// Index (0..=7) of the most significant set bit of a **non-zero** byte.
    ///
    /// Widened to `u32` before `leading_zeros` deliberately, and it must stay
    /// that way: mrustc's C backend lowers `leading_zeros` on integers narrower
    /// than 32 bits as a *32-bit* count, so under it `0x01u8.leading_zeros()`
    /// is 31 rather than 7 (and `0x80u8` gives 24 rather than 0; `u16` is wrong
    /// the same way). The obvious `7 - byte.leading_zeros()` therefore wraps to
    /// about 4.29e9 on any mrustc-built binary - which is how the PowerPC build
    /// came to report every ext volume's last used block as the end of the
    /// volume, silently imaging 32 MiB where the desktop imaged 4.5 MiB. It was
    /// invisible because `count_ones` *is* lowered correctly, so bit *counts*
    /// agreed everywhere and only positions were wrong.
    ///
    /// `u32`/`u64` `leading_zeros` are correct under mrustc, so this form is
    /// identical on both compilers. Verified directly against mrustc, not
    /// assumed.
    #[inline]
    fn top_set_bit(byte: u8) -> u32 {
        debug_assert!(byte != 0, "top_set_bit requires a non-zero byte");
        31 - (byte as u32).leading_zeros()
    }

    /// Find the index of the highest set bit, scanning from the end.
    ///
    /// Returns `None` if no bits are set.
    pub fn highest_set_bit(&self) -> Option<u64> {
        if self.bit_count == 0 {
            return None;
        }

        let full_bytes = (self.bit_count / 8) as usize;
        let remaining_bits = (self.bit_count % 8) as u32;

        // Check partial last byte first
        if remaining_bits > 0 && full_bytes < self.data.len() {
            let mask = (1u8 << remaining_bits) - 1;
            let masked = self.data[full_bytes] & mask;
            if masked != 0 {
                let top_bit = Self::top_set_bit(masked);
                return Some(full_bytes as u64 * 8 + top_bit as u64);
            }
        }

        // Scan full bytes from end to start
        for i in (0..full_bytes).rev() {
            if self.data[i] != 0 {
                let top_bit = Self::top_set_bit(self.data[i]);
                return Some(i as u64 * 8 + top_bit as u64);
            }
        }

        None
    }

    /// Find the index of the highest **clear** (0) bit, scanning from the
    /// end. Useful for "set = free" bitmaps (UFS / JFS / AFFS) where the
    /// last-allocated position translates to the highest clear bit.
    ///
    /// Returns `None` if all bits in the valid range are set.
    pub fn highest_clear_bit(&self) -> Option<u64> {
        if self.bit_count == 0 {
            return None;
        }

        let full_bytes = (self.bit_count / 8) as usize;
        let remaining_bits = (self.bit_count % 8) as u32;

        // Partial last byte first — invert before scanning so we can reuse
        // the same `leading_zeros` trick.
        if remaining_bits > 0 && full_bytes < self.data.len() {
            let mask = (1u8 << remaining_bits) - 1;
            let inv = (!self.data[full_bytes]) & mask;
            if inv != 0 {
                let top_bit = Self::top_set_bit(inv);
                // `mask` guarantees `top_bit < remaining_bits`, so the
                // result is already inside the valid range.
                return Some(full_bytes as u64 * 8 + top_bit as u64);
            }
        }

        for i in (0..full_bytes).rev() {
            let inv = !self.data[i];
            if inv != 0 {
                let top_bit = Self::top_set_bit(inv);
                return Some(i as u64 * 8 + top_bit as u64);
            }
        }

        None
    }

    /// Find the index of the lowest clear (0) bit, scanning from the start.
    ///
    /// Returns `None` if all bits are set.
    pub fn lowest_clear_bit(&self) -> Option<u64> {
        if self.bit_count == 0 {
            return None;
        }

        let full_bytes = (self.bit_count / 8) as usize;
        let remaining_bits = (self.bit_count % 8) as u32;

        // Scan full bytes
        for i in 0..full_bytes {
            if self.data[i] != 0xFF {
                let bit = self.data[i].trailing_ones();
                return Some(i as u64 * 8 + bit as u64);
            }
        }

        // Check partial last byte
        if remaining_bits > 0 && full_bytes < self.data.len() {
            let mask = (1u8 << remaining_bits) - 1;
            let masked = self.data[full_bytes] & mask;
            if masked != mask {
                // There's a clear bit in the valid range
                let bit = masked.trailing_ones();
                if bit < remaining_bits {
                    return Some(full_bytes as u64 * 8 + bit as u64);
                }
            }
        }

        None
    }

    /// Iterate over the indices of all set (1) bits in ascending order.
    pub fn iter_set_bits(&self) -> SetBitsIter<'_> {
        SetBitsIter {
            bitmap: self,
            index: 0,
        }
    }

    /// Iterate over the indices of all clear (0) bits in ascending order.
    pub fn iter_clear_bits(&self) -> ClearBitsIter<'_> {
        ClearBitsIter {
            bitmap: self,
            index: 0,
        }
    }
}

/// Iterator over set bit indices in a `BitmapReader`.
pub struct SetBitsIter<'a> {
    bitmap: &'a BitmapReader<'a>,
    index: u64,
}

impl Iterator for SetBitsIter<'_> {
    type Item = u64;

    fn next(&mut self) -> Option<u64> {
        while self.index < self.bitmap.bit_count {
            let i = self.index;
            self.index += 1;
            if self.bitmap.is_bit_set(i) {
                return Some(i);
            }
        }
        None
    }
}

/// Iterator over clear bit indices in a `BitmapReader`.
pub struct ClearBitsIter<'a> {
    bitmap: &'a BitmapReader<'a>,
    index: u64,
}

impl Iterator for ClearBitsIter<'_> {
    type Item = u64;

    fn next(&mut self) -> Option<u64> {
        while self.index < self.bitmap.bit_count {
            let i = self.index;
            self.index += 1;
            if !self.bitmap.is_bit_set(i) {
                return Some(i);
            }
        }
        None
    }
}

// --- Bitmap mutation helpers (for filesystem editing) ---

/// Set the bit at `index` in a mutable bitmap (little-endian bit order).
///
/// Panics if `index` is out of range for the data slice.
#[inline]
pub fn bitmap_set_bit(data: &mut [u8], index: u64) {
    let byte_idx = (index / 8) as usize;
    let bit_idx = (index % 8) as u32;
    data[byte_idx] |= 1u8 << bit_idx;
}

/// Clear the bit at `index` in a mutable bitmap (little-endian bit order).
///
/// Panics if `index` is out of range for the data slice.
#[inline]
pub fn bitmap_clear_bit(data: &mut [u8], index: u64) {
    let byte_idx = (index / 8) as usize;
    let bit_idx = (index % 8) as u32;
    data[byte_idx] &= !(1u8 << bit_idx);
}

/// Find the first clear (0) bit in the bitmap within `bit_count` valid bits.
///
/// Returns `None` if all valid bits are set.
pub fn bitmap_find_clear_bit(data: &[u8], bit_count: u64) -> Option<u64> {
    let full_bytes = (bit_count / 8) as usize;
    let remaining_bits = (bit_count % 8) as u32;

    for (i, &byte) in data[..full_bytes].iter().enumerate() {
        if byte != 0xFF {
            let bit = byte.trailing_ones();
            return Some(i as u64 * 8 + bit as u64);
        }
    }

    if remaining_bits > 0 && full_bytes < data.len() {
        let mask = (1u8 << remaining_bits) - 1;
        let masked = data[full_bytes] & mask;
        if masked != mask {
            let bit = masked.trailing_ones();
            if bit < remaining_bits {
                return Some(full_bytes as u64 * 8 + bit as u64);
            }
        }
    }

    None
}

#[cfg(test)]
mod tests {
    use super::*;

    /// `top_set_bit` must agree with a from-scratch bit scan for every byte.
    ///
    /// This pins the widening in `top_set_bit`. The natural-looking
    /// `7 - byte.leading_zeros()` is correct on rustc and *wrong* on mrustc,
    /// which lowers `leading_zeros` for sub-32-bit integers as a 32-bit count -
    /// so `0x01u8` yields 31, the subtraction wraps to ~4.29e9, and every
    /// caller silently believes the last used block is off the end of the
    /// volume. A native `cargo test` cannot see that difference, so this test
    /// guards the *shape* of the computation: it fails only if someone changes
    /// the arithmetic itself, and the doc comment on `top_set_bit` explains why
    /// the obvious simplification must not be applied.
    #[test]
    fn top_set_bit_matches_a_plain_scan_for_every_byte() {
        for byte in 1u8..=255 {
            let expected = (0..8u32).filter(|b| byte & (1 << b) != 0).max().unwrap();
            assert_eq!(
                BitmapReader::top_set_bit(byte),
                expected,
                "top_set_bit({byte:#04x})"
            );
        }
    }

    /// The whole-bitmap entry points, over every single-bit pattern: a wrong
    /// top-bit calculation would show up as an off-by-N index here.
    #[test]
    fn highest_set_and_clear_bits_are_positionally_correct() {
        for bit in 0..8u64 {
            let data = [1u8 << bit];
            let bm = BitmapReader::new(&data, 8);
            assert_eq!(bm.highest_set_bit(), Some(bit), "highest_set_bit for {bit}");

            // Inverse pattern: exactly one clear bit, at the same position.
            let inv = [!(1u8 << bit)];
            let bmi = BitmapReader::new(&inv, 8);
            assert_eq!(
                bmi.highest_clear_bit(),
                Some(bit),
                "highest_clear_bit for {bit}"
            );
        }

        // Multi-byte: the highest set bit is in an earlier byte than the last.
        let data = [0x00u8, 0b0001_0000, 0x00];
        assert_eq!(
            BitmapReader::new(&data, 24).highest_set_bit(),
            Some(8 + 4),
            "highest set bit should be byte 1, bit 4"
        );
    }

    #[test]
    fn test_single_byte_set_bits() {
        let data = [0b10100101u8]; // bits 0,2,5,7 set
        let bm = BitmapReader::new(&data, 8);
        assert!(bm.is_bit_set(0));
        assert!(!bm.is_bit_set(1));
        assert!(bm.is_bit_set(2));
        assert!(!bm.is_bit_set(3));
        assert!(!bm.is_bit_set(4));
        assert!(bm.is_bit_set(5));
        assert!(!bm.is_bit_set(6));
        assert!(bm.is_bit_set(7));
        assert_eq!(bm.count_set_bits(), 4);
        assert_eq!(bm.count_clear_bits(), 4);
        assert_eq!(bm.highest_set_bit(), Some(7));
    }

    #[test]
    fn test_partial_byte() {
        let data = [0xFF];
        let bm = BitmapReader::new(&data, 5); // only 5 bits valid
        assert_eq!(bm.count_set_bits(), 5); // not 8
        assert_eq!(bm.count_clear_bits(), 0);
        assert_eq!(bm.highest_set_bit(), Some(4));
    }

    #[test]
    fn test_partial_byte_with_clear_bits() {
        // 0b11110011 = bits 0,1,4,5,6,7 set; bits 2,3 clear
        let data = [0b11110011u8];
        let bm = BitmapReader::new(&data, 6); // only bits 0-5 valid
                                              // Valid bits: 0=1, 1=1, 2=0, 3=0, 4=1, 5=1 → 4 set
        assert_eq!(bm.count_set_bits(), 4);
        assert_eq!(bm.count_clear_bits(), 2);
        assert_eq!(bm.highest_set_bit(), Some(5));
    }

    #[test]
    fn test_iter_set_bits() {
        let data = [0b00000101u8]; // bits 0, 2
        let bm = BitmapReader::new(&data, 8);
        let bits: Vec<u64> = bm.iter_set_bits().collect();
        assert_eq!(bits, vec![0, 2]);
    }

    #[test]
    fn test_iter_clear_bits() {
        let data = [0b11111010u8]; // bits 1,3,4,5,6,7 set; bits 0,2 clear
        let bm = BitmapReader::new(&data, 8);
        let bits: Vec<u64> = bm.iter_clear_bits().collect();
        assert_eq!(bits, vec![0, 2]);
    }

    #[test]
    fn test_iter_set_bits_partial() {
        let data = [0b11111111u8];
        let bm = BitmapReader::new(&data, 3); // only bits 0,1,2 valid
        let bits: Vec<u64> = bm.iter_set_bits().collect();
        assert_eq!(bits, vec![0, 1, 2]);
    }

    #[test]
    fn test_multi_byte() {
        let data = [0xFF, 0x00, 0b00000011]; // byte 0: all set, byte 1: none, byte 2: bits 0,1
        let bm = BitmapReader::new(&data, 24);
        assert_eq!(bm.count_set_bits(), 10); // 8 + 0 + 2
        assert_eq!(bm.count_clear_bits(), 14);
        assert_eq!(bm.highest_set_bit(), Some(17)); // byte 2, bit 1 → index 16+1=17
    }

    #[test]
    fn test_multi_byte_partial() {
        let data = [0xFF, 0xFF, 0b00001111];
        let bm = BitmapReader::new(&data, 20); // 16 full + 4 partial
        assert_eq!(bm.count_set_bits(), 20); // 8+8+4
        assert_eq!(bm.highest_set_bit(), Some(19));
    }

    #[test]
    fn test_highest_set_bit_middle() {
        let data = [0x00, 0b10000000, 0x00]; // only bit 15 set (byte 1, bit 7)
        let bm = BitmapReader::new(&data, 24);
        assert_eq!(bm.highest_set_bit(), Some(15));
    }

    #[test]
    fn test_highest_set_bit_none() {
        let data = [0x00, 0x00];
        let bm = BitmapReader::new(&data, 16);
        assert_eq!(bm.highest_set_bit(), None);
    }

    #[test]
    fn test_highest_clear_bit_basic() {
        // 0b10110101: bits 1, 3, 6 clear; bits 0,2,4,5,7 set.
        let data = [0b10110101u8];
        let bm = BitmapReader::new(&data, 8);
        assert_eq!(bm.highest_clear_bit(), Some(6));
    }

    #[test]
    fn test_highest_clear_bit_all_set() {
        let data = [0xFF, 0xFF];
        let bm = BitmapReader::new(&data, 16);
        assert_eq!(bm.highest_clear_bit(), None);
    }

    #[test]
    fn test_highest_clear_bit_partial() {
        // Last byte partial: only bits 0,1,2 valid; byte is 0xFF, all set.
        let data = [0xFE, 0xFF];
        let bm = BitmapReader::new(&data, 11);
        // Byte 0 has bit 0 clear → highest clear bit is 0.
        assert_eq!(bm.highest_clear_bit(), Some(0));
    }

    #[test]
    fn test_highest_clear_bit_partial_with_clear() {
        // Byte 1 partial (bits 8,9,10 valid). Bit 10 set, bit 9 clear, bit 8 clear.
        let data = [0xFF, 0b00000101]; // byte 1: bits 0, 2 set; bits 1, 3..7 clear
        let bm = BitmapReader::new(&data, 11);
        // Valid bits in byte 1: bit 0 (=set), 1 (=clear), 2 (=set). Highest clear = 1 → index 9.
        assert_eq!(bm.highest_clear_bit(), Some(9));
    }

    #[test]
    fn test_highest_clear_bit_middle_byte() {
        // byte 0 = 0xFF (no clear), byte 1 = 0b01111111 (bit 7 clear),
        // byte 2 = 0xFF.
        let data = [0xFF, 0b01111111, 0xFF];
        let bm = BitmapReader::new(&data, 24);
        // Highest clear bit: byte 1 bit 7 = index 15.
        assert_eq!(bm.highest_clear_bit(), Some(15));
    }

    #[test]
    fn test_highest_clear_bit_empty() {
        let data: [u8; 0] = [];
        let bm = BitmapReader::new(&data, 0);
        assert_eq!(bm.highest_clear_bit(), None);
    }

    #[test]
    fn test_lowest_clear_bit() {
        let data = [0xFF, 0b11111110]; // byte 0: all set, byte 1: bit 0 clear
        let bm = BitmapReader::new(&data, 16);
        assert_eq!(bm.lowest_clear_bit(), Some(8)); // first clear bit is at index 8
    }

    #[test]
    fn test_lowest_clear_bit_first() {
        let data = [0b11111110]; // bit 0 clear
        let bm = BitmapReader::new(&data, 8);
        assert_eq!(bm.lowest_clear_bit(), Some(0));
    }

    #[test]
    fn test_lowest_clear_bit_none() {
        let data = [0xFF, 0xFF];
        let bm = BitmapReader::new(&data, 16);
        assert_eq!(bm.lowest_clear_bit(), None);
    }

    #[test]
    fn test_lowest_clear_bit_partial() {
        // All bits set in the full byte, partial byte has a clear bit
        let data = [0xFF, 0b11111011]; // bit 2 of byte 1 is clear
        let bm = BitmapReader::new(&data, 12); // 8 full + 4 partial
                                               // Valid bits in byte 1: 0=1, 1=1, 2=0, 3=1 → bit 10 is clear
        assert_eq!(bm.lowest_clear_bit(), Some(10));
    }

    #[test]
    fn test_empty_bitmap() {
        let data = [];
        let bm = BitmapReader::new(&data, 0);
        assert_eq!(bm.count_set_bits(), 0);
        assert_eq!(bm.count_clear_bits(), 0);
        assert_eq!(bm.highest_set_bit(), None);
        assert_eq!(bm.lowest_clear_bit(), None);
        assert!(!bm.is_bit_set(0));
        assert_eq!(bm.iter_set_bits().count(), 0);
        assert_eq!(bm.iter_clear_bits().count(), 0);
    }

    #[test]
    fn test_out_of_range() {
        let data = [0xFF];
        let bm = BitmapReader::new(&data, 8);
        assert!(!bm.is_bit_set(8));
        assert!(!bm.is_bit_set(100));
    }

    #[test]
    fn test_bit_count() {
        let data = [0xFF; 4];
        let bm = BitmapReader::new(&data, 30);
        assert_eq!(bm.bit_count(), 30);
    }

    #[test]
    fn test_large_bitmap() {
        // Simulate a 1024-block bitmap (128 bytes), all allocated
        let data = vec![0xFF; 128];
        let bm = BitmapReader::new(&data, 1024);
        assert_eq!(bm.count_set_bits(), 1024);
        assert_eq!(bm.count_clear_bits(), 0);
        assert_eq!(bm.highest_set_bit(), Some(1023));
        assert_eq!(bm.lowest_clear_bit(), None);
    }

    #[test]
    fn test_large_bitmap_sparse() {
        // 1024-bit bitmap with only every 8th bit set (bit 0 of each byte)
        let data = vec![0x01; 128];
        let bm = BitmapReader::new(&data, 1024);
        assert_eq!(bm.count_set_bits(), 128);
        assert_eq!(bm.count_clear_bits(), 896);
        assert_eq!(bm.highest_set_bit(), Some(1016)); // byte 127, bit 0 → 127*8=1016
        assert_eq!(bm.lowest_clear_bit(), Some(1)); // bit 1 of byte 0
    }

    #[test]
    fn test_iter_set_bits_multi_byte() {
        let data = [0b00000001, 0b10000000]; // bit 0 and bit 15
        let bm = BitmapReader::new(&data, 16);
        let bits: Vec<u64> = bm.iter_set_bits().collect();
        assert_eq!(bits, vec![0, 15]);
    }

    #[test]
    fn test_iter_clear_bits_all_set() {
        let data = [0xFF];
        let bm = BitmapReader::new(&data, 8);
        let bits: Vec<u64> = bm.iter_clear_bits().collect();
        assert!(bits.is_empty());
    }

    #[test]
    fn test_iter_set_bits_all_clear() {
        let data = [0x00];
        let bm = BitmapReader::new(&data, 8);
        let bits: Vec<u64> = bm.iter_set_bits().collect();
        assert!(bits.is_empty());
    }

    // --- Mutation helper tests ---

    #[test]
    fn test_bitmap_set_bit() {
        let mut data = [0u8; 2];
        bitmap_set_bit(&mut data, 0);
        assert_eq!(data[0], 0b00000001);
        bitmap_set_bit(&mut data, 7);
        assert_eq!(data[0], 0b10000001);
        bitmap_set_bit(&mut data, 8);
        assert_eq!(data[1], 0b00000001);
        bitmap_set_bit(&mut data, 15);
        assert_eq!(data[1], 0b10000001);
    }

    #[test]
    fn test_bitmap_clear_bit() {
        let mut data = [0xFF, 0xFF];
        bitmap_clear_bit(&mut data, 0);
        assert_eq!(data[0], 0b11111110);
        bitmap_clear_bit(&mut data, 7);
        assert_eq!(data[0], 0b01111110);
        bitmap_clear_bit(&mut data, 8);
        assert_eq!(data[1], 0b11111110);
    }

    #[test]
    fn test_bitmap_set_clear_roundtrip() {
        let mut data = [0u8; 4];
        for i in 0..32u64 {
            bitmap_set_bit(&mut data, i);
        }
        assert_eq!(data, [0xFF; 4]);
        for i in 0..32u64 {
            bitmap_clear_bit(&mut data, i);
        }
        assert_eq!(data, [0x00; 4]);
    }

    #[test]
    fn test_bitmap_find_clear_bit_basic() {
        let data = [0xFF, 0b11111110]; // first clear bit is 8
        assert_eq!(bitmap_find_clear_bit(&data, 16), Some(8));
    }

    #[test]
    fn test_bitmap_find_clear_bit_first() {
        let data = [0b11111110]; // bit 0 clear
        assert_eq!(bitmap_find_clear_bit(&data, 8), Some(0));
    }

    #[test]
    fn test_bitmap_find_clear_bit_none() {
        let data = [0xFF, 0xFF];
        assert_eq!(bitmap_find_clear_bit(&data, 16), None);
    }

    #[test]
    fn test_bitmap_find_clear_bit_partial() {
        let data = [0xFF, 0b00001111]; // bits 12-15 are clear, but only 12 valid
        assert_eq!(bitmap_find_clear_bit(&data, 12), None); // all 12 valid bits are set
    }

    #[test]
    fn test_bitmap_find_clear_bit_in_partial() {
        let data = [0xFF, 0b11111011]; // bit 10 clear
        assert_eq!(bitmap_find_clear_bit(&data, 12), Some(10));
    }
}