use crate::ifds_resident_types::IfdsResidentDispatch;
mod parallel_scratch;
mod prepared_csr;
mod single_scratch;
pub use parallel_scratch::{
allocate_resident_ifds_parallel_scratch,
allocate_resident_ifds_parallel_scratch_with_capacities,
allocate_resident_ifds_parallel_scratch_with_iterations, free_resident_ifds_parallel_scratch,
};
pub use prepared_csr::{
free_resident_prepared_ifds_csr, prepare_ifds_csr_resident_via,
upload_prepared_ifds_csr_resident,
};
pub use single_scratch::{
allocate_resident_ifds_scratch, allocate_resident_ifds_scratch_with_seed_capacity,
free_resident_ifds_scratch,
};
pub(super) fn free_unique_resident_resources<D>(
dispatch: &D,
resources: impl IntoIterator<Item = (&'static str, D::Resource)>,
) -> Result<(), String>
where
D: IfdsResidentDispatch,
{
let mut unique = Vec::new();
for (label, resource) in resources {
if unique
.iter()
.any(|(_, seen): &(&'static str, D::Resource)| seen == &resource)
{
continue;
}
unique.try_reserve(1).map_err(|error| {
format!(
"weir IFDS resident free tracking could not reserve one resource slot: {error}. Fix: shard resident cleanup before teardown."
)
})?;
unique.push((label, resource));
}
let mut first_error = None;
for (label, resource) in unique {
if let Err(error) = dispatch.free_resident(resource) {
first_error
.get_or_insert_with(|| format!("weir IFDS resident free {label} failed: {error}"));
}
}
match first_error {
Some(error) => Err(error),
None => Ok(()),
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::Mutex;
use vyre_foundation::ir::Program;
#[derive(Default)]
struct FreeRecorder {
freed: Mutex<Vec<u64>>,
}
impl IfdsResidentDispatch for FreeRecorder {
type Resource = u64;
fn resident_backend_id(&self) -> &'static str {
"free_recorder"
}
fn allocate_resident(&self, _byte_len: usize) -> Result<Self::Resource, String> {
unreachable!("allocate_resident is not used by free tracking tests")
}
fn upload_resident(&self, _resource: &Self::Resource, _bytes: &[u8]) -> Result<(), String> {
unreachable!("upload_resident is not used by free tracking tests")
}
fn upload_resident_many(
&self,
_uploads: &[(&Self::Resource, &[u8])],
) -> Result<(), String> {
unreachable!("upload_resident_many is not used by free tracking tests")
}
fn download_resident(&self, _resource: &Self::Resource) -> Result<Vec<u8>, String> {
unreachable!("download_resident is not used by free tracking tests")
}
fn download_resident_into(
&self,
_resource: &Self::Resource,
_output: &mut Vec<u8>,
) -> Result<(), String> {
unreachable!("download_resident_into is not used by free tracking tests")
}
fn download_resident_range(
&self,
_resource: &Self::Resource,
_byte_offset: usize,
_byte_len: usize,
) -> Result<Vec<u8>, String> {
unreachable!("download_resident_range is not used by free tracking tests")
}
fn download_resident_range_into(
&self,
_resource: &Self::Resource,
_byte_offset: usize,
_byte_len: usize,
_output: &mut Vec<u8>,
) -> Result<(), String> {
unreachable!("download_resident_range_into is not used by free tracking tests")
}
fn free_resident(&self, resource: Self::Resource) -> Result<(), String> {
self.freed
.lock()
.expect("free recorder mutex should not be poisoned")
.push(resource);
Ok(())
}
fn dispatch_resident(
&self,
_program: &Program,
_resources: &[Self::Resource],
_grid_override: Option<[u32; 3]>,
) -> Result<(), String> {
unreachable!("dispatch_resident is not used by free tracking tests")
}
}
#[test]
fn generated_free_unique_resident_resources_releases_first_seen_handles_once() {
for case in 0..4096u64 {
let resources = [
("a", case & 0x3f),
("b", (case >> 2) & 0x3f),
("c", (case >> 4) & 0x3f),
("d", case & 0x3f),
("e", (case >> 8) & 0x3f),
("f", (case >> 2) & 0x3f),
];
let recorder = FreeRecorder::default();
free_unique_resident_resources(&recorder, resources)
.expect("generated unique free should succeed");
let mut expected = Vec::new();
for (_, resource) in resources {
if !expected.contains(&resource) {
expected.push(resource);
}
}
let freed = recorder
.freed
.lock()
.expect("free recorder mutex should not be poisoned")
.clone();
assert_eq!(freed, expected, "case={case}");
}
}
}