1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150

/*
 * Copyright 2022, Isaac Woods
 * SPDX-License-Identifier: MPL-2.0
 */

//! It's useful to be able to model an integral type such as `u32` as being a series of bits,
//! instead of a whole number. There are, of course, the usual bitwise operators for simple stuff,
//! but the `Bitmap` trait provides more complex, specific operations that are useful for bitmaps.
//!
//! A common use of the `Bitmap` trait is for memory allocators to track an area of pages, where
//! each bit represents a page. You might, for example, want to find a series of `n` zeros (which
//! would mark an area of `n` pages that are free to allocate) - the `alloc` method provides this
//! functionality.

use bit_field::{BitArray, BitField};
use core::{fmt::Debug, mem};
use num_traits::PrimInt;

pub trait Bitmap: Sized {
    /// Find `n` consecutive unset bits, set them and return the index of the first bit.
    fn alloc(&mut self, n: usize) -> Option<usize>;

    /// Free `n` previously allocated bits, starting at `index`.
    fn free(&mut self, index: usize, n: usize);
}

impl<T> Bitmap for T
where
    T: PrimInt + BitField + Debug,
{
    fn alloc(&mut self, n: usize) -> Option<usize> {
        let num_bits = 8 * mem::size_of::<Self>();
        let mask = Self::from((Self::one() << n) - Self::one()).unwrap();

        /*
         * For each bit before there are no longer `n` bits to the end, take the next `n` bits
         * and and them with a mask of `n` ones. If the result is zero, all the bits in
         * the slice must be 0 and so we've found a run of `n` zeros.
         */
        for i in 0..(num_bits - n) {
            if self.get_bits(i..(i + n)) & mask == Self::zero() {
                self.set_bits(i..(i + n), mask);
                return Some(i);
            }
        }

        None
    }

    fn free(&mut self, index: usize, n: usize) {
        assert_eq!(self.get_bits(index..(index + n)), Self::from((Self::one() << n) - Self::one()).unwrap());
        self.set_bits(index..(index + n), Self::zero());
    }
}

/// Like `Bitmap`, but for arrays. This is unfortunately needed due to conflicting implementations.
pub trait BitmapSlice: Sized {
    /// Find `n` consecutive unset bits, set them and return the index of the first bit.
    /// This is useful for memory managers using `Bitmap` to track free frames / pages.
    fn alloc(self, n: usize) -> Option<usize>;

    /// Free `n` previously allocated bits, starting at `index`.
    fn free(self, index: usize, n: usize);
}

impl BitmapSlice for &mut [u8] {
    fn alloc(self, n: usize) -> Option<usize> {
        let num_bits = 8 * self.len();
        let mask = (1 << n) - 1;

        for i in 0..(num_bits - n) {
            if self.get_bits(i..(i + n)) & mask == 0 {
                self.set_bits(i..(i + n), mask);
                return Some(i);
            }
        }

        None
    }

    fn free(self, index: usize, n: usize) {
        assert_eq!(self.get_bits(index..(index + n)), (1 << n) - 1);
        self.set_bits(index..(index + n), 0);
    }
}

#[test]
fn test_bitmap() {
    assert_eq!((0b10001u16).alloc(3), Some(1));
    assert_eq!((0b11_0000_1_000_111u16).alloc(4), Some(7));
    assert_eq!((0b1111_1111u8).alloc(1), None);
    assert_eq!((0b0110_1010u8).alloc(2), None);
}

#[test]
fn test_bitmap_array_simple() {
    /*
     * These might be a bit counterintuitive at first, because `BitArray` treats the array as a
     * little-endian set of bytes, so the LSB of the first byte is bit 0.
     */
    assert_eq!([0xff, 0xff, 0xff, 0xff, 0xff].alloc(3), None);
    assert_eq!([0xfe, 0xff, 0xff, 0xff].alloc(1), Some(0));
    assert_eq!([0b1111_0001, 0xff, 0xff].alloc(3), Some(1));
    assert_eq!([0b1111_1001, 0xff, 0xff].alloc(3), None);
}

#[test]
fn test_bitmap_array_multiple() {
    let mut bitmap = [0b1111_1100, 0xff, 0xff];

    // Make the first allocation
    let first = bitmap.alloc(1);
    assert_eq!(first, Some(0));
    assert_eq!(bitmap, [0b1111_1101, 0xff, 0xff]);

    // Make a second allocation
    let second = bitmap.alloc(1);
    assert_eq!(second, Some(1));
    assert_eq!(bitmap, [0b1111_1111, 0xff, 0xff]);

    // Try to make a third allocation - it should fail
    assert_eq!(bitmap.alloc(1), None);
}

#[test]
fn test_bitmap_free() {
    let mut bitmap = [0b1000_1000, 0xff, 0xff];

    // Make the first allocation
    let first = bitmap.alloc(2);
    assert_eq!(first, Some(0));
    assert_eq!(bitmap, [0b1000_1011, 0xff, 0xff]);

    // Make a second allocation
    let second = bitmap.alloc(3);
    assert_eq!(second, Some(4));
    assert_eq!(bitmap, [0b1111_1011, 0xff, 0xff]);

    // Make an allocation that doesn't fit
    assert_eq!(bitmap.alloc(3), None);

    // Free the first allocation
    bitmap.free(first.unwrap(), 2);
    assert_eq!(bitmap, [0b1111_1000, 0xff, 0xff]);

    // Make another allocation that now fits
    let third = bitmap.alloc(3);
    assert_eq!(third, Some(0));
    assert_eq!(bitmap, [0b1111_1111, 0xff, 0xff]);
}