From 3e4f1c6ec2958ee9f72d90c56e194ea3a2edf1be Mon Sep 17 00:00:00 2001 From: Patrick Owen Date: Sun, 2 Aug 2026 23:45:50 -0400 Subject: [PATCH] Refactor offset-allocator to improve readability This commit makes the following changes: * Add documentation to better explain how the allocator works * Use the slab crate for managing free node indexes * Set up a newtype for SmallFloat * Create `SmallFloatMap`, which SmallFloat is used as an index for * Create `BinsMap` as a layer of abstraction for `used_bins` etc. * Set up a custom NodeIndexOption type to remove nonmax crate dependency --- .gitignore | 1 + Cargo.toml | 2 +- src/allocator.rs | 619 +++++++++++++++++++++++++++++++++++++++++++++ src/bins_map.rs | 146 +++++++++++ src/ext.rs | 9 - src/lib.rs | 592 +------------------------------------------ src/node_index.rs | 98 +++++++ src/small_float.rs | 252 ++++++++++++++---- src/tests.rs | 276 -------------------- 9 files changed, 1072 insertions(+), 923 deletions(-) create mode 100644 src/allocator.rs create mode 100644 src/bins_map.rs delete mode 100644 src/ext.rs create mode 100644 src/node_index.rs delete mode 100644 src/tests.rs diff --git a/.gitignore b/.gitignore index ea8c4bf..96ef6c0 100644 --- a/.gitignore +++ b/.gitignore @@ -1 +1,2 @@ /target +Cargo.lock diff --git a/Cargo.toml b/Cargo.toml index 4e4b5ac..9547127 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -13,4 +13,4 @@ keywords = ["memory-management"] [dependencies] log = "0.4" -nonmax = "0.5" +slab = "0.4.12" diff --git a/src/allocator.rs b/src/allocator.rs new file mode 100644 index 0000000..aaa1c0b --- /dev/null +++ b/src/allocator.rs @@ -0,0 +1,619 @@ +// offset-allocator/src/allocator.rs + +use std::fmt::{Debug, Formatter}; + +use slab::Slab; +use log::debug; + +use crate::{ + bins_map::BinsMap, + node_index::{NodeIndex, NodeIndexOption}, + small_float::{SmallFloat, SmallFloatMap}, +}; + +/// An allocator that manages a single contiguous chunk of space and hands out +/// portions of it as requested. Since this allocator does not support alignment, it is recommended +/// to interpret these allocations in whatever unit is most convenient, which will likely not be "bytes" +pub struct Allocator { + /// The total size of the buffer + size: u32, + /// The maximum number of "nodes", or continuous blocks the allocator can handle. The actual supported number of allocations is less than this. + max_nodes: u32, + /// The total amount of remaining available space in the buffer. Fragmentation and rounding means that an allocation of this size is not always possible, + /// but as long as this is non-zero, and `max_nodes` isn't exceeded, it's always possible to create an allocation of size 1. + free_storage: u32, + /// A [`BinsMap`] that keeps track of all nodes that are not part of an existing allocation + bins_map: BinsMap, + /// Maintains the mapping from [`NodeIndex`] to [`Node`] + nodes: NodeSlab, +} + +/// A single allocation. +#[derive(Clone, Copy)] +pub struct Allocation { + /// The location of this allocation within the buffer. + pub offset: u32, + /// The node index associated with this allocation. + metadata: NI, +} + +/// Provides a summary of the state of the allocator, including space remaining. +#[derive(Debug)] +pub struct StorageReport { + /// The amount of free space left. + pub total_free_space: u32, + /// The maximum potential size of a single contiguous allocation. + pub largest_free_region: u32, +} + +/// Provides a detailed accounting of each bin within the allocator. +#[derive(Debug, Default)] +pub struct StorageReportFull { + /// Each bin within the allocator. + pub free_regions: SmallFloatMap, +} + +/// A detailed accounting of each allocator bin. +#[derive(Clone, Copy, Debug, Default)] +pub struct StorageReportFullRegion { + /// The size of the bin, in units. + pub size: u32, + /// The number of allocations in the bin. + pub count: u32, +} + +#[derive(Clone, Copy)] +struct Node { + /// The offset of the node in the buffer + data_offset: u32, + /// The size of the node in the buffer + data_size: u32, + /// Nodes representing free space are added to bins based on their size. Each bin can store an arbitrary number of nodes, + /// so we used a linked list. This stores the previous node in the bin. This field is meaningless when the node is used in an active allocation. + bin_list_prev: NodeIndexOption, + /// Nodes representing free space are added to bins based on their size. Each bin can store an arbitrary number of nodes, + /// so we used a linked list. This stores the next node in the bin. This field is meaningless when the node is used in an active allocation. + bin_list_next: NodeIndexOption, + /// The entire buffer is split up into several nodes, some marking an allocation and others marking free space. + /// Neighboring nodes in this buffer point to each other in a linked list. This field stores the index of the previous neighboring node. + neighbor_prev: NodeIndexOption, + /// The entire buffer is split up into several nodes, some marking an allocation and others marking free space. + /// Neighboring nodes in this buffer point to each other in a linked list. This field stores the index of the next neighboring node. + neighbor_next: NodeIndexOption, + /// Whether the node is used in an active allocation + used: bool, // Note: One possible enhancement to reduce the size of `Node` is to merge this with another field as a bit flag. +} + +/// A slab that stores [`Node`]s by their [`NodeIndex`]. The maximum capacity of this slab depends on the size of the [`NodeIndex`]. +struct NodeSlab(Slab>); + +impl NodeSlab { + /// Construct a new, empty `NodeSlab` + #[inline] + pub fn new() -> Self { + // TODO: To avoid potential pauses, we may want to pre-allocate the maximum number of allocations we need. + NodeSlab(Slab::new()) + } + + /// Return the number of stored nodes + #[inline] + pub fn len(&self) -> u32 { + self.0.len() as u32 + } + + /// Insert a node into the slab, returning the index associated with it + #[inline] + pub fn insert(&mut self, node: Node) -> NI { + assert!(self.len() != u32::MAX); + NI::from_u32(self.0.insert(node) as u32) + } + + /// Remove and return the node associated with the index + #[inline] + pub fn remove(&mut self, index: NI) -> Node { + self.0.remove(index.to_usize()) + } +} + +impl std::ops::Index for NodeSlab { + type Output = Node; + + #[inline] + fn index(&self, index: NI) -> &Self::Output { + &self.0[index.to_usize()] + } +} + +impl std::ops::IndexMut for NodeSlab { + #[inline] + fn index_mut(&mut self, index: NI) -> &mut Self::Output { + &mut self.0[index.to_usize()] + } +} + +impl Allocator { + /// Creates a new allocator, managing a contiguous block of memory of `size` + /// units, with the maximum allocations set as high as possible. + pub fn new(size: u32) -> Self { + Allocator::with_max_nodes(size, NI::NUM_VALID) + } + + /// Creates a new allocator, managing a contiguous block of memory of `size` + /// units, with the given number of maximum nodes. + /// + /// Note that even if no memory is freed, the maximum number of allocations + /// allowed is 1 less than the maximum number of nodes, since a node is needed + /// to keep track of the remaining free space. If memory is freed, due to fragmentation, + /// it is not guaranteed that another allocation will become available. + /// + /// Note also that the maximum number of nodes must be at most + /// [`NodeIndex::NUM_VALID`] and at least 1. If this restriction is violated, this + /// constructor will panic. + pub fn with_max_nodes(size: u32, max_nodes: u32) -> Self { + assert!(max_nodes > 0); + assert!(max_nodes <= NI::NUM_VALID); + + let mut this = Self { + size, + max_nodes, + free_storage: 0, + bins_map: BinsMap::default(), + nodes: NodeSlab::new(), + }; + this.insert_node_into_bin(size, 0); + this + } + + /// Clears out all allocations. + pub fn reset(&mut self) { + *self = Self::with_max_nodes(self.size, self.max_nodes); + } + + /// Allocates a block of `size` elements and returns its allocation. + /// + /// If there's not enough contiguous space for this allocation, returns + /// None. + pub fn allocate(&mut self, size: u32) -> Option> { + // Out of allocations? + if self.nodes.len() >= self.max_nodes { + return None; + } + + // Round up when finding the bin index to ensure that any node in that bin can hold the allocation + let min_bin_index = SmallFloat::from_u32_round_up(size); + let bin_index = self.bins_map.min_occupied_since(min_bin_index)?; + + // Pop the top node of the bin from the linked list + let node_index = self.bins_map[bin_index].unwrap(); + let node = &mut self.nodes[node_index]; + let node_total_size = node.data_size; + node.data_size = size; + node.used = true; + self.bins_map + .replace_bin_node(bin_index, node.bin_list_next); + if let Some(bin_list_next) = node.bin_list_next.to_option() { + self.nodes[bin_list_next].bin_list_prev = NodeIndexOption::NONE; + } + self.free_storage -= node_total_size; + debug!( + "Free storage: {} (-{}) (allocate)", + self.free_storage, node_total_size + ); + + // Push back remainder N elements to a (usually) lower bin + let remainder_size = node_total_size - size; + if remainder_size > 0 { + let Node { + data_offset, + neighbor_next, + .. + } = self.nodes[node_index]; + + let new_node_index = self.insert_node_into_bin(remainder_size, data_offset + size); + + // Link nodes next to each other so that we can merge them later if both are free + // And update the old next neighbor to point to the new node (in middle) + let node = &mut self.nodes[node_index]; + if let Some(neighbor_next) = node.neighbor_next.to_option() { + self.nodes[neighbor_next].neighbor_prev = NodeIndexOption::some(new_node_index); + } + self.nodes[new_node_index].neighbor_prev = NodeIndexOption::some(node_index); + self.nodes[new_node_index].neighbor_next = neighbor_next; + self.nodes[node_index].neighbor_next = NodeIndexOption::some(new_node_index); + } + + let node = &mut self.nodes[node_index]; + Some(Allocation { + offset: node.data_offset, + metadata: node_index, + }) + } + + /// Frees an allocation, returning the data to the heap. + /// + /// If the allocation has already been freed, the behavior is unspecified. + /// It may or may not panic. Note that the memory safety of the allocator *itself* will be + /// uncompromised, even on double free. + pub fn free(&mut self, allocation: Allocation) { + let node_index = allocation.metadata; + + // Merge with neighbors… + let Node { + data_offset: mut offset, + data_size: mut size, + used, + .. + } = self.nodes[node_index]; + + // Double delete check + assert!(used); + + if let Some(neighbor_prev) = self.nodes[node_index].neighbor_prev.to_option() { + if !self.nodes[neighbor_prev].used { + // Previous (contiguous) free node: Change offset to previous + // node offset. Sum sizes + let prev_node = &self.nodes[neighbor_prev]; + offset = prev_node.data_offset; + size += prev_node.data_size; + + let prev_node = &self.nodes[neighbor_prev]; + debug_assert_eq!(prev_node.neighbor_next, NodeIndexOption::some(node_index)); + self.nodes[node_index].neighbor_prev = prev_node.neighbor_prev; + self.remove_node_from_bin(neighbor_prev); + } + } + + if let Some(neighbor_next) = self.nodes[node_index].neighbor_next.to_option() { + if !self.nodes[neighbor_next].used { + // Next (contiguous) free node: Offset remains the same. Sum + // sizes. + let next_node = &self.nodes[neighbor_next]; + size += next_node.data_size; + + let next_node = &self.nodes[neighbor_next]; + debug_assert_eq!(next_node.neighbor_prev, NodeIndexOption::some(node_index)); + self.nodes[node_index].neighbor_next = next_node.neighbor_next; + self.remove_node_from_bin(neighbor_next); + } + } + + let Node { + neighbor_next, + neighbor_prev, + .. + } = self.nodes[node_index]; + + self.nodes.remove(node_index); + + // Insert the (combined) free node to bin + let combined_node_index = self.insert_node_into_bin(size, offset); + + // Connect neighbors with the new combined node + if let Some(neighbor_next) = neighbor_next.to_option() { + self.nodes[combined_node_index].neighbor_next = NodeIndexOption::some(neighbor_next); + self.nodes[neighbor_next].neighbor_prev = NodeIndexOption::some(combined_node_index); + } + if let Some(neighbor_prev) = neighbor_prev.to_option() { + self.nodes[combined_node_index].neighbor_prev = NodeIndexOption::some(neighbor_prev); + self.nodes[neighbor_prev].neighbor_next = NodeIndexOption::some(combined_node_index); + } + } + + /// Creates a new free [`Node`] and inserts it at the head of the appropriate bin. Note that the caller of this + /// function is responsible for linking node in the "neighbor" linked list. + fn insert_node_into_bin(&mut self, size: u32, data_offset: u32) -> NI { + // Round down when finding the bin index to ensure that the node being put in that bin can hold any allocation associated with that bin + let bin_index = SmallFloat::from_u32_round_down(size); + + // Create a new node and insert on top of the bin linked list + let top_node_index = self.bins_map[bin_index]; + let node_index = self.nodes.insert(Node { + data_offset, + data_size: size, + bin_list_prev: NodeIndexOption::NONE, + bin_list_next: top_node_index, + neighbor_prev: NodeIndexOption::NONE, + neighbor_next: NodeIndexOption::NONE, + used: false, + }); + if let Some(top_node_index) = top_node_index.to_option() { + self.nodes[top_node_index].bin_list_prev = NodeIndexOption::some(node_index); + } + self.bins_map + .replace_bin_node(bin_index, NodeIndexOption::some(node_index)); + + self.free_storage += size; + debug!( + "Free storage: {} (+{}) (insert_node_into_bin)", + self.free_storage, size + ); + node_index + } + + /// Deletes a [`Node`], removing it from the bin. Note that the caller of this + /// function is responsible for fixing up links in the "neighbor" linked list, and it is recommended + /// that this fixup occur before this function is called. + fn remove_node_from_bin(&mut self, node_index: NI) { + // Copy the node to work around borrow check. + let node = self.nodes[node_index]; + + match node.bin_list_prev.to_option() { + Some(bin_list_prev) => { + // Easy case: We have previous node. Just remove this node from the middle of the list. + self.nodes[bin_list_prev].bin_list_next = node.bin_list_next; + if let Some(bin_list_next) = node.bin_list_next.to_option() { + self.nodes[bin_list_next].bin_list_prev = node.bin_list_prev; + } + } + None => { + // Hard case: We are the first node in a bin. Find the bin. + + // Round down when finding the bin index to ensure consistency with `insert_node_into_bin` + let bin_index = SmallFloat::from_u32_round_down(node.data_size); + + self.bins_map + .replace_bin_node(bin_index, node.bin_list_next); + if let Some(bin_list_next) = node.bin_list_next.to_option() { + self.nodes[bin_list_next].bin_list_prev = NodeIndexOption::NONE; + } + } + } + + self.nodes.remove(node_index); + + self.free_storage -= node.data_size; + debug!( + "Free storage: {} (-{}) (remove_node_from_bin)", + self.free_storage, node.data_size + ); + } + + /// Returns the *used* size of an allocation. + /// + /// For this allocator, this always equals the size requested at allocation time. + pub fn allocation_size(&self, allocation: Allocation) -> u32 { + self.nodes[allocation.metadata].data_size + } + + /// Returns a structure containing the amount of free space remaining, as + /// well as the largest amount that can be allocated at once. + pub fn storage_report(&self) -> StorageReport { + if self.nodes.len() >= self.max_nodes { + // Out of allocations? -> Zero free space + return StorageReport { + total_free_space: 0, + largest_free_region: 0, + }; + } + + let largest_free_region = self.bins_map.max_occupied().map_or(0, |x| x.to_u32()); + debug_assert!(self.free_storage >= largest_free_region); + + StorageReport { + total_free_space: self.free_storage, + largest_free_region, + } + } + + /// Returns detailed information about the number of allocations in each bin. + pub fn storage_report_full(&self) -> StorageReportFull { + let mut report = StorageReportFull::default(); + for i in SmallFloat::values() { + let mut count = 0; + let mut maybe_node_index = self.bins_map[i]; + while let Some(node_index) = maybe_node_index.to_option() { + maybe_node_index = self.nodes[node_index].bin_list_next; + count += 1; + } + report.free_regions[i] = StorageReportFullRegion { + size: i.to_u32(), + count, + } + } + report + } +} + +impl Debug for Allocator { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + self.storage_report().fmt(f) + } +} + +/// Returns the minimum allocator size needed to hold an object of the given size. +pub fn min_allocator_size(needed_object_size: u32) -> u32 { + SmallFloat::from_u32_round_up(needed_object_size).to_u32() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn basic_offset_allocator() { + let mut allocator: Allocator = Allocator::new(1024 * 1024 * 256); + let a = allocator.allocate(1337).unwrap(); + let offset = a.offset; + assert_eq!(offset, 0); + allocator.free(a); + } + + #[test] + fn allocate_offset_allocator_simple() { + let mut allocator: Allocator = Allocator::new(1024 * 1024 * 256); + + // Free merges neighbor empty nodes. Next allocation should also have offset = 0 + let a = allocator.allocate(0).unwrap(); + assert_eq!(a.offset, 0); + + let b = allocator.allocate(1).unwrap(); + assert_eq!(b.offset, 0); + + let c = allocator.allocate(123).unwrap(); + assert_eq!(c.offset, 1); + + let d = allocator.allocate(1234).unwrap(); + assert_eq!(d.offset, 124); + + allocator.free(a); + allocator.free(b); + allocator.free(c); + allocator.free(d); + + // End: Validate that allocator has no fragmentation left. Should be 100% clean. + let validate_all = allocator.allocate(1024 * 1024 * 256).unwrap(); + assert_eq!(validate_all.offset, 0); + allocator.free(validate_all); + } + + #[test] + fn allocate_offset_allocator_merge_trivial() { + let mut allocator: Allocator = Allocator::new(1024 * 1024 * 256); + + // Free merges neighbor empty nodes. Next allocation should also have offset = 0 + let a = allocator.allocate(1337).unwrap(); + assert_eq!(a.offset, 0); + allocator.free(a); + + let b = allocator.allocate(1337).unwrap(); + assert_eq!(b.offset, 0); + allocator.free(b); + + // End: Validate that allocator has no fragmentation left. Should be 100% clean. + let validate_all = allocator.allocate(1024 * 1024 * 256).unwrap(); + assert_eq!(validate_all.offset, 0); + allocator.free(validate_all); + } + + #[test] + fn allocate_offset_allocator_reuse_trivial() { + let mut allocator: Allocator = Allocator::new(1024 * 1024 * 256); + + // Allocator should reuse node freed by A since the allocation C fits in the same bin (using pow2 size to be sure) + let a = allocator.allocate(1024).unwrap(); + assert_eq!(a.offset, 0); + + let b = allocator.allocate(3456).unwrap(); + assert_eq!(b.offset, 1024); + + allocator.free(a); + + let c = allocator.allocate(1024).unwrap(); + assert_eq!(c.offset, 0); + + allocator.free(c); + allocator.free(b); + + // End: Validate that allocator has no fragmentation left. Should be 100% clean. + let validate_all = allocator.allocate(1024 * 1024 * 256).unwrap(); + assert_eq!(validate_all.offset, 0); + allocator.free(validate_all); + } + + #[test] + fn allocate_offset_allocator_reuse_complex() { + let mut allocator: Allocator = Allocator::new(1024 * 1024 * 256); + + // Allocator should not reuse node freed by A since the allocation C doesn't fits in the same bin + // However node D and E fit there and should reuse node from A + let a = allocator.allocate(1024).unwrap(); + assert_eq!(a.offset, 0); + + let b = allocator.allocate(3456).unwrap(); + assert_eq!(b.offset, 1024); + + allocator.free(a); + + let c = allocator.allocate(2345).unwrap(); + assert_eq!(c.offset, 1024 + 3456); + + let d = allocator.allocate(456).unwrap(); + assert_eq!(d.offset, 0); + + let e = allocator.allocate(512).unwrap(); + assert_eq!(e.offset, 456); + + let report = allocator.storage_report(); + assert_eq!( + report.total_free_space, + 1024 * 1024 * 256 - 3456 - 2345 - 456 - 512 + ); + assert_ne!(report.largest_free_region, report.total_free_space); + + allocator.free(c); + allocator.free(d); + allocator.free(b); + allocator.free(e); + + // End: Validate that allocator has no fragmentation left. Should be 100% clean. + let validate_all = allocator.allocate(1024 * 1024 * 256).unwrap(); + assert_eq!(validate_all.offset, 0); + allocator.free(validate_all); + } + + #[test] + fn allocate_offset_allocator_zero_fragmentation() { + let mut allocator: Allocator = Allocator::new(1024 * 1024 * 256); + + // Allocate 256x 1MB. Should fit. Then free four random slots and reallocate four slots. + // Plus free four contiguous slots an allocate 4x larger slot. All must be zero fragmentation! + let mut allocations: [_; 256] = std::array::from_fn(|i| { + let allocation = allocator.allocate(1024 * 1024).unwrap(); + assert_eq!(allocation.offset, i as u32 * 1024 * 1024); + allocation + }); + + let report = allocator.storage_report(); + assert_eq!(report.total_free_space, 0); + assert_eq!(report.largest_free_region, 0); + + // Free four random slots + allocator.free(allocations[243]); + allocator.free(allocations[5]); + allocator.free(allocations[123]); + allocator.free(allocations[95]); + + // Free four contiguous slots (allocator must merge) + allocator.free(allocations[151]); + allocator.free(allocations[152]); + allocator.free(allocations[153]); + allocator.free(allocations[154]); + + allocations[243] = allocator.allocate(1024 * 1024).unwrap(); + allocations[5] = allocator.allocate(1024 * 1024).unwrap(); + allocations[123] = allocator.allocate(1024 * 1024).unwrap(); + allocations[95] = allocator.allocate(1024 * 1024).unwrap(); + allocations[151] = allocator.allocate(1024 * 1024 * 4).unwrap(); // 4x larger + + for (i, allocation) in allocations.iter().enumerate() { + if !(152..155).contains(&i) { + allocator.free(*allocation); + } + } + + let report2 = allocator.storage_report(); + assert_eq!(report2.total_free_space, 1024 * 1024 * 256); + assert_eq!(report2.largest_free_region, 1024 * 1024 * 256); + + // End: Validate that allocator has no fragmentation left. Should be 100% clean. + let validate_all = allocator.allocate(1024 * 1024 * 256).unwrap(); + assert_eq!(validate_all.offset, 0); + allocator.free(validate_all); + } + + #[test] + fn min_allocator_size() { + // Randomly generated integers on a log distribution, σ = 10. + static TEST_OBJECT_SIZES: [u32; 42] = [ + 0, 1, 2, 3, 4, 5, 8, 17, 23, 36, 51, 68, 87, 151, 165, 167, 201, 223, 306, 346, 394, + 411, 806, 969, 1404, 1798, 2236, 4281, 4745, 13989, 21095, 26594, 27146, 29679, 144685, + 153878, 495127, 727999, 1377073, 9440387, 41994490, 68520116, + ]; + + for needed_object_size in TEST_OBJECT_SIZES { + let allocator_size = super::min_allocator_size(needed_object_size); + let mut allocator: Allocator = Allocator::new(allocator_size); + assert!(allocator.allocate(needed_object_size).is_some()); + } + } +} diff --git a/src/bins_map.rs b/src/bins_map.rs new file mode 100644 index 0000000..82255bf --- /dev/null +++ b/src/bins_map.rs @@ -0,0 +1,146 @@ +// offset-allocator/src/bins_map.rs + +use crate::{ + node_index::{NodeIndex, NodeIndexOption}, + small_float::{SmallFloat, SmallFloatMap}, +}; + +const NUM_TOP_BINS: usize = 32; +const TOP_BINS_INDEX_SHIFT: u32 = 3; +const LEAF_BINS_INDEX_MASK: u32 = 7; + +/// A map from each bin to the node at the head of the linked list for that bin. The name of this struct is `BinsMap` instead of `BinMap` to avoid confusion with binary maps. +pub struct BinsMap { + /// A bit-vector showing which `occupied_bins` entries are nonzero, used for faster lookup of nonempty bins + occupied_bins_top: u32, + /// An array of 32 bit-vectors that show which bins are nonempty, used for faster lookup of nonempty bins + occupied_bins: [u8; NUM_TOP_BINS], + /// A map that points to the head node of each bin + bins: SmallFloatMap>, +} + +impl Default for BinsMap { + fn default() -> Self { + Self { + occupied_bins_top: 0, + occupied_bins: [0; NUM_TOP_BINS], + bins: SmallFloatMap::default(), + } + } +} + +impl BinsMap { + /// Returns the minimum bin index greater than or equal to `min` that corresponds to a nonempty bin + pub fn min_occupied_since(&self, min: SmallFloat) -> Option { + /// Out of bits at position greater than or equal to `start_bit_index`, Returns the position of the + /// lowest-position bit that is set to 1. Return `None` if there is no such bit. + fn find_lowest_bit_set_after(bit_mask: u32, start_bit_index: u32) -> Option { + let mask_before_start_index = (1 << start_bit_index) - 1; + let mask_after_start_index = !mask_before_start_index; + let bits_after = bit_mask & mask_after_start_index; + if bits_after == 0 { + None + } else { + Some(bits_after.trailing_zeros()) + } + } + + let min_top_bin_index = min.reinterpret_as_u32() >> TOP_BINS_INDEX_SHIFT; + let min_leaf_bin_index = min.reinterpret_as_u32() & LEAF_BINS_INDEX_MASK; + + let mut top_bin_index = min_top_bin_index; + let mut leaf_bin_index = None; + + // If top bin exists, scan its leaf bin. This can fail (NO_SPACE). + if (self.occupied_bins_top & (1 << top_bin_index)) != 0 { + leaf_bin_index = find_lowest_bit_set_after( + self.occupied_bins[top_bin_index as usize] as _, + min_leaf_bin_index, + ); + } + + // If we didn't find space in top bin, we search top bin from +1 + let leaf_bin_index = match leaf_bin_index { + Some(leaf_bin_index) => leaf_bin_index, + None => { + top_bin_index = + find_lowest_bit_set_after(self.occupied_bins_top, min_top_bin_index + 1)?; + + // All leaf bins here fit the alloc, since the top bin was + // rounded up. Start leaf search from bit 0. + // + // NOTE: This search can't fail since at least one leaf bit was + // set because the top bit was set. + self.occupied_bins[top_bin_index as usize].trailing_zeros() + } + }; + + Some(SmallFloat::reinterpret_u32( + (top_bin_index << TOP_BINS_INDEX_SHIFT) | leaf_bin_index, + )) + } + + /// Returns the maximum bin index that corresponds to a nonempty bin + pub fn max_occupied(&self) -> Option { + if self.occupied_bins_top == 0 { + return None; + } + let top_bin_index = self.occupied_bins_top.ilog2(); + let leaf_bin_index = (self.occupied_bins[top_bin_index as usize] as u32).ilog2(); + Some(SmallFloat::reinterpret_u32( + (top_bin_index << TOP_BINS_INDEX_SHIFT) | leaf_bin_index, + )) + } + + /// Replace the [`NodeIndexOption`] pointed to by the specific bin index with a new [`NodeIndexOption`] + pub fn replace_bin_node( + &mut self, + bin_index: SmallFloat, + node: NodeIndexOption, + ) -> NodeIndexOption { + let old_node = std::mem::replace(&mut self.bins[bin_index], node); + if node.is_none() && !old_node.is_none() { + // Newly empty + self.mark_bin_empty(bin_index); + } else if !node.is_none() && old_node.is_none() { + // Newly filled + self.mark_bin_occupied(bin_index); + } + old_node + } + + /// Internal method to ensure that [`Self::occupied_bins`] and [`Self::occupied_bins_top`] are correct + /// after a bin has been emptied out + fn mark_bin_empty(&mut self, bin_index: SmallFloat) { + let top_bin_index = bin_index.reinterpret_as_u32() >> TOP_BINS_INDEX_SHIFT; + let leaf_bin_index = bin_index.reinterpret_as_u32() & LEAF_BINS_INDEX_MASK; + + // Remove a leaf bin mask bit + self.occupied_bins[top_bin_index as usize] &= !(1 << u32::from(leaf_bin_index)); + + // All leaf bins empty? + if self.occupied_bins[top_bin_index as usize] == 0 { + // Remove a top bin mask bit + self.occupied_bins_top &= !(1 << top_bin_index); + } + } + + /// Internal method to ensure that [`Self::occupied_bins`] and [`Self::occupied_bins_top`] are correct + /// after an empty bin has been occupied + fn mark_bin_occupied(&mut self, bin_index: SmallFloat) { + let top_bin_index = bin_index.reinterpret_as_u32() >> TOP_BINS_INDEX_SHIFT; + let leaf_bin_index = bin_index.reinterpret_as_u32() & LEAF_BINS_INDEX_MASK; + + // Set bin mask bits + self.occupied_bins[top_bin_index as usize] |= 1 << leaf_bin_index; + self.occupied_bins_top |= 1 << top_bin_index; + } +} + +impl std::ops::Index for BinsMap { + type Output = NodeIndexOption; + + fn index(&self, index: SmallFloat) -> &Self::Output { + &self.bins[index] + } +} diff --git a/src/ext.rs b/src/ext.rs deleted file mode 100644 index e0d73ae..0000000 --- a/src/ext.rs +++ /dev/null @@ -1,9 +0,0 @@ -//! Extension functions not present in the original C++ `OffsetAllocator`. - -use crate::small_float; - -/// Returns the minimum allocator size needed to hold an object of the given -/// size. -pub fn min_allocator_size(needed_object_size: u32) -> u32 { - small_float::float_to_uint(small_float::uint_to_float_round_up(needed_object_size)) -} diff --git a/src/lib.rs b/src/lib.rs index 98ab025..8a92499 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -4,589 +4,11 @@ #![deny(unsafe_code)] #![warn(missing_docs)] -use std::fmt::{Debug, Display, Formatter, Result as FmtResult}; - -use log::debug; -use nonmax::{NonMaxU16, NonMaxU32}; - -pub mod ext; - +mod allocator; +mod bins_map; +mod node_index; mod small_float; - -#[cfg(test)] -mod tests; - -const NUM_TOP_BINS: usize = 32; -const BINS_PER_LEAF: usize = 8; -const TOP_BINS_INDEX_SHIFT: u32 = 3; -const LEAF_BINS_INDEX_MASK: u32 = 7; -const NUM_LEAF_BINS: usize = NUM_TOP_BINS * BINS_PER_LEAF; - -/// Determines the number of allocations that the allocator supports. -/// -/// By default, [`Allocator`] and related functions use `u32`, which allows for -/// `u32::MAX - 1` allocations. You can, however, use `u16` instead, which -/// causes the allocator to use less memory but limits the number of allocations -/// within a single allocator to at most 65,534. -pub trait NodeIndex: Clone + Copy + Default { - /// The `NonMax` version of this type. - /// - /// This is used extensively to optimize `enum` representations. - type NonMax: NodeIndexNonMax + TryFrom + Into; - - /// The maximum value representable in this type. - const MAX: u32; - - /// Converts from a unsigned 32-bit integer to an instance of this type. - fn from_u32(val: u32) -> Self; - - /// Converts this type to an unsigned machine word. - fn to_usize(self) -> usize; -} - -/// The `NonMax` version of the [`NodeIndex`]. -/// -/// For example, for `u32`, the `NonMax` version is [`NonMaxU32`]. -pub trait NodeIndexNonMax: Clone + Copy + PartialEq + Default + Debug + Display { - /// Converts this type to an unsigned machine word. - fn to_usize(self) -> usize; -} - -/// An allocator that manages a single contiguous chunk of space and hands out -/// portions of it as requested. -pub struct Allocator -where - NI: NodeIndex, -{ - size: u32, - max_allocs: u32, - free_storage: u32, - - used_bins_top: u32, - used_bins: [u8; NUM_TOP_BINS], - bin_indices: [Option; NUM_LEAF_BINS], - - nodes: Vec>, - free_nodes: Vec, - free_offset: u32, -} - -/// A single allocation. -#[derive(Clone, Copy)] -pub struct Allocation -where - NI: NodeIndex, -{ - /// The location of this allocation within the buffer. - pub offset: NI, - /// The node index associated with this allocation. - metadata: NI::NonMax, -} - -/// Provides a summary of the state of the allocator, including space remaining. -#[derive(Debug)] -pub struct StorageReport { - /// The amount of free space left. - pub total_free_space: u32, - /// The maximum potential size of a single contiguous allocation. - pub largest_free_region: u32, -} - -/// Provides a detailed accounting of each bin within the allocator. -#[derive(Debug)] -pub struct StorageReportFull { - /// Each bin within the allocator. - pub free_regions: [StorageReportFullRegion; NUM_LEAF_BINS], -} - -/// A detailed accounting of each allocator bin. -#[derive(Clone, Copy, Debug, Default)] -pub struct StorageReportFullRegion { - /// The size of the bin, in units. - pub size: u32, - /// The number of allocations in the bin. - pub count: u32, -} - -#[derive(Clone, Copy, Default)] -struct Node -where - NI: NodeIndex, -{ - data_offset: u32, - data_size: u32, - bin_list_prev: Option, - bin_list_next: Option, - neighbor_prev: Option, - neighbor_next: Option, - used: bool, // TODO: Merge as bit flag -} - -// Utility functions -fn find_lowest_bit_set_after(bit_mask: u32, start_bit_index: u32) -> Option { - let mask_before_start_index = (1 << start_bit_index) - 1; - let mask_after_start_index = !mask_before_start_index; - let bits_after = bit_mask & mask_after_start_index; - if bits_after == 0 { - None - } else { - NonMaxU32::try_from(bits_after.trailing_zeros()).ok() - } -} - -impl Allocator -where - NI: NodeIndex, -{ - /// Creates a new allocator, managing a contiguous block of memory of `size` - /// units, with a default reasonable number of maximum allocations. - pub fn new(size: u32) -> Self { - Allocator::with_max_allocs(size, u32::min(128 * 1024, NI::MAX - 1)) - } - - /// Creates a new allocator, managing a contiguous block of memory of `size` - /// units, with the given number of maximum allocations. - /// - /// Note that the maximum number of allocations must be less than - /// [`NodeIndex::MAX`] minus one. If this restriction is violated, this - /// constructor will panic. - pub fn with_max_allocs(size: u32, max_allocs: u32) -> Self { - assert!(max_allocs < NI::MAX - 1); - - let mut this = Self { - size, - max_allocs, - free_storage: 0, - used_bins_top: 0, - free_offset: 0, - used_bins: [0; NUM_TOP_BINS], - bin_indices: [None; NUM_LEAF_BINS], - nodes: vec![], - free_nodes: vec![], - }; - this.reset(); - this - } - - /// Clears out all allocations. - pub fn reset(&mut self) { - self.free_storage = 0; - self.used_bins_top = 0; - self.free_offset = self.max_allocs - 1; - - self.used_bins.iter_mut().for_each(|bin| *bin = 0); - - self.bin_indices.iter_mut().for_each(|index| *index = None); - - self.nodes = vec![Node::default(); self.max_allocs as usize]; - - // Freelist is a stack. Nodes in inverse order so that [0] pops first. - self.free_nodes = (0..self.max_allocs) - .map(|i| { - NI::NonMax::try_from(NI::from_u32(self.max_allocs - i - 1)).unwrap_or_default() - }) - .collect(); - - // Start state: Whole storage as one big node - // Algorithm will split remainders and push them back as smaller nodes - self.insert_node_into_bin(self.size, 0); - } - - /// Allocates a block of `size` elements and returns its allocation. - /// - /// If there's not enough contiguous space for this allocation, returns - /// None. - pub fn allocate(&mut self, size: u32) -> Option> { - // Out of allocations? - if self.free_offset == 0 { - return None; - } - - // Round up to bin index to ensure that alloc >= bin - // Gives us min bin index that fits the size - let min_bin_index = small_float::uint_to_float_round_up(size); - - let min_top_bin_index = min_bin_index >> TOP_BINS_INDEX_SHIFT; - let min_leaf_bin_index = min_bin_index & LEAF_BINS_INDEX_MASK; - - let mut top_bin_index = min_top_bin_index; - let mut leaf_bin_index = None; - - // If top bin exists, scan its leaf bin. This can fail (NO_SPACE). - if (self.used_bins_top & (1 << top_bin_index)) != 0 { - leaf_bin_index = find_lowest_bit_set_after( - self.used_bins[top_bin_index as usize] as _, - min_leaf_bin_index, - ); - } - - // If we didn't find space in top bin, we search top bin from +1 - let leaf_bin_index = match leaf_bin_index { - Some(leaf_bin_index) => leaf_bin_index, - None => { - top_bin_index = - find_lowest_bit_set_after(self.used_bins_top, min_top_bin_index + 1)?.into(); - - // All leaf bins here fit the alloc, since the top bin was - // rounded up. Start leaf search from bit 0. - // - // NOTE: This search can't fail since at least one leaf bit was - // set because the top bit was set. - NonMaxU32::try_from(self.used_bins[top_bin_index as usize].trailing_zeros()) - .unwrap() - } - }; - - let bin_index = (top_bin_index << TOP_BINS_INDEX_SHIFT) | u32::from(leaf_bin_index); - - // Pop the top node of the bin. Bin top = node.next. - let node_index = self.bin_indices[bin_index as usize].unwrap(); - let node = &mut self.nodes[node_index.to_usize()]; - let node_total_size = node.data_size; - node.data_size = size; - node.used = true; - self.bin_indices[bin_index as usize] = node.bin_list_next; - if let Some(bin_list_next) = node.bin_list_next { - self.nodes[bin_list_next.to_usize()].bin_list_prev = None; - } - self.free_storage -= node_total_size; - debug!( - "Free storage: {} (-{}) (allocate)", - self.free_storage, node_total_size - ); - - // Bin empty? - if self.bin_indices[bin_index as usize].is_none() { - // Remove a leaf bin mask bit - self.used_bins[top_bin_index as usize] &= !(1 << u32::from(leaf_bin_index)); - - // All leaf bins empty? - if self.used_bins[top_bin_index as usize] == 0 { - // Remove a top bin mask bit - self.used_bins_top &= !(1 << top_bin_index); - } - } - - // Push back remainder N elements to a lower bin - let remainder_size = node_total_size - size; - if remainder_size > 0 { - let Node { - data_offset, - neighbor_next, - .. - } = self.nodes[node_index.to_usize()]; - - let new_node_index = self.insert_node_into_bin(remainder_size, data_offset + size); - - // Link nodes next to each other so that we can merge them later if both are free - // And update the old next neighbor to point to the new node (in middle) - let node = &mut self.nodes[node_index.to_usize()]; - if let Some(neighbor_next) = node.neighbor_next { - self.nodes[neighbor_next.to_usize()].neighbor_prev = Some(new_node_index); - } - self.nodes[new_node_index.to_usize()].neighbor_prev = Some(node_index); - self.nodes[new_node_index.to_usize()].neighbor_next = neighbor_next; - self.nodes[node_index.to_usize()].neighbor_next = Some(new_node_index); - } - - let node = &mut self.nodes[node_index.to_usize()]; - Some(Allocation { - offset: NI::from_u32(node.data_offset), - metadata: node_index, - }) - } - - /// Frees an allocation, returning the data to the heap. - /// - /// If the allocation has already been freed, the behavior is unspecified. - /// It may or may not panic. Note that, because this crate contains no - /// unsafe code, the memory safe of the allocator *itself* will be - /// uncompromised, even on double free. - pub fn free(&mut self, allocation: Allocation) { - let node_index = allocation.metadata; - - // Merge with neighbors… - let Node { - data_offset: mut offset, - data_size: mut size, - used, - .. - } = self.nodes[node_index.to_usize()]; - - // Double delete check - assert!(used); - - if let Some(neighbor_prev) = self.nodes[node_index.to_usize()].neighbor_prev { - if !self.nodes[neighbor_prev.to_usize()].used { - // Previous (contiguous) free node: Change offset to previous - // node offset. Sum sizes - let prev_node = &self.nodes[neighbor_prev.to_usize()]; - offset = prev_node.data_offset; - size += prev_node.data_size; - - // Remove node from the bin linked list and put it in the - // freelist - self.remove_node_from_bin(neighbor_prev); - - let prev_node = &self.nodes[neighbor_prev.to_usize()]; - debug_assert_eq!(prev_node.neighbor_next, Some(node_index)); - self.nodes[node_index.to_usize()].neighbor_prev = prev_node.neighbor_prev; - } - } - - if let Some(neighbor_next) = self.nodes[node_index.to_usize()].neighbor_next { - if !self.nodes[neighbor_next.to_usize()].used { - // Next (contiguous) free node: Offset remains the same. Sum - // sizes. - let next_node = &self.nodes[neighbor_next.to_usize()]; - size += next_node.data_size; - - // Remove node from the bin linked list and put it in the - // freelist - self.remove_node_from_bin(neighbor_next); - - let next_node = &self.nodes[neighbor_next.to_usize()]; - debug_assert_eq!(next_node.neighbor_prev, Some(node_index)); - self.nodes[node_index.to_usize()].neighbor_next = next_node.neighbor_next; - } - } - - let Node { - neighbor_next, - neighbor_prev, - .. - } = self.nodes[node_index.to_usize()]; - - // Insert the removed node to freelist - debug!( - "Putting node {} into freelist[{}] (free)", - node_index, - self.free_offset + 1 - ); - self.free_offset += 1; - self.free_nodes[self.free_offset as usize] = node_index; - - // Insert the (combined) free node to bin - let combined_node_index = self.insert_node_into_bin(size, offset); - - // Connect neighbors with the new combined node - if let Some(neighbor_next) = neighbor_next { - self.nodes[combined_node_index.to_usize()].neighbor_next = Some(neighbor_next); - self.nodes[neighbor_next.to_usize()].neighbor_prev = Some(combined_node_index); - } - if let Some(neighbor_prev) = neighbor_prev { - self.nodes[combined_node_index.to_usize()].neighbor_prev = Some(neighbor_prev); - self.nodes[neighbor_prev.to_usize()].neighbor_next = Some(combined_node_index); - } - } - - fn insert_node_into_bin(&mut self, size: u32, data_offset: u32) -> NI::NonMax { - // Round down to bin index to ensure that bin >= alloc - let bin_index = small_float::uint_to_float_round_down(size); - - let top_bin_index = bin_index >> TOP_BINS_INDEX_SHIFT; - let leaf_bin_index = bin_index & LEAF_BINS_INDEX_MASK; - - // Bin was empty before? - if self.bin_indices[bin_index as usize].is_none() { - // Set bin mask bits - self.used_bins[top_bin_index as usize] |= 1 << leaf_bin_index; - self.used_bins_top |= 1 << top_bin_index; - } - - // Take a freelist node and insert on top of the bin linked list (next = old top) - let top_node_index = self.bin_indices[bin_index as usize]; - let free_offset = self.free_offset; - let node_index = self.free_nodes[free_offset as usize]; - self.free_offset -= 1; - debug!( - "Getting node {} from freelist[{}]", - node_index, - self.free_offset + 1 - ); - self.nodes[node_index.to_usize()] = Node { - data_offset, - data_size: size, - bin_list_next: top_node_index, - ..Node::default() - }; - if let Some(top_node_index) = top_node_index { - self.nodes[top_node_index.to_usize()].bin_list_prev = Some(node_index); - } - self.bin_indices[bin_index as usize] = Some(node_index); - - self.free_storage += size; - debug!( - "Free storage: {} (+{}) (insert_node_into_bin)", - self.free_storage, size - ); - node_index - } - - fn remove_node_from_bin(&mut self, node_index: NI::NonMax) { - // Copy the node to work around borrow check. - let node = self.nodes[node_index.to_usize()]; - - match node.bin_list_prev { - Some(bin_list_prev) => { - // Easy case: We have previous node. Just remove this node from the middle of the list. - self.nodes[bin_list_prev.to_usize()].bin_list_next = node.bin_list_next; - if let Some(bin_list_next) = node.bin_list_next { - self.nodes[bin_list_next.to_usize()].bin_list_prev = node.bin_list_prev; - } - } - None => { - // Hard case: We are the first node in a bin. Find the bin. - - // Round down to bin index to ensure that bin >= alloc - let bin_index = small_float::uint_to_float_round_down(node.data_size); - - let top_bin_index = (bin_index >> TOP_BINS_INDEX_SHIFT) as usize; - let leaf_bin_index = (bin_index & LEAF_BINS_INDEX_MASK) as usize; - - self.bin_indices[bin_index as usize] = node.bin_list_next; - if let Some(bin_list_next) = node.bin_list_next { - self.nodes[bin_list_next.to_usize()].bin_list_prev = None; - } - - // Bin empty? - if self.bin_indices[bin_index as usize].is_none() { - // Remove a leaf bin mask bit - self.used_bins[top_bin_index as usize] &= !(1 << leaf_bin_index); - - // All leaf bins empty? - if self.used_bins[top_bin_index as usize] == 0 { - // Remove a top bin mask bit - self.used_bins_top &= !(1 << top_bin_index); - } - } - } - } - - // Insert the node to freelist - debug!( - "Putting node {} into freelist[{}] (remove_node_from_bin)", - node_index, - self.free_offset + 1 - ); - self.free_offset += 1; - self.free_nodes[self.free_offset as usize] = node_index; - - self.free_storage -= node.data_size; - debug!( - "Free storage: {} (-{}) (remove_node_from_bin)", - self.free_storage, node.data_size - ); - } - - /// Returns the *used* size of an allocation. - /// - /// Note that this may be larger than the size requested at allocation time, - /// due to rounding. - pub fn allocation_size(&self, allocation: Allocation) -> u32 { - self.nodes - .get(allocation.metadata.to_usize()) - .map(|node| node.data_size) - .unwrap_or_default() - } - - /// Returns a structure containing the amount of free space remaining, as - /// well as the largest amount that can be allocated at once. - pub fn storage_report(&self) -> StorageReport { - let mut largest_free_region = 0; - let mut free_storage = 0; - - // Out of allocations? -> Zero free space - if self.free_offset > 0 { - free_storage = self.free_storage; - if self.used_bins_top > 0 { - let top_bin_index = 31 - self.used_bins_top.leading_zeros(); - let leaf_bin_index = - 31 - (self.used_bins[top_bin_index as usize] as u32).leading_zeros(); - largest_free_region = small_float::float_to_uint( - (top_bin_index << TOP_BINS_INDEX_SHIFT) | leaf_bin_index, - ); - debug_assert!(free_storage >= largest_free_region); - } - } - - StorageReport { - total_free_space: free_storage, - largest_free_region, - } - } - - /// Returns detailed information about the number of allocations in each - /// bin. - pub fn storage_report_full(&self) -> StorageReportFull { - let mut report = StorageReportFull::default(); - for i in 0..NUM_LEAF_BINS { - let mut count = 0; - let mut maybe_node_index = self.bin_indices[i]; - while let Some(node_index) = maybe_node_index { - maybe_node_index = self.nodes[node_index.to_usize()].bin_list_next; - count += 1; - } - report.free_regions[i] = StorageReportFullRegion { - size: small_float::float_to_uint(i as u32), - count, - } - } - report - } -} - -impl Default for StorageReportFull { - fn default() -> Self { - Self { - free_regions: [Default::default(); NUM_LEAF_BINS], - } - } -} - -impl Debug for Allocator -where - NI: NodeIndex, -{ - fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult { - self.storage_report().fmt(f) - } -} - -impl NodeIndex for u32 { - type NonMax = NonMaxU32; - const MAX: u32 = u32::MAX; - - fn from_u32(val: u32) -> Self { - val - } - - fn to_usize(self) -> usize { - self as usize - } -} - -impl NodeIndex for u16 { - type NonMax = NonMaxU16; - const MAX: u32 = u16::MAX as u32; - - fn from_u32(val: u32) -> Self { - val as u16 - } - - fn to_usize(self) -> usize { - self as usize - } -} - -impl NodeIndexNonMax for NonMaxU32 { - fn to_usize(self) -> usize { - u32::from(self) as usize - } -} - -impl NodeIndexNonMax for NonMaxU16 { - fn to_usize(self) -> usize { - u16::from(self) as usize - } -} +pub use allocator::{ + min_allocator_size, Allocation, Allocator, StorageReport, StorageReportFull, + StorageReportFullRegion, +}; diff --git a/src/node_index.rs b/src/node_index.rs new file mode 100644 index 0000000..9622054 --- /dev/null +++ b/src/node_index.rs @@ -0,0 +1,98 @@ +// offset-allocator/src/node_index.rs + +use std::fmt::{Debug, Display}; + +/// The index used to identify nodes in the allocator. Determines the number of allocations +/// that the allocator supports. +/// +/// By default, [`Allocator`] and related functions use `u32`, which allows for +/// `u32::MAX` allocations. You can, however, use `u16` instead, which +/// causes the allocator to use less memory but limits the number of allocations +/// within a single allocator to at most 65,535. +pub trait NodeIndex: Display + Debug + Clone + Copy + PartialEq + Eq { + /// An invalid representation in its type, used as the `None` type of `NodeIndexOption`. + const INVALID: Self; + + /// The number of indexes, consecutive starting from 0, that are valid representations + const NUM_VALID: u32; + + /// Converts from a unsigned 32-bit integer to an instance of this type. + fn from_u32(val: u32) -> Self; + + /// Converts this type to an unsigned machine word. + fn to_usize(self) -> usize; +} + +/// A type much like [`Option`] but made to use the maximum integer value as the `None` +/// value instead of requiring a separate discriminant. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct NodeIndexOption(NI); + +impl NodeIndexOption { + /// Equivalent to [`Option::None`] + pub const NONE: Self = NodeIndexOption(NodeIndex::INVALID); + + /// Initializes what is equivalent to an [`Option::Some`] for the given node index + pub fn some(inner: NI) -> Self { + Self(inner) + } + + /// Converts to the [`Option`] type for easier processing + #[inline] + pub fn to_option(self) -> Option { + if self == Self::NONE { + None + } else { + Some(self.0) + } + } + + /// Whether the option holds no value + #[inline] + pub fn is_none(self) -> bool { + self == Self::NONE + } + + /// Returns the value contained within the option. Panics if there is no such option. + #[inline] + pub fn unwrap(self) -> NI { + assert!(self != Self::NONE); + self.0 + } +} + +impl Default for NodeIndexOption { + fn default() -> Self { + Self::NONE + } +} + +impl NodeIndex for u32 { + const INVALID: u32 = u32::MAX; + + const NUM_VALID: u32 = Self::INVALID; + + fn from_u32(val: u32) -> Self { + assert!(val < Self::NUM_VALID); + val + } + + fn to_usize(self) -> usize { + self as usize + } +} + +impl NodeIndex for u16 { + const INVALID: u16 = u16::MAX; + + const NUM_VALID: u32 = Self::INVALID as u32; + + fn from_u32(val: u32) -> Self { + assert!(val < Self::NUM_VALID); + val as u16 + } + + fn to_usize(self) -> usize { + self as usize + } +} diff --git a/src/small_float.rs b/src/small_float.rs index 563869a..b755968 100644 --- a/src/small_float.rs +++ b/src/small_float.rs @@ -1,65 +1,213 @@ // offset-allocator/src/small_float.rs -pub const MANTISSA_BITS: u32 = 3; -pub const MANTISSA_VALUE: u32 = 1 << MANTISSA_BITS; -pub const MANTISSA_MASK: u32 = MANTISSA_VALUE - 1; - -// Bin sizes follow floating point (exponent + mantissa) distribution (piecewise linear log approx) -// This ensures that for each size class, the average overhead percentage stays the same -pub fn uint_to_float_round_up(size: u32) -> u32 { - let mut exp = 0; - let mut mantissa; - - if size < MANTISSA_VALUE { - // Denorm: 0..(MANTISSA_VALUE-1) - mantissa = size - } else { - // Normalized: Hidden high bit always 1. Not stored. Just like float. - let leading_zeros = size.leading_zeros(); - let highest_set_bit = 31 - leading_zeros; - - let mantissa_start_bit = highest_set_bit - MANTISSA_BITS; - exp = mantissa_start_bit + 1; - mantissa = (size >> mantissa_start_bit) & MANTISSA_MASK; - - let low_bits_mask = (1 << mantissa_start_bit) - 1; - - // Round up! - if (size & low_bits_mask) != 0 { - mantissa += 1; +//! This module handles operations around [`SmallFloat`], a custom struct that represents an +//! 8-bit unsigned floating point value that represents integer values using a 3-bit mantissa +//! and a 5-bit exponent. Each of these 256 values correspond to a specific bin, which determines +//! the size of the allocations supported by each bin. + +/// An 8-bit unsigned floating point value representing an integer using a 3-bit mantissa and a 5-bit exponent +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct SmallFloat(u32); + +impl SmallFloat { + /// The number of bits that represent the mantissa + const MANTISSA_BITS: u32 = 3; + + /// The number of bits that represent the exponent + const EXPONENT_BITS: u32 = 5; + + /// The number of possible values that can be stored in this `SmallFloat` + const NUM_VALUES: usize = 1 << (Self::MANTISSA_BITS + Self::EXPONENT_BITS); + + /// The number of possible mantissa values. This number is a power of 2. + const MANTISSA_VALUE: u32 = 1 << Self::MANTISSA_BITS; + + /// A mask that can be bitwise-anded with the float to get just the mantissa + const MANTISSA_MASK: u32 = Self::MANTISSA_VALUE - 1; + + /// All possible values of a [`SmallFloat`] from smallest to largest + pub fn values() -> impl ExactSizeIterator { + (0..Self::NUM_VALUES).map(|i| Self(i as u32)) + } + + /// The least [`SmallFloat`] greater than or equal to the given value + pub fn from_u32_round_up(value: u32) -> Self { + let mut exp = 0; + let mut mantissa; + + if value < Self::MANTISSA_VALUE { + // Denorm: 0..(MANTISSA_VALUE-1) + mantissa = value + } else { + // Normalized: Hidden high bit always 1. Not stored. Just like float. + let highest_set_bit = value.ilog2(); + let mantissa_start_bit = highest_set_bit - Self::MANTISSA_BITS; + exp = mantissa_start_bit + 1; + mantissa = (value >> mantissa_start_bit) & Self::MANTISSA_MASK; + + let low_bits_mask = (1 << mantissa_start_bit) - 1; + + // Round up! + if (value & low_bits_mask) != 0 { + mantissa += 1; + } } + + // Using `+` instead of `|` allows mantissa->exp overflow for round up + SmallFloat((exp << Self::MANTISSA_BITS) + mantissa) } - // + allows mantissa->exp overflow for round up - (exp << MANTISSA_BITS) + mantissa + /// The greatest [`SmallFloat`] less than or equal to the given value + pub fn from_u32_round_down(value: u32) -> Self { + let mut exp = 0; + let mantissa; + + if value < Self::MANTISSA_VALUE { + // Denorm: 0..(MANTISSA_VALUE-1) + mantissa = value + } else { + // Normalized: Hidden high bit always 1. Not stored. Just like float. + let highest_set_bit = value.ilog2(); + let mantissa_start_bit = highest_set_bit - Self::MANTISSA_BITS; + exp = mantissa_start_bit + 1; + mantissa = (value >> mantissa_start_bit) & Self::MANTISSA_MASK; + } + + SmallFloat((exp << Self::MANTISSA_BITS) | mantissa) + } + + /// The `u32` that holds the same value as the [`SmallFloat`] + pub fn to_u32(self) -> u32 { + let exponent = self.0 >> Self::MANTISSA_BITS; + let mantissa = self.0 & Self::MANTISSA_MASK; + if exponent == 0 { + mantissa + } else { + (mantissa | Self::MANTISSA_VALUE) << (exponent - 1) + } + } + + /// Reinterprets the bits of the [`SmallFloat`] as a `u32` instead + #[inline] + pub fn reinterpret_as_u32(self) -> u32 { + self.0 + } + + /// Reinterprets the bits of the `u32` as a [`SmallFloat`] instead + #[inline] + pub fn reinterpret_u32(data: u32) -> Self { + Self(data) + } } -pub fn uint_to_float_round_down(size: u32) -> u32 { - let mut exp = 0; - let mantissa; - - if size < MANTISSA_VALUE { - // Denorm: 0..(MANTISSA_VALUE-1) - mantissa = size - } else { - // Normalized: Hidden high bit always 1. Not stored. Just like float. - let leading_zeros = size.leading_zeros(); - let highest_set_bit = 31 - leading_zeros; - - let mantissa_start_bit = highest_set_bit - MANTISSA_BITS; - exp = mantissa_start_bit + 1; - mantissa = (size >> mantissa_start_bit) & MANTISSA_MASK; +/// A map whose key is a [`SmallFloat`]. Internally represented as an array +#[derive(Debug)] +pub struct SmallFloatMap([T; SmallFloat::NUM_VALUES]); + +impl Default for SmallFloatMap { + fn default() -> Self { + Self([T::default(); SmallFloat::NUM_VALUES]) } +} + +impl std::ops::Index for SmallFloatMap { + type Output = T; - (exp << MANTISSA_BITS) | mantissa + fn index(&self, index: SmallFloat) -> &Self::Output { + &self.0[index.0 as usize] + } } -pub fn float_to_uint(float_value: u32) -> u32 { - let exponent = float_value >> MANTISSA_BITS; - let mantissa = float_value & MANTISSA_MASK; - if exponent == 0 { - mantissa - } else { - (mantissa | MANTISSA_VALUE) << (exponent - 1) +impl std::ops::IndexMut for SmallFloatMap { + fn index_mut(&mut self, index: SmallFloat) -> &mut Self::Output { + &mut self.0[index.0 as usize] + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn small_float_uint_to_float() { + // Denorms, exp=1 and exp=2 + mantissa = 0 are all precise. + // NOTE: Assuming 8 value (3 bit) mantissa. + // If this test fails, please change this assumption! + let precise_number_count = 17; + for i in 0..precise_number_count { + let round_up = SmallFloat::from_u32_round_up(i); + let round_down = SmallFloat::from_u32_round_down(i); + assert_eq!(SmallFloat::reinterpret_u32(i), round_up); + assert_eq!(SmallFloat::reinterpret_u32(i), round_down); + } + + // Test some random picked numbers + struct NumberFloatUpDown { + number: u32, + up: SmallFloat, + down: SmallFloat, + } + + let test_data = [ + NumberFloatUpDown { + number: 17, + up: SmallFloat::reinterpret_u32(17), + down: SmallFloat::reinterpret_u32(16), + }, + NumberFloatUpDown { + number: 118, + up: SmallFloat::reinterpret_u32(39), + down: SmallFloat::reinterpret_u32(38), + }, + NumberFloatUpDown { + number: 1024, + up: SmallFloat::reinterpret_u32(64), + down: SmallFloat::reinterpret_u32(64), + }, + NumberFloatUpDown { + number: 65536, + up: SmallFloat::reinterpret_u32(112), + down: SmallFloat::reinterpret_u32(112), + }, + NumberFloatUpDown { + number: 529445, + up: SmallFloat::reinterpret_u32(137), + down: SmallFloat::reinterpret_u32(136), + }, + NumberFloatUpDown { + number: 1048575, + up: SmallFloat::reinterpret_u32(144), + down: SmallFloat::reinterpret_u32(143), + }, + ]; + + for v in test_data { + let round_up = SmallFloat::from_u32_round_up(v.number); + let round_down = SmallFloat::from_u32_round_down(v.number); + assert_eq!(round_up, v.up); + assert_eq!(round_down, v.down); + } + } + + #[test] + fn small_float_float_to_uint() { + // Denorms, exp=1 and exp=2 + mantissa = 0 are all precise. + // NOTE: Assuming 8 value (3 bit) mantissa. + // If this test fails, please change this assumption! + let precise_number_count = 17; + for i in 0..precise_number_count { + let v = SmallFloat::reinterpret_u32(i).to_u32(); + assert_eq!(i, v); + } + + // Test that float->uint->float conversion is precise for all numbers + // NOTE: Test values < 240. 240->4G = overflows 32 bit integer + for i in (0..240).map(|i| SmallFloat::reinterpret_u32(i)) { + let v = i.to_u32(); + let round_up = SmallFloat::from_u32_round_up(v); + let round_down = SmallFloat::from_u32_round_down(v); + assert_eq!(i, round_up); + assert_eq!(i, round_down); + } } } diff --git a/src/tests.rs b/src/tests.rs deleted file mode 100644 index 31d274b..0000000 --- a/src/tests.rs +++ /dev/null @@ -1,276 +0,0 @@ -// offset-allocator/src/tests.rs - -use std::array; - -use crate::{ext, small_float, Allocator}; - -#[test] -fn small_float_uint_to_float() { - // Denorms, exp=1 and exp=2 + mantissa = 0 are all precise. - // NOTE: Assuming 8 value (3 bit) mantissa. - // If this test fails, please change this assumption! - let precise_number_count = 17; - for i in 0..precise_number_count { - let round_up = small_float::uint_to_float_round_up(i); - let round_down = small_float::uint_to_float_round_down(i); - assert_eq!(i, round_up); - assert_eq!(i, round_down); - } - - // Test some random picked numbers - struct NumberFloatUpDown { - number: u32, - up: u32, - down: u32, - } - - let test_data = [ - NumberFloatUpDown { - number: 17, - up: 17, - down: 16, - }, - NumberFloatUpDown { - number: 118, - up: 39, - down: 38, - }, - NumberFloatUpDown { - number: 1024, - up: 64, - down: 64, - }, - NumberFloatUpDown { - number: 65536, - up: 112, - down: 112, - }, - NumberFloatUpDown { - number: 529445, - up: 137, - down: 136, - }, - NumberFloatUpDown { - number: 1048575, - up: 144, - down: 143, - }, - ]; - - for v in test_data { - let round_up = small_float::uint_to_float_round_up(v.number); - let round_down = small_float::uint_to_float_round_down(v.number); - assert_eq!(round_up, v.up); - assert_eq!(round_down, v.down); - } -} - -#[test] -fn small_float_float_to_uint() { - // Denorms, exp=1 and exp=2 + mantissa = 0 are all precise. - // NOTE: Assuming 8 value (3 bit) mantissa. - // If this test fails, please change this assumption! - let precise_number_count = 17; - for i in 0..precise_number_count { - let v = small_float::float_to_uint(i); - assert_eq!(i, v); - } - - // Test that float->uint->float conversion is precise for all numbers - // NOTE: Test values < 240. 240->4G = overflows 32 bit integer - for i in 0..240 { - let v = small_float::float_to_uint(i); - let round_up = small_float::uint_to_float_round_up(v); - let round_down = small_float::uint_to_float_round_down(v); - assert_eq!(i, round_up); - assert_eq!(i, round_down); - } -} - -#[test] -fn basic_offset_allocator() { - let mut allocator = Allocator::new(1024 * 1024 * 256); - let a = allocator.allocate(1337).unwrap(); - let offset: u32 = a.offset; - assert_eq!(offset, 0); - allocator.free(a); -} - -#[test] -fn allocate_offset_allocator_simple() { - let mut allocator: Allocator = Allocator::new(1024 * 1024 * 256); - - // Free merges neighbor empty nodes. Next allocation should also have offset = 0 - let a = allocator.allocate(0).unwrap(); - assert_eq!(a.offset, 0); - - let b = allocator.allocate(1).unwrap(); - assert_eq!(b.offset, 0); - - let c = allocator.allocate(123).unwrap(); - assert_eq!(c.offset, 1); - - let d = allocator.allocate(1234).unwrap(); - assert_eq!(d.offset, 124); - - allocator.free(a); - allocator.free(b); - allocator.free(c); - allocator.free(d); - - // End: Validate that allocator has no fragmentation left. Should be 100% clean. - let validate_all = allocator.allocate(1024 * 1024 * 256).unwrap(); - assert_eq!(validate_all.offset, 0); - allocator.free(validate_all); -} - -#[test] -fn allocate_offset_allocator_merge_trivial() { - let mut allocator: Allocator = Allocator::new(1024 * 1024 * 256); - - // Free merges neighbor empty nodes. Next allocation should also have offset = 0 - let a = allocator.allocate(1337).unwrap(); - assert_eq!(a.offset, 0); - allocator.free(a); - - let b = allocator.allocate(1337).unwrap(); - assert_eq!(b.offset, 0); - allocator.free(b); - - // End: Validate that allocator has no fragmentation left. Should be 100% clean. - let validate_all = allocator.allocate(1024 * 1024 * 256).unwrap(); - assert_eq!(validate_all.offset, 0); - allocator.free(validate_all); -} - -#[test] -fn allocate_offset_allocator_reuse_trivial() { - let mut allocator: Allocator = Allocator::new(1024 * 1024 * 256); - - // Allocator should reuse node freed by A since the allocation C fits in the same bin (using pow2 size to be sure) - let a = allocator.allocate(1024).unwrap(); - assert_eq!(a.offset, 0); - - let b = allocator.allocate(3456).unwrap(); - assert_eq!(b.offset, 1024); - - allocator.free(a); - - let c = allocator.allocate(1024).unwrap(); - assert_eq!(c.offset, 0); - - allocator.free(c); - allocator.free(b); - - // End: Validate that allocator has no fragmentation left. Should be 100% clean. - let validate_all = allocator.allocate(1024 * 1024 * 256).unwrap(); - assert_eq!(validate_all.offset, 0); - allocator.free(validate_all); -} - -#[test] -fn allocate_offset_allocator_reuse_complex() { - let mut allocator: Allocator = Allocator::new(1024 * 1024 * 256); - - // Allocator should not reuse node freed by A since the allocation C doesn't fits in the same bin - // However node D and E fit there and should reuse node from A - let a = allocator.allocate(1024).unwrap(); - assert_eq!(a.offset, 0); - - let b = allocator.allocate(3456).unwrap(); - assert_eq!(b.offset, 1024); - - allocator.free(a); - - let c = allocator.allocate(2345).unwrap(); - assert_eq!(c.offset, 1024 + 3456); - - let d = allocator.allocate(456).unwrap(); - assert_eq!(d.offset, 0); - - let e = allocator.allocate(512).unwrap(); - assert_eq!(e.offset, 456); - - let report = allocator.storage_report(); - assert_eq!( - report.total_free_space, - 1024 * 1024 * 256 - 3456 - 2345 - 456 - 512 - ); - assert_ne!(report.largest_free_region, report.total_free_space); - - allocator.free(c); - allocator.free(d); - allocator.free(b); - allocator.free(e); - - // End: Validate that allocator has no fragmentation left. Should be 100% clean. - let validate_all = allocator.allocate(1024 * 1024 * 256).unwrap(); - assert_eq!(validate_all.offset, 0); - allocator.free(validate_all); -} - -#[test] -fn allocate_offset_allocator_zero_fragmentation() { - let mut allocator: Allocator = Allocator::new(1024 * 1024 * 256); - - // Allocate 256x 1MB. Should fit. Then free four random slots and reallocate four slots. - // Plus free four contiguous slots an allocate 4x larger slot. All must be zero fragmentation! - let mut allocations: [_; 256] = array::from_fn(|i| { - let allocation = allocator.allocate(1024 * 1024).unwrap(); - assert_eq!(allocation.offset, i as u32 * 1024 * 1024); - allocation - }); - - let report = allocator.storage_report(); - assert_eq!(report.total_free_space, 0); - assert_eq!(report.largest_free_region, 0); - - // Free four random slots - allocator.free(allocations[243]); - allocator.free(allocations[5]); - allocator.free(allocations[123]); - allocator.free(allocations[95]); - - // Free four contiguous slots (allocator must merge) - allocator.free(allocations[151]); - allocator.free(allocations[152]); - allocator.free(allocations[153]); - allocator.free(allocations[154]); - - allocations[243] = allocator.allocate(1024 * 1024).unwrap(); - allocations[5] = allocator.allocate(1024 * 1024).unwrap(); - allocations[123] = allocator.allocate(1024 * 1024).unwrap(); - allocations[95] = allocator.allocate(1024 * 1024).unwrap(); - allocations[151] = allocator.allocate(1024 * 1024 * 4).unwrap(); // 4x larger - - for (i, allocation) in allocations.iter().enumerate() { - if !(152..155).contains(&i) { - allocator.free(*allocation); - } - } - - let report2 = allocator.storage_report(); - assert_eq!(report2.total_free_space, 1024 * 1024 * 256); - assert_eq!(report2.largest_free_region, 1024 * 1024 * 256); - - // End: Validate that allocator has no fragmentation left. Should be 100% clean. - let validate_all = allocator.allocate(1024 * 1024 * 256).unwrap(); - assert_eq!(validate_all.offset, 0); - allocator.free(validate_all); -} - -#[test] -fn ext_min_allocator_size() { - // Randomly generated integers on a log distribution, σ = 10. - static TEST_OBJECT_SIZES: [u32; 42] = [ - 0, 1, 2, 3, 4, 5, 8, 17, 23, 36, 51, 68, 87, 151, 165, 167, 201, 223, 306, 346, 394, 411, - 806, 969, 1404, 1798, 2236, 4281, 4745, 13989, 21095, 26594, 27146, 29679, 144685, 153878, - 495127, 727999, 1377073, 9440387, 41994490, 68520116, - ]; - - for needed_object_size in TEST_OBJECT_SIZES { - let allocator_size = ext::min_allocator_size(needed_object_size); - let mut allocator: Allocator = Allocator::new(allocator_size); - assert!(allocator.allocate(needed_object_size).is_some()); - } -}