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

//! [`Buddies`](crate): A buddies allocator.
//!
//! This can be used to allocate blocks of different sizes from a single
//! fixed-width block, and is useful in bare-metal physical memory allocation
//! (which is how Linux does things).

#![no_std]

extern crate bitvec;
use bitvec::prelude::*;

/// [`RawBuddies`]: A slightly unsafe buddy allocator.
///
/// A small size and no standard library dependency is traded for an unsafe
/// structure. A safe shell can be constructed around this for built-in
/// allocation of resources as well as a safe allocation result.
pub struct RawBuddies<T> {
    /// The number of buddies.
    num: usize,
    /// A pointer to the first data element (size 2^num).
    data: *mut T,
    /// A pointer to the first bitspace byte.
    bits: *mut u8,
}

impl<T> RawBuddies<T> {
    /// Creates a new [`RawBuddies`].
    ///
    /// ### Conditions
    /// `data` and `bits` are not dropped as long as the instantiation lives.  
    /// `data` is at least of length `2^num`. It may be uninitialized.  
    /// `bits` is at least of length `2^num/8` (i.e it holds `2^num` bits). It
    /// must only contain `0`s (i.e `false`s).
    pub unsafe fn new(num: usize, data: *mut T, bits: *mut u8) -> Self {
        Self {
            num,
            data,
            bits,
        }
    }

    /// Checks if a block of size `2^n` `T`s can be allocated.
    ///
    /// ### Panics
    /// Panics if the block size is too large (`>= buddies`).
    pub fn can_allocate(&self, n: usize) -> bool {
        assert!(n < self.num);
        // Check if matching buddy contains free zones
        self.buddymap_ref(n).any()
    }

    /// Allocates a block of size `2^n` `T`s.
    ///
    /// Note for safe shells: You want to convert the pointer to a slice such
    /// that multiple (mutable) slices can be held simultaneously.
    ///
    /// Returns the reference as well as the block index (for freeing later).
    ///
    /// ### Panics
    /// Panics if the block size is too large (`>= buddies`).
    pub fn allocate(&mut self, n: usize) -> Option<(*mut T, usize)> {
        assert!(n < self.num);
        // Find a free zone for the buddy (auto-exit on fail)
        let pos = self.buddymap_ref(n).iter().position(|s| s)?;
        // Mark the network
        self.set_network(n, pos, true);
        // Return the block
        Some((unsafe { self.data.add(pos << n) }, pos))
    }

    /// Frees a given block by index and size.
    ///
    /// ### Panics
    /// Panics if the block size is too large (`>= buddies`).  
    /// Panics if the index is too large (`>= 2^(buddies-size-1)`).  
    /// Panics if the block was already free (possible double-free).  
    pub fn free(&mut self, n: usize, pos: usize) {
        assert!(n < self.num);
        assert!(pos < (1usize << (self.num - n - 1)));
        assert!(self.buddymap_ref(n)[pos]);
        // Drop the data
        unsafe { core::ptr::drop_in_place(self.data.add(pos << n)); }
        // Free the network
        self.set_network(n, pos, false);
    }

    /// Retrieves a bit slice for a certain buddy immutably.
    ///
    /// Primarily defined for [`can_allocate`].
    fn buddymap_ref(&self, n: usize) -> &BitSlice {
        assert!(n < self.num);
        // Index is 2^(num-n) from end
        let bits: &BitSlice = unsafe {
            core::slice::from_raw_parts(self.bits, 1usize << self.num)
        }.into();
        &bits[
            (1usize << self.num) - (1usize << (self.num - n))
         .. (1usize << self.num) - (1usize << (self.num - n - 1))
        ]
    }

    /// Retrieves a bit slice for a certain buddy mutably.
    fn buddymap_mut(&mut self, n: usize) -> &mut BitSlice {
        assert!(n < self.num);
        // Index is 2^(num-n) from end
        let bits: &mut BitSlice = unsafe {
            core::slice::from_raw_parts_mut(self.bits, 1usize << self.num)
        }.into();
        &mut bits[
            (1usize << self.num) - (1usize << (self.num - n))
         .. (1usize << self.num) - (1usize << (self.num - n - 1))
        ]
    }

    /// Sets a 'network' of bits around a single set.
    ///
    /// Primarily defined for [`allocate`] and [`free`].
    fn set_network(&mut self, n: usize, i: usize, v: bool) {
        assert!(n < self.num);
        assert!(i < (1 << n));
        // Begin by setting the lower network of bits
        for b in 0..n {
            self.buddymap_mut(b)[i << (n-b) .. (i+1) << (n-b)].set_all(v);
        }
        // Set the higher network of bits
        // v == true: Keep setting until we hit another true
        // v == false: Keep setting until the other is true
        if v {
            for b in n .. self.num {
                let map = self.buddymap_mut(b);
                if map[i >> (b-n)] {
                    break;
                }
                map.set(i >> (b-n), true);
            }
        } else {
            for b in n .. self.num {
                let map = self.buddymap_mut(b);
                map.set(i >> (b-n), false);
                if map[(i >> (b-n)) ^ 1] {
                    break;
                }
            }
        }
    }
}