myriad-rs 0.2.0

use polka::{HANDLE_NONE, HANDLE_SLOT_MAX};
use alloc::{boxed::Box, format, string::String, vec, vec::Vec};
use hashbrown::HashMap;

// Route a trace line to the installed text sink, if any. No-op when unset
// (the no_std default), so diagnostics cost nothing in production builds.
macro_rules! htrace {
    ($self:ident, $($arg:tt)*) => {
        if let Some(f) = $self.trace_out { f(&format!($($arg)*)); }
    };
}

#[inline(always)]
pub fn handle_parts(raw: u64) -> (u32, u32) {
    (((raw >> 24) & 0x00FF_FFFF) as u32, (raw & 0x00FF_FFFF) as u32)
}

#[inline(always)]
pub fn make_handle(slot: u32, generation: u32) -> u64 {
    ((slot as u64) << 24) | (generation as u64 & 0x00FF_FFFF)
}

pub const HEAP_BYTES_PER_SLOT: usize = 8 + 1; // data u64 + amortized mask bit

pub struct Cell {
    pub data: Box<[u64]>,
    // Bit i (LSB of word i/64) = 1 → data[i] is a handle.
    pub mask: Box<[u64]>,
}

pub struct Heap {
    cells: Vec<Option<Cell>>,
    rc: Vec<u32>,
    generation: Vec<u32>,
    free_list: Vec<u32>,
    bytes_used: usize,
    trace_slot: Option<u32>,
    trace_all: bool,
    trace_out: Option<fn(&str)>,
    pub trace_pc: usize,
    buffer_pool: HashMap<usize, Vec<Cell>>,
}

#[inline]
fn mask_words_for(size: usize) -> usize { (size + 63) / 64 }

#[inline]
pub fn mask_bit(mask: &[u64], i: usize) -> bool {
    let w = i / 64;
    let b = i % 64;
    mask.get(w).map_or(false, |x| (x >> b) & 1 == 1)
}

#[inline]
pub fn mask_set(mask: &mut [u64], i: usize, on: bool) {
    let w = i / 64;
    let b = i % 64;
    if let Some(x) = mask.get_mut(w) {
        if on { *x |= 1u64 << b; } else { *x &= !(1u64 << b); }
    }
}

impl Heap {
    pub fn new() -> Self {
        Self {
            cells: Vec::new(),
            rc: Vec::new(),
            generation: Vec::new(),
            free_list: Vec::new(),
            bytes_used: 0,
            trace_slot: None,
            trace_all: false,
            trace_out: None,
            trace_pc: 0,
            buffer_pool: HashMap::new(),
        }
    }

    pub fn set_trace(&mut self, slot: Option<u32>, all: bool, out: fn(&str)) {
        self.trace_slot = slot;
        self.trace_all = all;
        self.trace_out = Some(out);
    }

    pub fn bytes_used(&self) -> usize { self.bytes_used }

    #[inline]
    fn traced(&self, slot: u32) -> bool { self.trace_all || self.trace_slot == Some(slot) }

    fn make_cell(size: usize, init_mask: &[u64]) -> Cell {
        let data: Box<[u64]> = vec![HANDLE_NONE; size].into_boxed_slice();
        let nwords = mask_words_for(size);
        let mut mask = vec![0u64; nwords];
        let copy_n = init_mask.len().min(nwords);
        mask[..copy_n].copy_from_slice(&init_mask[..copy_n]);
        Cell { data, mask: mask.into_boxed_slice() }
    }

    fn take_cell(&mut self, size: usize, init_mask: &[u64]) -> Cell {
        if let Some(mut cell) = self.buffer_pool.get_mut(&size).and_then(|v| v.pop()) {
            for d in cell.data.iter_mut() { *d = HANDLE_NONE; }
            for w in cell.mask.iter_mut() { *w = 0; }
            let copy_n = init_mask.len().min(cell.mask.len());
            cell.mask[..copy_n].copy_from_slice(&init_mask[..copy_n]);
            cell
        } else {
            Self::make_cell(size, init_mask)
        }
    }

    fn recycle(&mut self, cell: Cell) {
        self.buffer_pool.entry(cell.data.len()).or_default().push(cell);
    }

    fn alloc_inner(&mut self, size: usize, init_mask: &[u64]) -> Result<(u32, u32), String> {
        let cell = self.take_cell(size, init_mask);
        self.bytes_used = self.bytes_used.saturating_add(
            size * 8 + mask_words_for(size) * 8
        );
        let (h, g) = if let Some(h) = self.free_list.pop() {
            let idx = h as usize;
            self.cells[idx] = Some(cell);
            self.rc[idx] = 1;
            self.generation[idx] = self.generation[idx].wrapping_add(1);
            (h, self.generation[idx])
        } else {
            if self.cells.len() as u32 > HANDLE_SLOT_MAX {
                return Err(format!(
                    "heap.alloc: slot index would exceed handle limit {}",
                    HANDLE_SLOT_MAX
                ));
            }
            self.cells.push(Some(cell));
            self.rc.push(1);
            self.generation.push(0);
            let h = (self.cells.len() - 1) as u32;
            (h, 0)
        };
        if self.traced(h) {
            htrace!(self, "[ALLOC] slot {} gen {} size {}", h, g, size);
        }
        Ok((h, g))
    }

    pub fn alloc(&mut self, size: usize) -> (u32, u32) {
        self.alloc_inner(size, &[]).expect("alloc slot overflow")
    }

    pub fn alloc_with_mask(&mut self, size: usize, init_mask: &[u64]) -> (u32, u32) {
        self.alloc_inner(size, init_mask).expect("alloc slot overflow")
    }

    pub fn try_alloc(&mut self, size: usize) -> Result<(u32, u32), String> {
        self.alloc_inner(size, &[])
    }

    pub fn try_alloc_with_mask(&mut self, size: usize, init_mask: &[u64]) -> Result<(u32, u32), String> {
        self.alloc_inner(size, init_mask)
    }

    fn check(&self, slot: u32, generation: u32, op: &str) -> Result<usize, String> {
        let idx = slot as usize;
        if idx >= self.cells.len() {
            return Err(format!("{}: invalid slot {}", op, slot));
        }
        if self.cells[idx].is_none() {
            return Err(format!("{}: use-after-free of slot {}", op, slot));
        }
        if self.generation[idx] != generation {
            return Err(format!(
                "{}: stale handle for slot {} (have generation {}, live generation {})",
                op, slot, generation, self.generation[idx]
            ));
        }
        Ok(idx)
    }

    // Returns (raw u64, is_handle_bit).
    pub fn ld(&self, slot: u32, generation: u32, offset: usize) -> Result<(u64, bool), String> {
        let idx = self.check(slot, generation, "ld")?;
        let cell = self.cells[idx].as_ref().unwrap();
        if offset >= cell.data.len() {
            return Err(format!("ld: offset {} out of bounds (size {})", offset, cell.data.len()));
        }
        let val = cell.data[offset];
        let is_handle = mask_bit(&cell.mask, offset);
        Ok((val, is_handle))
    }

    // Store raw u64; caller supplies whether it is a handle. Returns (old_data, old_is_handle).
    pub fn st(
        &mut self, slot: u32, generation: u32,
        offset: usize, val: u64, is_handle: bool,
    ) -> Result<(u64, bool), String> {
        let idx = self.check(slot, generation, "st")?;
        let cell = self.cells[idx].as_mut().unwrap();
        if offset >= cell.data.len() {
            return Err(format!("st: offset {} out of bounds (size {})", offset, cell.data.len()));
        }
        let old = cell.data[offset];
        let old_is_handle = mask_bit(&cell.mask, offset);
        cell.data[offset] = val;
        mask_set(&mut cell.mask, offset, is_handle);
        Ok((old, old_is_handle))
    }

    pub fn cell_data(&self, slot: u32, generation: u32) -> Result<&[u64], String> {
        let idx = self.check(slot, generation, "cell")?;
        Ok(&self.cells[idx].as_ref().unwrap().data)
    }

    pub fn cell_data_mut(&mut self, slot: u32, generation: u32) -> Result<&mut [u64], String> {
        let idx = self.check(slot, generation, "cell_mut")?;
        Ok(&mut self.cells[idx].as_mut().unwrap().data)
    }

    pub fn cell_mask(&self, slot: u32, generation: u32) -> Result<&[u64], String> {
        let idx = self.check(slot, generation, "cell_mask")?;
        Ok(&self.cells[idx].as_ref().unwrap().mask)
    }

    pub fn size(&self, slot: u32, generation: u32) -> Result<usize, String> {
        let idx = self.check(slot, generation, "size")?;
        Ok(self.cells[idx].as_ref().unwrap().data.len())
    }

    #[inline]
    pub fn is_live(&self, slot: u32, generation: u32) -> bool {
        let idx = slot as usize;
        idx < self.cells.len()
            && self.cells[idx].is_some()
            && self.generation[idx] == generation
    }

    pub fn rc_inc_handle(&mut self, raw: u64) -> Result<(), String> {
        if raw == HANDLE_NONE { return Ok(()); }
        let (s, g) = handle_parts(raw);
        self.rc_inc(s, g)
    }

    pub fn rc_dec_handle(&mut self, raw: u64) -> Result<(), String> {
        if raw == HANDLE_NONE { return Ok(()); }
        let (s, g) = handle_parts(raw);
        self.rc_dec(s, g)?;
        Ok(())
    }

    pub fn rc_inc(&mut self, slot: u32, generation: u32) -> Result<(), String> {
        let trace = self.traced(slot);
        let idx = self.check(slot, generation, "rc_inc")?;
        self.rc[idx] = self.rc[idx]
            .checked_add(1)
            .ok_or_else(|| format!("rc_inc: refcount overflow on slot {}", slot))?;
        if trace {
            htrace!(self, "[RC_INC] pc {} slot {} gen {} -> rc {}", self.trace_pc, slot, generation, self.rc[idx]);
        }
        Ok(())
    }

    // At rc=0: recursively rc_dec child handles per cell mask, then reclaim.
    pub fn rc_dec(&mut self, slot: u32, generation: u32) -> Result<bool, String> {
        let trace = self.traced(slot);
        if trace {
            let live_gen = self.generation.get(slot as usize).copied().unwrap_or(0);
            let is_live = self.cells.get(slot as usize).map(|c| c.is_some()).unwrap_or(false);
            htrace!(self, "[RC_DEC] pc {} slot {} gen {} (live_gen {} live={})", self.trace_pc, slot, generation, live_gen, is_live);
        }
        let idx = self.check(slot, generation, "rc_dec")?;
        if self.rc[idx] == 0 {
            return Err(format!("rc_dec: refcount underflow on slot {}", slot));
        }
        self.rc[idx] -= 1;
        if trace {
            htrace!(self, "[RC_DEC] slot {} -> rc {}", slot, self.rc[idx]);
        }
        if self.rc[idx] != 0 {
            return Ok(false);
        }
        let cell = self.cells[idx].take().unwrap();
        let size = cell.data.len();
        self.bytes_used = self.bytes_used.saturating_sub(size * 8 + mask_words_for(size) * 8);
        self.free_list.push(slot);
        if trace {
            htrace!(self, "[RC_DEC] slot {} FREED via rc_dec", slot);
        }
        for i in 0..size {
            if mask_bit(&cell.mask, i) {
                let v = cell.data[i];
                if v != HANDLE_NONE {
                    let s = ((v >> 24) & 0x00FF_FFFF) as u32;
                    let g = (v & 0x00FF_FFFF) as u32;
                    self.rc_dec(s, g)?;
                }
            }
        }
        self.recycle(cell);
        Ok(true)
    }

    // Idempotent against rc=0 reclaim. Used by region_pop.
    pub fn force_free(&mut self, slot: u32, generation: u32) -> Result<(), String> {
        let trace = self.traced(slot);
        let idx = slot as usize;
        if idx >= self.cells.len() { return Ok(()); }
        if self.cells[idx].is_none() {
            if trace { htrace!(self, "[FORCE_FREE] slot {} gen {} -- already none", slot, generation); }
            return Ok(());
        }
        if self.generation[idx] != generation {
            if trace { htrace!(self, "[FORCE_FREE] slot {} gen {} -- generation mismatch (live {})", slot, generation, self.generation[idx]); }
            return Ok(());
        }
        if trace {
            htrace!(self, "[FORCE_FREE] slot {} gen {} rc was {}", slot, generation, self.rc[idx]);
        }
        let cell = self.cells[idx].take().unwrap();
        let size = cell.data.len();
        self.bytes_used = self.bytes_used.saturating_sub(size * 8 + mask_words_for(size) * 8);
        self.rc[idx] = 0;
        self.free_list.push(slot);
        for i in 0..size {
            if mask_bit(&cell.mask, i) {
                let v = cell.data[i];
                if v != HANDLE_NONE {
                    let s = ((v >> 24) & 0x00FF_FFFF) as u32;
                    let g = (v & 0x00FF_FFFF) as u32;
                    if let Err(e) = self.rc_dec(s, g) {
                        debug_assert!(false, "force_free cascade rc_dec failed: {}", e);
                    }
                }
            }
        }
        self.recycle(cell);
        Ok(())
    }

    pub fn live_count(&self) -> usize {
        self.cells.iter().filter(|c| c.is_some()).count()
    }

    pub fn rc(&self, slot: u32, generation: u32) -> Option<u32> {
        self.check(slot, generation, "rc").ok().map(|idx| self.rc[idx])
    }

    // Debug: every live cell as (slot, generation, rc, data, is_handle-per-slot).
    // Used by the leak-localization dump; not on any hot path.
    pub fn live_cells(&self) -> Vec<(u32, u32, u32, Vec<u64>, Vec<bool>)> {
        self.cells.iter().enumerate().filter_map(|(i, c)| {
            c.as_ref().map(|cell| {
                let handles: Vec<bool> = (0..cell.data.len())
                    .map(|j| mask_bit(&cell.mask, j))
                    .collect();
                (
                    i as u32,
                    self.generation.get(i).copied().unwrap_or(0),
                    self.rc.get(i).copied().unwrap_or(0),
                    cell.data.to_vec(),
                    handles,
                )
            })
        }).collect()
    }

    pub fn clear(&mut self) {
        self.cells.clear();
        self.rc.clear();
        self.generation.clear();
        self.free_list.clear();
        self.bytes_used = 0;
        self.buffer_pool.clear();
    }
}