use ic_stable_structures::{Memory, memory_manager::VirtualMemory};
use std::{cell::RefCell, rc::Rc};
pub(super) struct GrowthState<M: Memory> {
pub backing: Rc<M>,
pub bucket_size_pages: u16,
pub allocated_buckets: RefCell<u16>,
}
pub struct RuntimeMemory<M: Memory> {
pub(super) memory: VirtualMemory<Rc<M>>,
pub(super) growth: Rc<GrowthState<M>>,
}
impl<M: Memory> Clone for RuntimeMemory<M> {
fn clone(&self) -> Self {
Self {
memory: self.memory.clone(),
growth: Rc::clone(&self.growth),
}
}
}
impl<M: Memory> RuntimeMemory<M> {
pub fn grow(&self, pages: u64) -> Result<u64, super::RuntimeGrowError> {
use super::RuntimeGrowError;
let mut allocated = self
.growth
.allocated_buckets
.try_borrow_mut()
.map_err(|_| RuntimeGrowError::ReentrantAccess)?;
let old_pages = self.memory.size();
if pages == 0 {
return Ok(old_pages);
}
let new_pages = old_pages
.checked_add(pages)
.ok_or(RuntimeGrowError::ArithmeticOverflow)?;
let bucket_pages = u64::from(self.growth.bucket_size_pages);
let extra = new_pages.div_ceil(bucket_pages) - old_pages.div_ceil(bucket_pages);
let total = u64::from(*allocated)
.checked_add(extra)
.ok_or(RuntimeGrowError::ArithmeticOverflow)?;
if total > u64::from(super::layout::BUCKET_CAPACITY) {
return Err(RuntimeGrowError::BucketExhausted {
required_buckets: total,
capacity: super::layout::BUCKET_CAPACITY,
});
}
#[expect(
clippy::cast_possible_truncation,
reason = "admission bounds total by the u16 bucket capacity"
)]
let total_buckets = total as u16;
let required_pages = 1 + total * bucket_pages;
let physical_pages = self.growth.backing.size();
if required_pages > physical_pages
&& self.growth.backing.grow(required_pages - physical_pages) < 0
{
return Err(RuntimeGrowError::BackingRefused {
additional_pages: required_pages - physical_pages,
});
}
assert_eq!(
self.memory.grow(pages),
old_pages.cast_signed(),
"preflighted manager growth returns the previous virtual extent"
);
*allocated = total_buckets;
Ok(old_pages)
}
}
impl<M: Memory> Memory for RuntimeMemory<M> {
fn size(&self) -> u64 {
self.memory.size()
}
fn grow(&self, pages: u64) -> i64 {
Self::grow(self, pages).map_or(-1, u64::cast_signed)
}
fn read(&self, offset: u64, dst: &mut [u8]) {
self.memory.read(offset, dst);
}
#[expect(
unsafe_code,
reason = "delegate the upstream raw-read contract unchanged"
)]
unsafe fn read_unsafe(&self, offset: u64, dst: *mut u8, count: usize) {
unsafe { self.memory.read_unsafe(offset, dst, count) }
}
fn write(&self, offset: u64, src: &[u8]) {
self.memory.write(offset, src);
}
}