#[cfg(not(feature = "std"))]
use alloc::{vec, vec::Vec};
use crate::file_writer::{write_offset, write_undef_offset};
pub(crate) const TABLE_WIDTH: u16 = 4;
pub(crate) const STARTING_BLOCK_SIZE: u64 = 1024;
pub(crate) const MAX_DIRECT_BLOCK_SIZE: u64 = 65_536;
pub(crate) const MAX_HEAP_SIZE_BITS: u16 = 40;
pub(crate) const BLOCK_OFFSET_BYTES: usize = (MAX_HEAP_SIZE_BITS as usize).div_ceil(8);
pub(crate) const START_ROOT_ROWS: u16 = 1;
const WIDTH: u64 = TABLE_WIDTH as u64;
const START_BITS: u32 = 10;
const WIDTH_BITS: u32 = 2;
const MAX_DIRECT_BITS: u32 = 16;
const _: () = assert!(1u64 << START_BITS == STARTING_BLOCK_SIZE);
const _: () = assert!(1u64 << WIDTH_BITS == WIDTH);
const _: () = assert!(1u64 << MAX_DIRECT_BITS == MAX_DIRECT_BLOCK_SIZE);
const FIRST_ROW_BITS: u32 = START_BITS + WIDTH_BITS;
const MAX_DIRECT_ROWS: usize = (MAX_DIRECT_BITS - START_BITS + 2) as usize;
const MAX_ROOT_ROWS: usize = (MAX_HEAP_SIZE_BITS as u32 - FIRST_ROW_BITS + 1) as usize;
pub(crate) const MAX_HEAP_SPACE: u64 = 1u64 << MAX_HEAP_SIZE_BITS;
pub(crate) const fn direct_block_header(offset_size: u8) -> usize {
4 + 1 + offset_size as usize + BLOCK_OFFSET_BYTES + 4
}
pub(crate) const fn max_managed_object(offset_size: u8) -> usize {
(1usize << MAX_DIRECT_BITS) - direct_block_header(offset_size)
}
const fn indirect_block_size(nrows: u16, offset_size: u8) -> u64 {
4 + 1
+ offset_size as u64
+ BLOCK_OFFSET_BYTES as u64
+ nrows as u64 * WIDTH * offset_size as u64
+ 4
}
fn row_block_size(row: usize) -> u64 {
if row <= 1 {
STARTING_BLOCK_SIZE
} else {
STARTING_BLOCK_SIZE << (row - 1)
}
}
fn row_offset(row: usize) -> u64 {
if row == 0 {
0
} else {
(STARTING_BLOCK_SIZE * WIDTH) << (row - 1)
}
}
fn lookup(offset: u64) -> (usize, u64) {
if offset < STARTING_BLOCK_SIZE * WIDTH {
return (0, offset / STARTING_BLOCK_SIZE);
}
let high_bit = 63 - offset.leading_zeros();
let row = (high_bit - FIRST_ROW_BITS + 1) as usize;
(row, (offset - (1u64 << high_bit)) / row_block_size(row))
}
fn size_to_rows(size: u64) -> usize {
((63 - size.leading_zeros()) - FIRST_ROW_BITS + 1) as usize
}
fn locate(offset: u64) -> (u64, u64) {
let mut base = 0;
let mut local = offset;
loop {
let (row, col) = lookup(local);
let block_size = row_block_size(row);
let within = row_offset(row) + col * block_size;
if row < MAX_DIRECT_ROWS {
return (base + within, block_size);
}
base += within;
local -= within;
}
}
fn rows_capacity(nrows: usize, offset_size: u8) -> u64 {
let mut per_block = [0u64; MAX_ROOT_ROWS];
let mut total = 0;
for row in 0..nrows {
per_block[row] = if row < MAX_DIRECT_ROWS {
row_block_size(row) - direct_block_header(offset_size) as u64
} else {
let child_rows = size_to_rows(row_block_size(row));
WIDTH * per_block[..child_rows].iter().sum::<u64>()
};
total += WIDTH * per_block[row];
}
total
}
fn root_rows_covering(span: u64) -> Option<usize> {
(1..=MAX_ROOT_ROWS).find(|&n| row_offset(n) >= span)
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum PlanRefusal {
HeapSpace,
Host {
bytes: u64,
},
}
#[derive(Clone, Copy)]
enum Child {
Direct(usize),
Indirect(usize),
}
struct PlannedDirect {
heap_offset: u64,
size: u64,
region_offset: u64,
objects: Vec<usize>,
}
struct PlannedIndirect {
heap_offset: u64,
nrows: u16,
region_offset: u64,
entries: Vec<Option<Child>>,
}
pub(crate) struct ManagedPlan {
offset_size: u8,
offsets: Vec<u64>,
directs: Vec<PlannedDirect>,
indirects: Vec<PlannedIndirect>,
region_size: u64,
managed_space: u64,
allocated_space: u64,
allocation_iterator: u64,
free_space: u64,
}
impl ManagedPlan {
pub(crate) fn new(sizes: &[u64], offset_size: u8) -> Result<ManagedPlan, PlanRefusal> {
debug_assert!(
sizes
.iter()
.all(|&size| size <= max_managed_object(offset_size) as u64),
"an object too large for a managed block belongs in huge storage"
);
let header = direct_block_header(offset_size) as u64;
let mut directs: Vec<PlannedDirect> = Vec::new();
let mut offsets: Vec<u64> = Vec::with_capacity(sizes.len());
let mut cursor = 0;
let mut fill = 0;
for &size in sizes {
loop {
if let Some(block) = directs.last_mut() {
if block.size - fill >= size {
block.objects.push(offsets.len());
offsets.push(block.heap_offset + fill);
fill += size;
break;
}
}
if cursor >= MAX_HEAP_SPACE {
return Err(PlanRefusal::HeapSpace);
}
let (heap_offset, block_size) = locate(cursor);
debug_assert_eq!(
heap_offset, cursor,
"the walk visits whole blocks, in order"
);
cursor = heap_offset + block_size;
if block_size - header >= size {
directs.push(PlannedDirect {
heap_offset,
size: block_size,
region_offset: 0,
objects: Vec::new(),
});
fill = header;
}
}
}
if directs.is_empty() {
directs.push(PlannedDirect {
heap_offset: 0,
size: STARTING_BLOCK_SIZE,
region_offset: 0,
objects: Vec::new(),
});
cursor = STARTING_BLOCK_SIZE;
}
let root_is_direct = directs.len() == 1
&& directs[0].heap_offset == 0
&& directs[0].size == STARTING_BLOCK_SIZE;
let mut indirects = Vec::new();
let (managed_space, capacity) = if root_is_direct {
(STARTING_BLOCK_SIZE, STARTING_BLOCK_SIZE - header)
} else {
let nrows = root_rows_covering(cursor).ok_or(PlanRefusal::HeapSpace)?;
let mut placed = 0;
build_indirect(0, nrows, &directs, &mut placed, &mut indirects);
debug_assert_eq!(placed, directs.len(), "every block belongs to a slot");
(row_offset(nrows), rows_capacity(nrows, offset_size))
};
let mut region_size = 0;
for block in &mut indirects {
block.region_offset = region_size;
region_size += indirect_block_size(block.nrows, offset_size);
}
for block in &mut directs {
block.region_offset = region_size;
region_size += block.size;
}
if usize::try_from(region_size).is_err() {
return Err(PlanRefusal::Host { bytes: region_size });
}
let used: u64 = sizes.iter().sum();
debug_assert!(
used <= capacity,
"objects cannot exceed the blocks holding them"
);
Ok(ManagedPlan {
offset_size,
offsets,
allocated_space: directs.iter().map(|b| b.size).sum(),
directs,
indirects,
region_size,
managed_space,
allocation_iterator: if root_is_direct { 0 } else { cursor },
free_space: capacity - used,
})
}
pub(crate) fn region_size(&self) -> u64 {
self.region_size
}
pub(crate) fn root_address(&self, region_address: u64) -> u64 {
match self.indirects.last() {
Some(root) => region_address + root.region_offset,
None => region_address + self.directs[0].region_offset,
}
}
pub(crate) fn root_rows(&self) -> u16 {
self.indirects.last().map_or(0, |root| root.nrows)
}
pub(crate) fn managed_space(&self) -> u64 {
self.managed_space
}
pub(crate) fn allocated_space(&self) -> u64 {
self.allocated_space
}
pub(crate) fn allocation_iterator(&self) -> u64 {
self.allocation_iterator
}
pub(crate) fn free_space(&self) -> u64 {
self.free_space
}
pub(crate) fn heap_offset(&self, index: usize) -> u64 {
self.offsets[index]
}
pub(crate) fn serialize(
&self,
objects: &[&[u8]],
region_address: u64,
heap_header_address: u64,
) -> Vec<u8> {
let region_size = usize::try_from(self.region_size)
.expect("ManagedPlan::new refuses a region this host cannot address");
let mut region = vec![0u8; region_size];
for block in &self.indirects {
let mut bytes = Vec::new();
bytes.extend_from_slice(b"FHIB");
bytes.push(0); write_offset(&mut bytes, heap_header_address, self.offset_size);
write_heap_offset(&mut bytes, block.heap_offset);
for entry in block.entries.iter().copied() {
match entry {
Some(Child::Direct(at)) => {
let address = region_address + self.directs[at].region_offset;
write_offset(&mut bytes, address, self.offset_size);
}
Some(Child::Indirect(at)) => {
let address = region_address + self.indirects[at].region_offset;
write_offset(&mut bytes, address, self.offset_size);
}
None => write_undef_offset(&mut bytes, self.offset_size),
}
}
let checksum = crate::checksum::jenkins_lookup3(&bytes);
bytes.extend_from_slice(&checksum.to_le_bytes());
debug_assert_eq!(
bytes.len() as u64,
indirect_block_size(block.nrows, self.offset_size)
);
place(&mut region, block.region_offset, &bytes);
}
for block in &self.directs {
let size = usize::try_from(block.size).expect("a direct block is at most 64 KiB");
let mut bytes = Vec::with_capacity(size);
bytes.extend_from_slice(b"FHDB");
bytes.push(0); write_offset(&mut bytes, heap_header_address, self.offset_size);
write_heap_offset(&mut bytes, block.heap_offset);
let checksum_at = bytes.len();
bytes.extend_from_slice(&[0u8; 4]); debug_assert_eq!(bytes.len(), direct_block_header(self.offset_size));
for &object in &block.objects {
debug_assert_eq!(
block.heap_offset + bytes.len() as u64,
self.offsets[object],
"an object must be emitted at the offset it was planned at"
);
bytes.extend_from_slice(objects[object]);
}
bytes.resize(size, 0);
let checksum = crate::checksum::jenkins_lookup3(&bytes);
bytes[checksum_at..checksum_at + 4].copy_from_slice(&checksum.to_le_bytes());
place(&mut region, block.region_offset, &bytes);
}
region
}
}
fn place(region: &mut [u8], at: u64, bytes: &[u8]) {
let at = usize::try_from(at).expect("a region offset is bounded by the region size");
region[at..at + bytes.len()].copy_from_slice(bytes);
}
fn write_heap_offset(buf: &mut Vec<u8>, offset: u64) {
debug_assert!(offset < MAX_HEAP_SPACE, "heap offset overflows its field");
buf.extend_from_slice(&offset.to_le_bytes()[..BLOCK_OFFSET_BYTES]);
}
fn build_indirect(
base: u64,
nrows: usize,
directs: &[PlannedDirect],
placed: &mut usize,
indirects: &mut Vec<PlannedIndirect>,
) -> usize {
let mut entries = vec![None; nrows * TABLE_WIDTH as usize];
'rows: for row in 0..nrows {
let block_size = row_block_size(row);
for col in 0..TABLE_WIDTH as usize {
let Some(next) = directs.get(*placed) else {
break 'rows;
};
let slot_offset = base + row_offset(row) + col as u64 * block_size;
if next.heap_offset >= slot_offset + block_size {
continue;
}
entries[row * TABLE_WIDTH as usize + col] = Some(if row < MAX_DIRECT_ROWS {
debug_assert_eq!(
next.heap_offset, slot_offset,
"a block fills its whole slot"
);
*placed += 1;
Child::Direct(*placed - 1)
} else {
Child::Indirect(build_indirect(
slot_offset,
size_to_rows(block_size),
directs,
placed,
indirects,
))
});
}
}
indirects.push(PlannedIndirect {
heap_offset: base,
nrows: u16::try_from(nrows).expect("a row count is bounded by MAX_ROOT_ROWS"),
region_offset: 0,
entries,
});
indirects.len() - 1
}
#[cfg(test)]
mod tests {
use super::*;
const OFFSET_SIZE: u8 = 8;
fn largest_object() -> u64 {
MAX_DIRECT_BLOCK_SIZE - direct_block_header(OFFSET_SIZE) as u64
}
fn shapes() -> Vec<(&'static str, Vec<u64>)> {
let starting_capacity = STARTING_BLOCK_SIZE - direct_block_header(OFFSET_SIZE) as u64;
vec![
("empty", Vec::new()),
("one tiny object", vec![1]),
("one starting block, exactly full", vec![starting_capacity]),
("one byte past a starting block", vec![starting_capacity, 1]),
(
"many small objects",
(0..500).map(|i| 20 + i % 40).collect(),
),
(
"objects that skip the small rows",
vec![largest_object(); 3],
),
(
"enough large objects to nest indirect blocks",
vec![largest_object(); 40],
),
(
"large and small mixed",
(0..60)
.map(|i| if i % 3 == 0 { largest_object() } else { 100 })
.collect(),
),
]
}
#[test]
fn objects_sit_inside_the_blocks_planned_for_them() {
for (name, sizes) in shapes() {
let plan = ManagedPlan::new(&sizes, OFFSET_SIZE).expect("plannable");
let header = direct_block_header(OFFSET_SIZE) as u64;
let mut previous_end = 0;
for (index, &size) in sizes.iter().enumerate() {
let at = plan.heap_offset(index);
let block = plan
.directs
.iter()
.find(|b| at >= b.heap_offset && at < b.heap_offset + b.size)
.unwrap_or_else(|| panic!("{name}: object {index} is in no allocated block"));
assert!(
at >= block.heap_offset + header,
"{name}: object {index} overlaps its block's header"
);
assert!(
at + size <= block.heap_offset + block.size,
"{name}: object {index} runs past its block"
);
assert!(
at >= previous_end,
"{name}: object {index} overlaps its predecessor"
);
previous_end = at + size;
}
}
}
#[test]
fn blocks_are_doubling_table_slots_in_heap_order() {
for (name, sizes) in shapes() {
let plan = ManagedPlan::new(&sizes, OFFSET_SIZE).expect("plannable");
let mut previous_end = 0;
for block in &plan.directs {
assert_eq!(
locate(block.heap_offset),
(block.heap_offset, block.size),
"{name}: a block at {} is not a slot of that size",
block.heap_offset
);
assert!(block.heap_offset >= previous_end, "{name}: blocks overlap");
previous_end = block.heap_offset + block.size;
}
assert!(
plan.managed_space() >= previous_end,
"{name}: the declared managed space does not cover the blocks"
);
assert!(
plan.managed_space() <= MAX_HEAP_SPACE,
"{name}: the heap outgrew its own address space"
);
}
}
#[test]
fn no_indirect_block_is_childless() {
for (name, sizes) in shapes() {
let plan = ManagedPlan::new(&sizes, OFFSET_SIZE).expect("plannable");
for block in &plan.indirects {
assert!(
block.entries.iter().any(Option::is_some),
"{name}: an indirect block at {} has no children",
block.heap_offset
);
}
}
}
#[test]
fn the_tree_puts_every_block_at_the_slot_its_heap_offset_names() {
for (name, sizes) in shapes() {
let plan = ManagedPlan::new(&sizes, OFFSET_SIZE).expect("plannable");
let Some(root) = plan.indirects.len().checked_sub(1) else {
assert_eq!(plan.directs.len(), 1, "{name}: a direct root is one block");
assert_eq!(plan.directs[0].heap_offset, 0);
continue;
};
let mut reached = Vec::new();
walk(&plan, root, 0, &mut reached, name);
let planned: Vec<u64> = plan.directs.iter().map(|b| b.heap_offset).collect();
assert_eq!(
reached, planned,
"{name}: the walk missed or reordered blocks"
);
}
}
fn walk(plan: &ManagedPlan, index: usize, expected: u64, reached: &mut Vec<u64>, name: &str) {
let block = &plan.indirects[index];
assert_eq!(
block.heap_offset, expected,
"{name}: block at the wrong slot"
);
assert_eq!(
block.entries.len(),
block.nrows as usize * TABLE_WIDTH as usize
);
for (slot, entry) in block.entries.iter().enumerate() {
let row = slot / TABLE_WIDTH as usize;
let col = (slot % TABLE_WIDTH as usize) as u64;
let size = row_block_size(row);
let at = block.heap_offset + row_offset(row) + col * size;
match entry {
None => {}
Some(Child::Direct(child)) => {
assert!(
row < MAX_DIRECT_ROWS,
"{name}: a direct block in an indirect row"
);
let child = &plan.directs[*child];
assert_eq!((child.heap_offset, child.size), (at, size), "{name}");
reached.push(child.heap_offset);
}
Some(Child::Indirect(child)) => {
assert!(
row >= MAX_DIRECT_ROWS,
"{name}: an indirect block in a direct row"
);
walk(plan, *child, at, reached, name);
}
}
}
}
#[test]
fn the_header_statistics_describe_the_blocks_that_were_planned() {
let header = direct_block_header(OFFSET_SIZE) as u64;
for (name, sizes) in shapes() {
let plan = ManagedPlan::new(&sizes, OFFSET_SIZE).expect("plannable");
let mut at = 0;
let mut capacity = 0;
while at < plan.managed_space() {
let (start, size) = locate(at);
assert_eq!(start, at, "{name}: a slot does not begin where it is found");
capacity += size - header;
at = start + size;
}
assert_eq!(
at,
plan.managed_space(),
"{name}: the managed space is not a whole number of blocks"
);
let used: u64 = sizes.iter().sum();
assert_eq!(
plan.free_space(),
capacity - used,
"{name}: free space must count every block the rows describe, \
allocated or not, less the object bytes"
);
let indirect_bytes: u64 = plan
.indirects
.iter()
.map(|b| indirect_block_size(b.nrows, OFFSET_SIZE))
.sum();
assert_eq!(
plan.allocated_space() + indirect_bytes,
plan.region_size(),
"{name}: allocated space is not the direct blocks alone"
);
let last = plan.directs.last().expect("a heap has at least one block");
let expected = if plan.root_rows() == 0 {
0
} else {
last.heap_offset + last.size
};
assert_eq!(
plan.allocation_iterator(),
expected,
"{name}: the iterator must sit past the last allocated block, \
or at zero while the root is a bare direct block"
);
}
}
#[test]
fn the_root_runs_out_of_rows_exactly_at_the_heaps_address_space() {
assert_eq!(row_offset(MAX_ROOT_ROWS), MAX_HEAP_SPACE);
assert_eq!(root_rows_covering(MAX_HEAP_SPACE), Some(MAX_ROOT_ROWS));
assert_eq!(root_rows_covering(MAX_HEAP_SPACE - 1), Some(MAX_ROOT_ROWS));
assert_eq!(root_rows_covering(MAX_HEAP_SPACE + 1), None);
assert_eq!(root_rows_covering(row_offset(3)), Some(3));
assert_eq!(root_rows_covering(row_offset(3) + 1), Some(4));
}
#[test]
fn the_root_is_direct_only_while_one_starting_block_holds_everything() {
let capacity = STARTING_BLOCK_SIZE - direct_block_header(OFFSET_SIZE) as u64;
for (sizes, direct) in [
(vec![], true),
(vec![capacity], true),
(vec![capacity, 1], false),
(vec![capacity - 1, 1], true),
(vec![capacity + 1], false),
] {
let plan = ManagedPlan::new(&sizes, OFFSET_SIZE).expect("plannable");
assert_eq!(
plan.root_rows() == 0,
direct,
"{sizes:?} should{} have a direct root",
if direct { "" } else { " not" }
);
if direct {
assert_eq!(plan.managed_space(), STARTING_BLOCK_SIZE);
}
}
}
}