use std::cell::RefCell;
use vyre::ir::Program;
const DEFAULT_OWNED_INPUT_SLOT_CAPACITY: usize = 8;
pub(crate) const FIXPOINT_FRONTIER_INPUT_SLOT: usize = 5;
pub(crate) const FIXPOINT_FRONTIER_INPUT_SLOTS: &[usize] = &[FIXPOINT_FRONTIER_INPUT_SLOT];
pub(crate) struct OwnedDispatchInputs {
buffers: RefCell<Vec<Vec<u8>>>,
}
impl OwnedDispatchInputs {
pub(crate) fn new() -> Self {
Self {
buffers: RefCell::new(Vec::new()),
}
}
pub(crate) fn dispatch<F>(
&self,
program: &Program,
inputs: &[&[u8]],
grid: Option<[u32; 3]>,
dispatch: &F,
) -> Result<Vec<Vec<u8>>, String>
where
F: Fn(&Program, &[Vec<u8>], Option<[u32; 3]>) -> Result<Vec<Vec<u8>>, String>,
{
self.with_owned(inputs, |owned| dispatch(program, owned, grid))
}
pub(crate) fn dispatch_refreshing_slots<F>(
&self,
program: &Program,
inputs: &[&[u8]],
grid: Option<[u32; 3]>,
refresh_slots: &[usize],
dispatch: &F,
) -> Result<Vec<Vec<u8>>, String>
where
F: Fn(&Program, &[Vec<u8>], Option<[u32; 3]>) -> Result<Vec<Vec<u8>>, String>,
{
self.with_owned_refreshing_slots(inputs, refresh_slots, |owned| {
dispatch(program, owned, grid)
})
}
pub(crate) fn dispatch_refreshing_frontier<F>(
&self,
program: &Program,
inputs: &[&[u8]],
grid: Option<[u32; 3]>,
dispatch: &F,
) -> Result<Vec<Vec<u8>>, String>
where
F: Fn(&Program, &[Vec<u8>], Option<[u32; 3]>) -> Result<Vec<Vec<u8>>, String>,
{
self.dispatch_refreshing_slots(
program,
inputs,
grid,
FIXPOINT_FRONTIER_INPUT_SLOTS,
dispatch,
)
}
pub(crate) fn with_owned<R, F>(&self, inputs: &[&[u8]], call: F) -> Result<R, String>
where
F: FnOnce(&[Vec<u8>]) -> Result<R, String>,
{
let mut buffers = self.buffers.borrow_mut();
refresh_all(&mut buffers, inputs)?;
call(&buffers)
}
pub(crate) fn with_owned_refreshing_slots<R, F>(
&self,
inputs: &[&[u8]],
refresh_slots: &[usize],
call: F,
) -> Result<R, String>
where
F: FnOnce(&[Vec<u8>]) -> Result<R, String>,
{
let mut buffers = self.buffers.borrow_mut();
refresh_slots_only(&mut buffers, inputs, refresh_slots)?;
call(&buffers)
}
}
fn refresh_all(buffers: &mut Vec<Vec<u8>>, inputs: &[&[u8]]) -> Result<(), String> {
if buffers.len() < inputs.len() {
let reserved_slots = inputs.len().max(DEFAULT_OWNED_INPUT_SLOT_CAPACITY);
crate::staging_reserve::reserve_vec(buffers, reserved_slots, "owned dispatch input slot")?;
buffers.extend(std::iter::repeat_with(Vec::new).take(inputs.len() - buffers.len()));
} else if buffers.len() > inputs.len() {
buffers.truncate(inputs.len());
}
for (slot, input) in buffers.iter_mut().zip(inputs.iter().copied()) {
refresh_slot(slot, input)?;
}
Ok(())
}
fn refresh_slots_only(
buffers: &mut Vec<Vec<u8>>,
inputs: &[&[u8]],
refresh_slots: &[usize],
) -> Result<(), String> {
if buffers.len() != inputs.len() {
return refresh_all(buffers, inputs);
}
for &idx in refresh_slots {
let Some(slot) = buffers.get_mut(idx) else {
return Err(format!(
"owned dispatch input refresh slot {idx} is out of bounds for {} input(s). Fix: keep borrowed and owned ABI adapters aligned.",
buffers.len()
));
};
let Some(input) = inputs.get(idx) else {
return Err(format!(
"borrowed dispatch input refresh slot {idx} is out of bounds for {} input(s). Fix: keep borrowed and owned ABI adapters aligned.",
inputs.len()
));
};
refresh_slot(slot, input)?;
}
Ok(())
}
#[inline]
fn refresh_slot(slot: &mut Vec<u8>, input: &[u8]) -> Result<(), String> {
if slot.len() == input.len() {
slot.copy_from_slice(input);
return Ok(());
}
crate::staging_reserve::reserve_vec(slot, input.len(), "owned dispatch input byte")?;
slot.clear();
slot.extend_from_slice(input);
Ok(())
}
#[cfg(test)]
mod tests {
use super::OwnedDispatchInputs;
#[test]
fn full_refresh_reuses_outer_and_inner_slots() {
let cache = OwnedDispatchInputs::new();
let first_a = [1_u8, 2, 3, 4];
let first_b = [5_u8, 6, 7, 8];
let second_a = [9_u8, 10];
let second_b = [11_u8, 12];
let mut outer_ptr = std::ptr::null();
let mut first_inner_ptr = std::ptr::null();
let mut second_inner_ptr = std::ptr::null();
cache
.with_owned(&[&first_a, &first_b], |owned| {
outer_ptr = owned.as_ptr();
first_inner_ptr = owned[0].as_ptr();
second_inner_ptr = owned[1].as_ptr();
Ok(())
})
.unwrap();
cache
.with_owned(&[&second_a, &second_b], |owned| {
assert_eq!(owned.as_ptr(), outer_ptr);
assert_eq!(owned[0].as_ptr(), first_inner_ptr);
assert_eq!(owned[1].as_ptr(), second_inner_ptr);
assert_eq!(owned[0], second_a);
assert_eq!(owned[1], second_b);
Ok(())
})
.unwrap();
}
#[test]
fn full_refresh_growing_arity_preserves_existing_slots() {
let cache = OwnedDispatchInputs::new();
let first_a = [1_u8, 2, 3, 4];
let first_b = [5_u8, 6, 7, 8];
let second_a = [9_u8, 10, 11, 12];
let second_b = [13_u8, 14, 15, 16];
let second_c = [17_u8, 18, 19, 20];
let mut outer_ptr = std::ptr::null();
let mut first_inner_ptr = std::ptr::null();
let mut second_inner_ptr = std::ptr::null();
cache
.with_owned(&[&first_a, &first_b], |owned| {
outer_ptr = owned.as_ptr();
first_inner_ptr = owned[0].as_ptr();
second_inner_ptr = owned[1].as_ptr();
Ok(())
})
.unwrap();
cache
.with_owned(&[&second_a, &second_b, &second_c], |owned| {
assert_eq!(owned.as_ptr(), outer_ptr);
assert_eq!(owned[0].as_ptr(), first_inner_ptr);
assert_eq!(owned[1].as_ptr(), second_inner_ptr);
assert_eq!(owned[0], second_a);
assert_eq!(owned[1], second_b);
assert_eq!(owned[2], second_c);
Ok(())
})
.unwrap();
}
#[test]
fn hot_slot_refresh_preserves_invariant_inputs() {
let cache = OwnedDispatchInputs::new();
let invariant = [1_u8, 2, 3, 4];
let mutable_a = [5_u8, 6, 7, 8];
let mutable_b = [9_u8, 10, 11, 12];
let updated_a = [13_u8, 14];
let updated_b = [15_u8, 16];
let mut invariant_ptr = std::ptr::null();
let mut first_hot_ptr = std::ptr::null();
let mut second_hot_ptr = std::ptr::null();
cache
.with_owned_refreshing_slots(&[&invariant, &mutable_a, &mutable_b], &[1, 2], |owned| {
invariant_ptr = owned[0].as_ptr();
first_hot_ptr = owned[1].as_ptr();
second_hot_ptr = owned[2].as_ptr();
Ok(())
})
.unwrap();
cache
.with_owned_refreshing_slots(
&[&[99_u8][..], &updated_a, &updated_b],
&[1, 2],
|owned| {
assert_eq!(owned[0].as_ptr(), invariant_ptr);
assert_eq!(owned[1].as_ptr(), first_hot_ptr);
assert_eq!(owned[2].as_ptr(), second_hot_ptr);
assert_eq!(owned[0], invariant);
assert_eq!(owned[1], updated_a);
assert_eq!(owned[2], updated_b);
Ok(())
},
)
.unwrap();
}
#[test]
fn hot_slot_refresh_rejects_bad_slot_index() {
let cache = OwnedDispatchInputs::new();
let input = [1_u8, 2, 3, 4];
cache
.with_owned_refreshing_slots(&[&input], &[0], |_| Ok(()))
.unwrap();
let err = cache
.with_owned_refreshing_slots(&[&input], &[1], |_| Ok(()))
.unwrap_err();
assert!(err.contains("out of bounds"));
assert!(err.contains("Fix:"));
}
#[test]
fn owned_dispatch_input_cache_source_has_no_resize_based_slot_growth() {
let source = include_str!("dispatch_input_cache.rs");
let production = source
.split("#[cfg(test)]")
.next()
.expect("dispatch input cache production source must precede tests");
assert!(
!production.contains(".resize_with(")
&& production.contains("buffers.extend(std::iter::repeat_with(Vec::new)"),
"Fix: owned dispatch input slot staging must extend after fallible reserve instead of resize-driven growth."
);
}
}