alex/volcanic-mem

root / src / tlsf.rs

const SLI: u32 = 5;
const SL_COUNT: u32 = 1 << SLI;
const SMALL_SIZE: u64 = 256;
const SMALL_STEP: u64 = 64;
const SMALL_LISTS: u32 = (SMALL_SIZE / SMALL_STEP) as u32;
const CLASS_SHIFT: u32 = 7;
const CLASSES: usize = (64 - CLASS_SHIFT) as usize;

const FIT_BUDGET: usize = 256;
const NIL: u32 = u32::MAX;
const TAKEN: u32 = u32::MAX - 1;

#[derive(Clone, Copy)]
struct Node {
    offset: u64,
    size: u64,
    prev_phys: u32,
    next_phys: u32,
    /// `TAKEN` while allocated.
    prev_free: u32,
    /// Resource kind while allocated.
    next_free: u32,
}

#[inline]
fn class_of(size: u64) -> u32 {
    if size <= SMALL_SIZE {
        0
    } else {
        63 - size.leading_zeros() - CLASS_SHIFT
    }
}

#[inline]
fn second_of(size: u64, class: u32) -> u32 {
    if class == 0 {
        ((size - 1) / SMALL_STEP) as u32
    } else {
        (size >> (class + CLASS_SHIFT - SLI)) as u32 ^ SL_COUNT
    }
}

#[inline]
fn list_index(class: u32, second: u32) -> usize {
    if class == 0 {
        second as usize
    } else {
        (SMALL_LISTS + (class - 1) * SL_COUNT + second) as usize
    }
}

/// Smallest size whose list only holds blocks of at least `size` bytes.
#[inline]
fn round_up_to_list(size: u64) -> u64 {
    if size <= SMALL_SIZE {
        // Small lists step by 64 bytes but the first large lists by 8; do not skip past them.
        (size + SMALL_STEP - 1).min(SMALL_SIZE + 1)
    } else {
        size.saturating_add((1 << (63 - size.leading_zeros() - SLI)) - 1)
    }
}

#[inline]
const fn align_up(value: u64, align: u64) -> u64 {
    (value + align - 1) & !(align - 1)
}

/// Two-level segregated fit allocator over an abstract offset range.
///
/// Every operation is O(1) in the common case. Nodes live in a dense arena and link by
/// `u32` index. Adjacent free blocks are always coalesced, so the physical neighbors of a
/// free block are allocations. This lets buffer-image granularity be enforced by checking
/// just those two neighbors.
pub(crate) struct Tlsf {
    nodes: Vec<Node>,
    spare: u32,
    heads: Box<[u32]>,
    first_level: u64,
    second_level: [u32; CLASSES],
    /// Physically last node; while free it is the untouched end of the block.
    tail: u32,
    granularity: u64,
    size: u64,
    free: u64,
    allocations: u32,
}

impl Tlsf {
    pub fn new(size: u64, granularity: u64) -> Self {
        debug_assert!(size > 0 && granularity.is_power_of_two());
        let class = class_of(size);
        let mut tlsf = Self {
            nodes: Vec::with_capacity(16),
            spare: NIL,
            heads: vec![NIL; list_index(class, second_of(size, class)) + 1].into_boxed_slice(),
            first_level: 0,
            second_level: [0; CLASSES],
            tail: 0,
            granularity,
            size,
            free: 0,
            allocations: 0,
        };
        let root = tlsf.new_node(0, size, NIL, NIL);
        tlsf.insert_free(root);
        tlsf
    }

    #[inline]
    pub fn size(&self) -> u64 {
        self.size
    }

    #[inline]
    pub fn free_bytes(&self) -> u64 {
        self.free
    }

    #[inline]
    pub fn allocation_count(&self) -> u32 {
        self.allocations
    }

    #[inline]
    pub fn is_empty(&self) -> bool {
        self.allocations == 0
    }

    #[inline]
    pub fn allocation_size(&self, node: u32) -> u64 {
        self.nodes[node as usize].size
    }

    /// Returns `(node, offset)` of a fitting range, or `None` when no free block can hold it.
    #[inline]
    pub fn allocate(&mut self, size: u64, align: u64, kind: u32) -> Option<(u32, u64)> {
        self.allocate_tight(size, align, kind)
            .or_else(|| self.allocate_loose(size, align, kind))
    }

    /// O(1) probes for close fits that leave larger blocks intact. May fail while space remains.
    pub fn allocate_tight(&mut self, size: u64, align: u64, kind: u32) -> Option<(u32, u64)> {
        debug_assert!(size > 0 && align.is_power_of_two() && kind < TAKEN);
        if size > self.free {
            return None;
        }
        // Exact list head first: same-size churn reuses the hole it left behind.
        let exact_class = class_of(size);
        let exact_second = second_of(size, exact_class);
        if let Some(hit) = self.try_head(list_index(exact_class, exact_second), size, align, kind) {
            return Some(hit);
        }
        // Good fit: every block in this list holds `size`, though alignment may still reject it.
        if let Some(hit) = self.try_first_list_from(size, size, align, kind) {
            return Some(hit);
        }
        // Prefer the untouched end over splitting larger interior holes, which stay available
        // for later tight fits.
        if self.is_free(self.tail)
            && let Some(offset) = self.fit(self.tail, size, align, kind)
        {
            let tail = self.tail;
            self.commit(tail, offset, size, kind);
            return Some((tail, offset));
        }
        None
    }

    /// Completes [`Self::allocate_tight`], walking more candidates and splitting larger blocks.
    pub fn allocate_loose(&mut self, size: u64, align: u64, kind: u32) -> Option<(u32, u64)> {
        if size > self.free {
            return None;
        }
        // Bounded best-fit walk: ascending lists, at most `FIT_BUDGET` candidate checks.
        let exact_class = class_of(size);
        let mut budget = FIT_BUDGET;
        let mut next = Some((exact_class, second_of(size, exact_class)));
        while let Some((c, s)) = next
            && budget > 0
        {
            let mut node = self.heads[list_index(c, s)];
            while node != NIL && budget > 0 {
                budget -= 1;
                if let Some(offset) = self.fit(node, size, align, kind) {
                    self.commit(node, offset, size, kind);
                    return Some((node, offset));
                }
                node = self.nodes[node as usize].next_free;
            }
            next = self.first_list_from(c, s + 1);
        }
        // Every block in this list holds `size` at any alignment and granularity placement.
        let slack = if self.granularity > 1 {
            align.max(self.granularity) - 1 + self.granularity
        } else {
            align - 1
        };
        if slack > 0
            && let Some(hit) =
                self.try_first_list_from(size.saturating_add(slack), size, align, kind)
        {
            return Some(hit);
        }
        // Nearly full: exhaustive search over every block that might hold `size`.
        let class = class_of(size);
        let mut next = Some((class, second_of(size, class)));
        while let Some((c, s)) = next {
            if let Some(hit) = self.search_list(list_index(c, s), size, align, kind, usize::MAX) {
                return Some(hit);
            }
            next = self.first_list_from(c, s + 1);
        }
        None
    }

    #[inline]
    fn try_first_list_from(
        &mut self,
        min: u64,
        size: u64,
        align: u64,
        kind: u32,
    ) -> Option<(u32, u64)> {
        let rounded = round_up_to_list(min);
        let class = class_of(rounded);
        let (c, s) = self.first_list_from(class, second_of(rounded, class))?;
        self.try_head(list_index(c, s), size, align, kind)
    }

    #[inline]
    fn try_head(&mut self, list: usize, size: u64, align: u64, kind: u32) -> Option<(u32, u64)> {
        let head = *self.heads.get(list)?;
        if head == NIL {
            return None;
        }
        let offset = self.fit(head, size, align, kind)?;
        self.commit(head, offset, size, kind);
        Some((head, offset))
    }

    pub fn free(&mut self, node: u32) {
        let mut node = node;
        debug_assert!(!self.is_free(node), "double free of TLSF node {node}");
        self.allocations -= 1;

        let prev = self.nodes[node as usize].prev_phys;
        if prev != NIL && self.is_free(prev) {
            self.remove_free(prev);
            self.absorb_next(prev);
            node = prev;
        }
        let next = self.nodes[node as usize].next_phys;
        if next != NIL && self.is_free(next) {
            self.remove_free(next);
            self.absorb_next(node);
        }
        self.insert_free(node);
    }

    #[inline]
    fn is_free(&self, node: u32) -> bool {
        self.nodes[node as usize].prev_free != TAKEN
    }

    /// First non-empty list at or after `(class, second)` in size order.
    #[inline]
    fn first_list_from(&self, class: u32, second: u32) -> Option<(u32, u32)> {
        let class_idx = class as usize;
        if class_idx < CLASSES && second < SL_COUNT {
            let map = self.second_level[class_idx] & (!0 << second);
            if map != 0 {
                return Some((class, map.trailing_zeros()));
            }
        }
        let map = self.first_level & (!0 << (class + 1));
        if map == 0 {
            return None;
        }
        let class = map.trailing_zeros();
        Some((class, self.second_level[class as usize].trailing_zeros()))
    }

    #[inline]
    fn search_list(
        &mut self,
        list: usize,
        size: u64,
        align: u64,
        kind: u32,
        limit: usize,
    ) -> Option<(u32, u64)> {
        let mut node = *self.heads.get(list)?;
        for _ in 0..limit {
            if node == NIL {
                break;
            }
            if let Some(offset) = self.fit(node, size, align, kind) {
                self.commit(node, offset, size, kind);
                return Some((node, offset));
            }
            node = self.nodes[node as usize].next_free;
        }
        None
    }

    #[inline]
    fn fit(&self, node: u32, size: u64, align: u64, kind: u32) -> Option<u64> {
        let n = &self.nodes[node as usize];
        let end = n.offset + n.size;
        let mut offset = align_up(n.offset, align);
        if self.granularity > 1 {
            let page = !(self.granularity - 1);
            if n.prev_phys != NIL {
                let p = &self.nodes[n.prev_phys as usize];
                debug_assert_eq!(p.prev_free, TAKEN);
                if p.next_free != kind && (p.offset + p.size - 1) & page == offset & page {
                    offset = align_up(offset, self.granularity);
                }
            }
            if offset + size > end {
                return None;
            }
            if n.next_phys != NIL {
                let q = &self.nodes[n.next_phys as usize];
                if q.next_free != kind && (offset + size - 1) & page == q.offset & page {
                    return None;
                }
            }
        }
        (offset + size <= end).then_some(offset)
    }

    fn commit(&mut self, node: u32, offset: u64, size: u64, kind: u32) {
        self.remove_free(node);
        let Node {
            offset: start,
            size: len,
            mut prev_phys,
            mut next_phys,
            ..
        } = self.nodes[node as usize];

        if offset > start {
            let pad = self.new_node(start, offset - start, prev_phys, node);
            if prev_phys != NIL {
                self.nodes[prev_phys as usize].next_phys = pad;
            }
            self.insert_free(pad);
            prev_phys = pad;
        }
        let end = start + len;
        if offset + size < end {
            let rest = self.new_node(offset + size, end - offset - size, node, next_phys);
            if next_phys != NIL {
                self.nodes[next_phys as usize].prev_phys = rest;
            }
            self.insert_free(rest);
            if next_phys == NIL {
                self.tail = rest;
            }
            next_phys = rest;
        }
        let n = &mut self.nodes[node as usize];
        n.prev_phys = prev_phys;
        n.next_phys = next_phys;
        n.offset = offset;
        n.size = size;
        n.prev_free = TAKEN;
        n.next_free = kind;
        self.allocations += 1;
    }

    /// Merges `node.next_phys` into `node` and recycles it.
    #[inline]
    fn absorb_next(&mut self, node: u32) {
        let next = self.nodes[node as usize].next_phys;
        let Node {
            size, next_phys, ..
        } = self.nodes[next as usize];
        let n = &mut self.nodes[node as usize];
        n.size += size;
        n.next_phys = next_phys;
        if next_phys != NIL {
            self.nodes[next_phys as usize].prev_phys = node;
        } else {
            self.tail = node;
        }
        let dead = &mut self.nodes[next as usize];
        dead.size = 0;
        dead.prev_free = TAKEN;
        dead.next_free = self.spare;
        self.spare = next;
    }

    #[inline]
    fn new_node(&mut self, offset: u64, size: u64, prev_phys: u32, next_phys: u32) -> u32 {
        let node = Node {
            offset,
            size,
            prev_phys,
            next_phys,
            prev_free: TAKEN,
            next_free: NIL,
        };
        if self.spare != NIL {
            let index = self.spare;
            self.spare = self.nodes[index as usize].next_free;
            self.nodes[index as usize] = node;
            index
        } else {
            self.nodes.push(node);
            (self.nodes.len() - 1) as u32
        }
    }

    #[inline]
    fn insert_free(&mut self, node: u32) {
        let size = self.nodes[node as usize].size;
        let class = class_of(size);
        let second = second_of(size, class);
        let list = list_index(class, second);
        let head = self.heads[list];
        let n = &mut self.nodes[node as usize];
        n.prev_free = NIL;
        n.next_free = head;
        if head != NIL {
            self.nodes[head as usize].prev_free = node;
        }
        self.heads[list] = node;
        self.first_level |= 1 << class;
        self.second_level[class as usize] |= 1 << second;
        self.free += size;
    }

    #[inline]
    fn remove_free(&mut self, node: u32) {
        let Node {
            size,
            prev_free,
            next_free,
            ..
        } = self.nodes[node as usize];
        if next_free != NIL {
            self.nodes[next_free as usize].prev_free = prev_free;
        }
        if prev_free != NIL {
            self.nodes[prev_free as usize].next_free = next_free;
        } else {
            let class = class_of(size);
            let second = second_of(size, class);
            self.heads[list_index(class, second)] = next_free;
            if next_free == NIL {
                let map = &mut self.second_level[class as usize];
                *map &= !(1 << second);
                if *map == 0 {
                    self.first_level &= !(1 << class);
                }
            }
        }
        self.free -= size;
    }

    /// Checks every structural invariant. Panics with a description on violation.
    pub fn validate(&self) {
        let live = |i: usize| self.nodes[i].size != 0;
        let first = (0..self.nodes.len())
            .find(|&i| live(i) && self.nodes[i].prev_phys == NIL)
            .expect("no first node");

        let (mut offset, mut free, mut allocations, mut free_nodes) = (0, 0, 0, 0);
        let mut prev = NIL;
        let mut prev_free = false;
        let mut node = first as u32;
        while node != NIL {
            let n = &self.nodes[node as usize];
            assert_eq!(n.prev_phys, prev, "broken physical back link at {node}");
            assert_eq!(n.offset, offset, "gap or overlap at {node}");
            assert!(n.size > 0, "empty node {node}");
            let is_free = self.is_free(node);
            assert!(!(is_free && prev_free), "uncoalesced free nodes at {node}");
            if is_free {
                free += n.size;
                free_nodes += 1;
                let class = class_of(n.size);
                let second = second_of(n.size, class);
                assert!(
                    self.second_level[class as usize] & (1 << second) != 0,
                    "missing bitmap bit"
                );
                let mut cursor = self.heads[list_index(class, second)];
                while cursor != NIL && cursor != node {
                    cursor = self.nodes[cursor as usize].next_free;
                }
                assert_eq!(cursor, node, "free node {node} missing from its list");
            } else {
                allocations += 1;
            }
            prev_free = is_free;
            offset += n.size;
            prev = node;
            node = n.next_phys;
        }
        assert_eq!(offset, self.size, "nodes do not cover the block");
        assert_eq!(prev, self.tail, "tail is not the last node");
        assert_eq!(free, self.free, "free byte count mismatch");
        assert_eq!(allocations, self.allocations, "allocation count mismatch");

        let mut listed = 0;
        for (list, &head) in self.heads.iter().enumerate() {
            let mut cursor = head;
            let mut back = NIL;
            while cursor != NIL {
                let n = &self.nodes[cursor as usize];
                assert_eq!(
                    n.prev_free, back,
                    "broken free list back link in list {list}"
                );
                let class = class_of(n.size);
                assert_eq!(
                    list_index(class, second_of(n.size, class)),
                    list,
                    "node in wrong list"
                );
                listed += 1;
                back = cursor;
                cursor = n.next_free;
            }
        }
        assert_eq!(listed, free_nodes, "free lists hold unreachable nodes");
        for class in 0..CLASSES {
            let has = self.first_level & (1 << class) != 0;
            assert_eq!(
                has,
                self.second_level[class] != 0,
                "first level bitmap mismatch at {class}"
            );
        }
    }
}

#[cfg(test)]
mod tests {
    use super::Tlsf;
    use std::collections::BTreeMap;

    #[test]
    fn granularity_separates_linear_and_optimal_pages() {
        const GRANULARITY: u64 = 1024;
        let mut state = 0x1234_5678_u64;
        let mut rand = move |n: u64| {
            state ^= state << 13;
            state ^= state >> 7;
            state ^= state << 17;
            state % n
        };
        let mut tlsf = Tlsf::new(4 << 20, GRANULARITY);
        let mut live: Vec<(u32, u64)> = Vec::new();
        let mut ranges = BTreeMap::new();

        for step in 0..60_000 {
            if live.is_empty() || rand(100) < 55 {
                let bits = 4 + rand(12);
                let size = 1 + rand(1 << bits);
                let kind = rand(2) as u32;
                let align = 1 << rand(9);
                if let Some((node, offset)) = tlsf.allocate(size, align, kind) {
                    live.push((node, offset));
                    ranges.insert(offset, (offset + size, kind));
                }
            } else {
                let (node, offset) = live.swap_remove(rand(live.len() as u64) as usize);
                ranges.remove(&offset);
                tlsf.free(node);
            }
            if step % 256 == 0 {
                tlsf.validate();
                let mut prev: Option<(u64, u32)> = None;
                for (&start, &(end, kind)) in &ranges {
                    if let Some((prev_end, prev_kind)) = prev {
                        assert!(prev_end <= start, "overlap at {start}");
                        if prev_kind != kind {
                            assert_ne!(
                                (prev_end - 1) / GRANULARITY,
                                start / GRANULARITY,
                                "granularity conflict at {start}"
                            );
                        }
                    }
                    prev = Some((end, kind));
                }
            }
        }
    }
}