qualia_core_db/inference/runtime/kv/paged/
pool.rs

1use super::config::INVALID_BLOCK;
2
3#[derive(Clone, Copy, Debug, PartialEq, Eq)]
4pub enum PoolError {
5    Exhausted,
6    InvalidBlock,
7    ReferenceOverflow,
8    DoubleRelease,
9}
10
11#[derive(Clone, Copy, Debug, PartialEq, Eq)]
12pub enum CopyOnWrite {
13    Existing(u32),
14    Allocated(u32),
15    Copy { source: u32, destination: u32 },
16}
17
18/// Fixed-capacity physical-page allocator with reference counts.
19///
20/// Both vectors are fully sized during cold construction. Hot operations only mutate elements or
21/// the existing vector length; `free` has capacity `physical_blocks`, so `push` cannot reallocate.
22#[derive(Debug)]
23pub struct BlockPool {
24    free: Vec<u32>,
25    refs: Vec<u32>,
26}
27
28impl BlockPool {
29    pub fn new(physical_blocks: u32) -> Self {
30        let mut free = Vec::with_capacity(physical_blocks as usize);
31        free.extend((0..physical_blocks).rev());
32        Self {
33            free,
34            refs: vec![0; physical_blocks as usize],
35        }
36    }
37
38    pub fn allocate(&mut self) -> Result<u32, PoolError> {
39        let block = self.free.pop().ok_or(PoolError::Exhausted)?;
40        self.refs[block as usize] = 1;
41        Ok(block)
42    }
43
44    pub fn retain(&mut self, block: u32) -> Result<(), PoolError> {
45        let count = self
46            .refs
47            .get_mut(block as usize)
48            .ok_or(PoolError::InvalidBlock)?;
49        if *count == 0 {
50            return Err(PoolError::InvalidBlock);
51        }
52        *count = count.checked_add(1).ok_or(PoolError::ReferenceOverflow)?;
53        Ok(())
54    }
55
56    pub fn release(&mut self, block: u32) -> Result<(), PoolError> {
57        let count = self
58            .refs
59            .get_mut(block as usize)
60            .ok_or(PoolError::InvalidBlock)?;
61        if *count == 0 {
62            return Err(PoolError::DoubleRelease);
63        }
64        *count -= 1;
65        if *count == 0 {
66            debug_assert!(self.free.len() < self.free.capacity());
67            self.free.push(block);
68        }
69        Ok(())
70    }
71
72    pub fn ref_count(&self, block: u32) -> Option<u32> {
73        self.refs.get(block as usize).copied()
74    }
75
76    pub fn free_count(&self) -> usize {
77        self.free.len()
78    }
79
80    pub fn writable(&mut self, current: u32) -> Result<CopyOnWrite, PoolError> {
81        if current == INVALID_BLOCK {
82            return self.allocate().map(CopyOnWrite::Allocated);
83        }
84        match self.ref_count(current) {
85            Some(1) => Ok(CopyOnWrite::Existing(current)),
86            Some(count) if count > 1 => {
87                let destination = self.allocate()?;
88                self.release(current)?;
89                Ok(CopyOnWrite::Copy {
90                    source: current,
91                    destination,
92                })
93            }
94            _ => Err(PoolError::InvalidBlock),
95        }
96    }
97}
qualia_core_db/inference/runtime/kv/paged/pool.rs

qualia_core_db/inference/runtime/kv/paged/
pool.rs
