core: span allocation with shared free list and geometric growth

Amp-Thread-ID: https://ampcode.com/threads/T-01a0b2dd-db32-712d-99ce-12ce4a86d979
Co-authored-by: Heaust Azure <heaust.azure@gmail.com>

core/render_data.rs

@@ -1,9 +1,63 @@
-use std::collections::{BTreeMap, HashMap, HashSet};
+use std::collections::{BTreeMap, HashMap};
 
 use serde::Serialize;
 
 const ALIGNMENT: u32 = 64;
 
+/// First-fit interval allocator over `u32` offsets; adjacent free ranges coalesce on release.
+#[derive(Default)]
+struct FreeList(BTreeMap<u32, u32>);
+
+impl FreeList {
+    fn allocate(&mut self, size: u32, align: u32) -> Option<u32> {
+        let (block, block_size, offset) = self.0.iter().find_map(|(&block, &block_size)| {
+            let offset = block.checked_add(align - 1)? & !(align - 1);
+            (offset.checked_add(size)? <= block + block_size).then_some((block, block_size, offset))
+        })?;
+        self.0.remove(&block);
+        if offset > block {
+            self.0.insert(block, offset - block);
+        }
+        let end = offset + size;
+        if end < block + block_size {
+            self.0.insert(end, block + block_size - end);
+        }
+        Some(offset)
+    }
+
+    /// Claims exactly `[offset, offset + size)` when that range is entirely free.
+    fn take(&mut self, offset: u32, size: u32) -> bool {
+        let Some(&available) = self.0.get(&offset) else {
+            return false;
+        };
+        if available < size {
+            return false;
+        }
+        self.0.remove(&offset);
+        if available > size {
+            self.0.insert(offset + size, available - size);
+        }
+        true
+    }
+
+    fn release(&mut self, mut offset: u32, mut size: u32) {
+        if let Some((&previous, &previous_size)) = self.0.range(..offset).next_back() {
+            if previous + previous_size == offset {
+                self.0.remove(&previous);
+                offset = previous;
+                size += previous_size;
+            }
+        }
+        if let Some((&next, &next_size)) = self.0.range(offset..).next() {
+            if offset + size == next {
+                self.0.remove(&next);
+                size += next_size;
+            }
+        }
+        self.0.insert(offset, size);
+    }
+}
+
 #[derive(Clone, Serialize)]
 pub struct Rows {
     pub name: String,
@@ -14,18 +68,18 @@ pub struct Rows {
     pub bytes: u32,
 }
 
+/// Row spans handed out by `allocate_object`, keyed by their first row.
 #[derive(Default)]
 struct Slots {
-    active: HashSet<u32>,
-    free: Vec<u32>,
-    next: u32,
+    free: FreeList,
+    spans: HashMap<u32, u32>,
 }
 
 pub struct RenderData {
     arena: Box<[u8]>,
     arena_start: usize,
     base: u32,
-    free: BTreeMap<u32, u32>,
+    free: FreeList,
     rows: HashMap<String, Rows>,
     slots: HashMap<String, Slots>,
 }
@@ -42,7 +96,7 @@ impl RenderData {
             arena,
             arena_start: aligned - pointer,
             base: aligned as u32,
-            free: BTreeMap::from([(0, capacity)]),
+            free: FreeList(BTreeMap::from([(0, capacity)])),
             rows: HashMap::new(),
             slots: HashMap::new(),
         };
@@ -76,7 +130,14 @@ impl RenderData {
             if name == "signals" {
                 return Err("ROWS_BUILTIN");
             }
-            return self.grow_rows(current, rows);
+            let added = rows - current.rows;
+            let descriptor = self.grow_rows(current, rows)?;
+            self.slots
+                .get_mut(&descriptor.name)
+                .unwrap()
+                .free
+                .release(descriptor.rows - added, added);
+            return Ok(descriptor);
         }
         let bytes = rows.checked_mul(stride).ok_or("ARENA_OOM")?;
         let offset = self.reserve(bytes)?;
@@ -91,7 +152,9 @@ impl RenderData {
             bytes,
         };
         self.rows.insert(name, descriptor.clone());
-        self.slots.insert(descriptor.name.clone(), Slots::default());
+        let mut slots = Slots::default();
+        slots.free.release(0, rows);
+        self.slots.insert(descriptor.name.clone(), slots);
         Ok(descriptor)
     }
 
@@ -102,7 +165,7 @@ impl RenderData {
         if self
             .slots
             .get(name)
-            .is_some_and(|slots| !slots.active.is_empty())
+            .is_some_and(|slots| !slots.spans.is_empty())
         {
             return Err("ROWS_ACTIVE");
         }
@@ -128,39 +191,44 @@ impl RenderData {
         Some(&self.arena[start..start + rows.bytes as usize])
     }
 
-    pub fn allocate_object(&mut self, name: &str) -> Result<(u32, bool), &'static str> {
+    /// Reserves `count` contiguous rows and returns the first one; grows geometrically when full.
+    pub fn allocate_object(&mut self, name: &str, count: u32) -> Result<(u32, bool), &'static str> {
         if name == "signals" {
             return Err("ROWS_BUILTIN");
         }
-        let slots = self.slots.get(name).ok_or("ROWS_UNKNOWN")?;
-        let reused = slots.free.last().copied();
-        let id = reused.unwrap_or(slots.next);
-        let next = id.checked_add(1).ok_or("OBJECT_LIMIT")?;
-        let capacity = self.rows[name].rows;
-        let grew = id >= capacity;
-        if grew {
-            let descriptor = self.rows[name].clone();
-            self.grow_rows(descriptor, next)?;
+        if count == 0 {
+            return Err("ROWS");
         }
-        let slots = self.slots.get_mut(name).unwrap();
-        if reused.is_some() {
-            slots.free.pop();
-        } else {
-            slots.next = next;
-        }
-        slots.active.insert(id);
+        let slots = self.slots.get_mut(name).ok_or("ROWS_UNKNOWN")?;
+        let mut grew = false;
+        let id = match slots.free.allocate(count, 1) {
+            Some(id) => id,
+            None => {
+                let descriptor = self.rows[name].clone();
+                let rows = descriptor
+                    .rows
+                    .checked_mul(2)
+                    .ok_or("OBJECT_LIMIT")?
+                    .max(descriptor.rows.checked_add(count).ok_or("OBJECT_LIMIT")?);
+                let old_rows = descriptor.rows;
+                self.grow_rows(descriptor, rows)?;
+                grew = true;
+                let slots = self.slots.get_mut(name).unwrap();
+                slots.free.release(old_rows, rows - old_rows);
+                slots.free.allocate(count, 1).unwrap()
+            }
+        };
+        self.slots.get_mut(name).unwrap().spans.insert(id, count);
         Ok((id, grew))
     }
 
     pub fn delete_object(&mut self, name: &str, id: u32) -> Result<(), &'static str> {
         let slots = self.slots.get_mut(name).ok_or("ROWS_UNKNOWN")?;
-        if !slots.active.remove(&id) {
-            return Err("OBJECT_UNKNOWN");
-        }
+        let count = slots.spans.remove(&id).ok_or("OBJECT_UNKNOWN")?;
+        slots.free.release(id, count);
         let rows = self.rows[name].clone();
         let start = self.arena_start + (rows.offset - self.base + id * rows.stride) as usize;
-        self.arena[start..start + rows.stride as usize].fill(0);
-        slots.free.push(id);
+        self.arena[start..start + (count * rows.stride) as usize].fill(0);
         Ok(())
     }
 
@@ -169,15 +237,7 @@ impl RenderData {
         let old_offset = descriptor.offset - self.base;
         let old_end = old_offset + descriptor.bytes;
         let extra = bytes - descriptor.bytes;
-        if self
-            .free
-            .get(&old_end)
-            .is_some_and(|available| *available >= extra)
-        {
-            let available = self.free.remove(&old_end).unwrap();
-            if available > extra {
-                self.free.insert(old_end + extra, available - extra);
-            }
+        if self.free.take(old_end, extra) {
             self.arena[self.arena_start + old_end as usize
                 ..self.arena_start + (old_end + extra) as usize]
                 .fill(0);
@@ -244,54 +304,14 @@ impl RenderData {
     }
 
     fn reserve(&mut self, bytes: u32) -> Result<u32, &'static str> {
-        let (block, block_bytes, offset) = self
-            .free
-            .iter()
-            .find_map(|(&block, &block_bytes)| {
-                let offset = align(block)?;
-                (offset.checked_add(bytes)? <= block.checked_add(block_bytes)?).then_some((
-                    block,
-                    block_bytes,
-                    offset,
-                ))
-            })
-            .ok_or("ARENA_OOM")?;
-        self.free.remove(&block);
-        if offset > block {
-            self.free.insert(block, offset - block);
-        }
-        let end = offset + bytes;
-        let block_end = block + block_bytes;
-        if end < block_end {
-            self.free.insert(end, block_end - end);
-        }
-        Ok(offset)
+        self.free.allocate(bytes, ALIGNMENT).ok_or("ARENA_OOM")
     }
 
-    fn release(&mut self, mut offset: u32, mut bytes: u32) {
-        if let Some((&previous, &previous_bytes)) = self.free.range(..offset).next_back() {
-            if previous + previous_bytes == offset {
-                self.free.remove(&previous);
-                offset = previous;
-                bytes += previous_bytes;
-            }
-        }
-        if let Some((&next, &next_bytes)) = self.free.range(offset..).next() {
-            if offset + bytes == next {
-                self.free.remove(&next);
-                bytes += next_bytes;
-            }
-        }
-        self.free.insert(offset, bytes);
+    fn release(&mut self, offset: u32, bytes: u32) {
+        self.free.release(offset, bytes);
     }
 }
 
-fn align(value: u32) -> Option<u32> {
-    value
-        .checked_add(ALIGNMENT - 1)
-        .map(|value| value & !(ALIGNMENT - 1))
-}
-
 #[cfg(test)]
 #[path = "render_data_tests.rs"]
 mod tests;