yawn/core/render_data.rs

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((&block, &block_size)) = self.0.range(..=offset).next_back() else {
            return false;
        };
        if offset + size > block + block_size {
            return false;
        }
        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);
        }
        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, Debug, PartialEq, Serialize)]
pub struct Rows {
    pub name: String,
    pub rows: u32,
    pub stride: u32,
    pub format: String,
    pub offset: u32,
    pub bytes: u32,
}

/// Row spans handed out by `allocate_object`, keyed by their first row.
#[derive(Default)]
struct Slots {
    free: FreeList,
    spans: HashMap<u32, u32>,
}

pub struct RenderData {
    arena: Box<[u8]>,
    arena_start: usize,
    base: u32,
    free: FreeList,
    rows: HashMap<String, Rows>,
    slots: HashMap<String, Slots>,
    /// Relocated descriptors the writer has not committed yet; `rows` keeps the old storage live.
    pending: HashMap<String, Rows>,
}

impl RenderData {
    pub fn new(capacity: u32) -> Result<Self, &'static str> {
        if capacity < ALIGNMENT || capacity > u32::MAX - (ALIGNMENT - 1) {
            return Err("INIT");
        }
        let mut arena = vec![0; (capacity + ALIGNMENT - 1) as usize].into_boxed_slice();
        let pointer = arena.as_mut_ptr() as usize;
        let aligned = (pointer + ALIGNMENT as usize - 1) & !(ALIGNMENT as usize - 1);
        let mut data = Self {
            arena,
            arena_start: aligned - pointer,
            base: aligned as u32,
            free: FreeList(BTreeMap::from([(0, capacity)])),
            rows: HashMap::new(),
            slots: HashMap::new(),
            pending: HashMap::new(),
        };
        data.create_rows("signals".into(), 1, 32, "f32".into())?;
        data.write_signal(5, 1.0);
        Ok(data)
    }

    /// Creates or grows an array. The flag is set when the rows moved: the returned descriptor
    /// is then staged and the writer must copy the old bytes into it and `commit` it.
    pub fn create_rows(
        &mut self,
        name: String,
        rows: u32,
        stride: u32,
        format: String,
    ) -> Result<(Rows, bool), &'static str> {
        if name.is_empty()
            || rows == 0
            || stride < 16
            || stride % 16 != 0
            || !matches!(format.as_str(), "f32" | "u32" | "i32")
        {
            return Err("ROWS");
        }
        if let Some(current) = self.rows.get(&name).cloned() {
            if current.stride != stride || current.format != format {
                return Err("ROWS");
            }
            if rows <= current.rows {
                return Ok((current, false));
            }
            if name == "signals" {
                return Err("ROWS_BUILTIN");
            }
            let added = rows - current.rows;
            let (descriptor, moved) = self.grow_rows(current, rows)?;
            self.slots
                .get_mut(&descriptor.name)
                .unwrap()
                .free
                .release(descriptor.rows - added, added);
            return Ok((descriptor, moved));
        }
        let bytes = rows.checked_mul(stride).ok_or("ARENA_OOM")?;
        let offset = self.reserve(bytes)?;
        let end = offset + bytes;
        self.arena[self.arena_start + offset as usize..self.arena_start + end as usize].fill(0);
        let descriptor = Rows {
            name: name.clone(),
            rows,
            stride,
            format,
            offset: self.base + offset,
            bytes,
        };
        self.rows.insert(name, descriptor.clone());
        let mut slots = Slots::default();
        slots.free.release(0, rows);
        self.slots.insert(descriptor.name.clone(), slots);
        Ok((descriptor, false))
    }

    pub fn delete_rows(&mut self, name: &str) -> Result<(), &'static str> {
        if name == "signals" {
            return Err("ROWS_BUILTIN");
        }
        if self.pending.contains_key(name) {
            return Err("ROWS_PENDING");
        }
        if self
            .slots
            .get(name)
            .is_some_and(|slots| !slots.spans.is_empty())
        {
            return Err("ROWS_ACTIVE");
        }
        let rows = self.rows.remove(name).ok_or("ROWS_UNKNOWN")?;
        self.slots.remove(name);
        self.release(rows.offset - self.base, rows.bytes);
        Ok(())
    }

    pub fn rows(&self, name: &str) -> Option<&Rows> {
        self.rows.get(name)
    }

    pub fn descriptors(&self) -> Vec<Rows> {
        let mut rows = self.rows.values().cloned().collect::<Vec<_>>();
        rows.sort_by_key(|rows| rows.offset);
        rows
    }

    pub fn bytes(&self, name: &str) -> Option<&[u8]> {
        let rows = self.rows(name)?;
        let start = self.arena_start + (rows.offset - self.base) as usize;
        Some(&self.arena[start..start + rows.bytes as usize])
    }

    /// Reserves `count` contiguous rows and returns the first one; grows geometrically when full.
    /// The growth flags mirror `create_rows`: `(grew, moved)`.
    pub fn allocate_object(
        &mut self,
        name: &str,
        count: u32,
    ) -> Result<(u32, bool, bool), &'static str> {
        if name == "signals" {
            return Err("ROWS_BUILTIN");
        }
        if count == 0 {
            return Err("ROWS");
        }
        let slots = self.slots.get_mut(name).ok_or("ROWS_UNKNOWN")?;
        let mut grew = false;
        let mut moved = 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;
                moved = self.grow_rows(descriptor, rows)?.1;
                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, moved))
    }

    /// The descriptor the writer should use: staged rows win over committed ones.
    pub fn staged(&self, name: &str) -> Option<&Rows> {
        self.pending.get(name).or_else(|| self.rows.get(name))
    }

    /// Publishes staged relocations after the writer copied the old bytes into them.
    pub fn commit(&mut self, names: &[String]) -> Result<(), &'static str> {
        for name in names {
            let staged = self.pending.remove(name).ok_or("ROWS_PENDING")?;
            let old = self.rows.insert(name.clone(), staged).unwrap();
            self.release(old.offset - self.base, old.bytes);
        }
        if !self.pending.is_empty() {
            return Err("ROWS_PENDING");
        }
        Ok(())
    }

    /// Drops every staged relocation, for example when a batch fails half-way.
    pub fn discard_pending(&mut self) {
        for staged in std::mem::take(&mut self.pending).into_values() {
            self.release(staged.offset - self.base, staged.bytes);
            let rows = self.rows[&staged.name].rows;
            let slots = self.slots.get_mut(&staged.name).unwrap();
            assert!(slots.free.take(rows, staged.rows - rows));
        }
    }

    pub fn delete_object(&mut self, name: &str, id: u32) -> Result<(), &'static str> {
        let slots = self.slots.get_mut(name).ok_or("ROWS_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 + (count * rows.stride) as usize].fill(0);
        Ok(())
    }

    /// Extends in place when the next arena block is free; otherwise stages a zeroed region for
    /// the writer to fill and `commit`, because only the writer knows when its last write landed.
    fn grow_rows(&mut self, mut descriptor: Rows, rows: u32) -> Result<(Rows, bool), &'static str> {
        if self.pending.contains_key(&descriptor.name) {
            return Err("ROWS_PENDING");
        }
        let bytes = rows.checked_mul(descriptor.stride).ok_or("ARENA_OOM")?;
        let old_offset = descriptor.offset - self.base;
        let old_end = old_offset + descriptor.bytes;
        let extra = bytes - descriptor.bytes;
        let moved = !self.free.take(old_end, extra);
        if moved {
            let offset = self.reserve(bytes)?;
            let start = self.arena_start + offset as usize;
            self.arena[start..start + bytes as usize].fill(0);
            descriptor.offset = self.base + offset;
        } else {
            self.arena[self.arena_start + old_end as usize
                ..self.arena_start + (old_end + extra) as usize]
                .fill(0);
        }
        descriptor.rows = rows;
        descriptor.bytes = bytes;
        if moved {
            &mut self.pending
        } else {
            &mut self.rows
        }
        .insert(descriptor.name.clone(), descriptor.clone());
        Ok((descriptor, moved))
    }

    pub fn update_signals(&mut self, delta: f32, frame: u32, elapsed: f32, fps: u32) {
        let start = self.arena_start + (self.rows["signals"].offset - self.base) as usize;
        for (index, value) in [delta, frame as f32, elapsed, fps as f32]
            .iter()
            .enumerate()
        {
            self.arena[start + index * 4..start + index * 4 + 4]
                .copy_from_slice(&value.to_le_bytes());
        }
    }

    pub fn render_pending(&self) -> bool {
        self.signal(4) == 0.0 && self.signal(5) != 0.0 && self.signal(6) == 0.0
    }

    pub fn begin_render(&mut self) -> bool {
        if !self.render_pending() {
            return false;
        }
        self.write_signal(5, 0.0);
        true
    }

    pub fn mark_dirty(&mut self) {
        self.write_signal(5, 1.0);
    }

    pub fn loadout_ready(&mut self) {
        self.write_signal(6, 0.0);
        self.mark_dirty();
    }

    fn signal(&self, index: usize) -> f32 {
        let rows = &self.rows["signals"];
        let start = self.arena_start + (rows.offset - self.base) as usize + index * 4;
        f32::from_le_bytes(self.arena[start..start + 4].try_into().unwrap())
    }

    fn write_signal(&mut self, index: usize, value: f32) {
        let rows = &self.rows["signals"];
        let start = self.arena_start + (rows.offset - self.base) as usize + index * 4;
        self.arena[start..start + 4].copy_from_slice(&value.to_le_bytes());
    }

    fn reserve(&mut self, bytes: u32) -> Result<u32, &'static str> {
        self.free.allocate(bytes, ALIGNMENT).ok_or("ARENA_OOM")
    }

    fn release(&mut self, offset: u32, bytes: u32) {
        self.free.release(offset, bytes);
    }
}

#[cfg(test)]
#[path = "render_data_tests.rs"]
mod tests;