use super::*;
#[test]
fn fixed_point_batch_direct_ifds_reuses_prepare_and_solve_scratch() {
let prepare_calls = Cell::new(0usize);
let solve_phase = Cell::new(0usize);
let second_prepare_saw_retained_slots = Cell::new(false);
let second_solve_saw_retained_slot = Cell::new(false);
let dispatch = |_: &vyre::ir::Program,
inputs: &[&[u8]],
grid: Option<[u32; 3]>,
outputs: &mut Vec<Vec<u8>>| {
if inputs.len() < 6 {
return crate::fixed_point_closure::fallback_bitset_equal_dispatch(inputs, outputs);
}
match inputs.len() {
17 => {
assert_eq!(grid, Some([1, 1, 1]));
if prepare_calls.get() > 0 && outputs.len() == 4 {
second_prepare_saw_retained_slots.set(true);
}
prepare_calls.set(prepare_calls.get() + 1);
if outputs.len() < 4 {
outputs.resize_with(4, Vec::new);
}
for slot in outputs.iter_mut() {
slot.clear();
}
for word in [0u32, 1, 1] {
outputs[0].extend_from_slice(&word.to_le_bytes());
}
for word in [0u32, 0] {
outputs[1].extend_from_slice(&word.to_le_bytes());
}
outputs[2].extend_from_slice(&1u32.to_le_bytes());
outputs[3].extend_from_slice(&1u32.to_le_bytes());
}
7 => {
assert_eq!(grid, Some([2, 1, 1]));
if outputs.is_empty() {
outputs.push(Vec::with_capacity(8));
} else if solve_phase.get() >= 2 {
second_solve_saw_retained_slot.set(true);
}
let frontier = u32::from_le_bytes(inputs[5][0..4].try_into().unwrap());
let next = match solve_phase.get() {
0 => {
assert_eq!(frontier, 0b01);
solve_phase.set(1);
0b11u32
}
1 => {
assert_eq!(frontier, 0b11);
solve_phase.set(2);
0b11u32
}
2 => {
assert_eq!(frontier, 0b10);
solve_phase.set(3);
0b11u32
}
3 => {
assert_eq!(frontier, 0b11);
solve_phase.set(4);
0b11u32
}
call => panic!("unexpected direct IFDS dispatch {call}"),
};
outputs[0].clear();
outputs[0].extend_from_slice(&next.to_le_bytes());
}
len => panic!("unexpected IFDS input count {len}"),
}
Ok(())
};
let mut batch = FixedPointBatch::new(&dispatch);
let first = batch
.ifds(1, 2, 1, &[(0, 0, 1)], &[], &[], &[], &[(0, 0, 0)], 4)
.expect("first direct batch IFDS solve must converge");
let prepare_slots = batch.ifds_prepare_scratch().output_slot_count();
let solve_slots = batch.ifds_scratch().output_slot_count();
let frontier_capacity = batch.ifds_scratch().frontier_word_capacity();
let second = batch
.ifds(1, 2, 1, &[(0, 0, 1)], &[], &[], &[], &[(0, 1, 0)], 4)
.expect("second direct batch IFDS solve must reuse scratch and converge");
assert_eq!(first, vec![0, 1024]);
assert_eq!(second, vec![0, 1024]);
assert!(second_prepare_saw_retained_slots.get());
assert!(second_solve_saw_retained_slot.get());
assert_eq!(
batch.ifds_prepare_scratch().output_slot_count(),
prepare_slots
);
assert_eq!(batch.ifds_scratch().output_slot_count(), solve_slots);
assert_eq!(
batch.ifds_scratch().frontier_word_capacity(),
frontier_capacity
);
}
#[test]
fn fixed_point_batch_direct_many_ifds_prepares_once_and_reuses_solve_scratch() {
let solve_phase = Cell::new(0usize);
let prepare_calls = Cell::new(0usize);
let retained_slot_seen = Cell::new(0usize);
let dispatch = |_: &vyre::ir::Program,
inputs: &[&[u8]],
grid: Option<[u32; 3]>,
outputs: &mut Vec<Vec<u8>>| {
if inputs.len() < 6 {
return crate::fixed_point_closure::fallback_bitset_equal_dispatch(inputs, outputs);
}
match inputs.len() {
17 => {
assert_eq!(grid, Some([1, 1, 1]));
prepare_calls.set(prepare_calls.get() + 1);
outputs.clear();
outputs.resize_with(4, Vec::new);
for word in [0u32, 1, 1] {
outputs[0].extend_from_slice(&word.to_le_bytes());
}
for word in [0u32, 0] {
outputs[1].extend_from_slice(&word.to_le_bytes());
}
outputs[2].extend_from_slice(&1u32.to_le_bytes());
outputs[3].extend_from_slice(&1u32.to_le_bytes());
}
7 => {
assert_eq!(grid, Some([2, 1, 1]));
if outputs.is_empty() {
outputs.push(Vec::with_capacity(8));
} else {
retained_slot_seen.set(retained_slot_seen.get() + 1);
}
let frontier = u32::from_le_bytes(inputs[5][0..4].try_into().unwrap());
let next = match solve_phase.get() {
0 => {
assert_eq!(frontier, 0b01);
solve_phase.set(1);
0b11u32
}
1 => {
assert_eq!(frontier, 0b11);
solve_phase.set(2);
0b11u32
}
2 => {
assert_eq!(frontier, 0b10);
solve_phase.set(3);
0b11u32
}
3 => {
assert_eq!(frontier, 0b11);
solve_phase.set(4);
0b11u32
}
call => panic!("unexpected direct many IFDS dispatch {call}"),
};
outputs[0].clear();
outputs[0].extend_from_slice(&next.to_le_bytes());
}
len => panic!("unexpected IFDS input count {len}"),
}
Ok(())
};
let mut batch = FixedPointBatch::new(&dispatch);
let seed_a: &[(u32, u32, u32)] = &[(0, 0, 0)];
let seed_b: &[(u32, u32, u32)] = &[(0, 1, 0)];
let mut results = vec![Vec::with_capacity(2), Vec::with_capacity(2)];
let outer_ptr = results.as_ptr();
let row_a_ptr = results[0].as_ptr();
let row_b_ptr = results[1].as_ptr();
batch
.ifds_many_into(
1,
2,
1,
&[(0, 0, 1)],
&[],
&[],
&[],
&[seed_a, seed_b],
4,
&mut results,
)
.expect("direct many IFDS batch must prepare once and fill caller-owned rows");
assert_eq!(results, vec![vec![0, 1024], vec![0, 1024]]);
assert_eq!(results.as_ptr(), outer_ptr);
assert_eq!(results[0].as_ptr(), row_a_ptr);
assert_eq!(results[1].as_ptr(), row_b_ptr);
assert_eq!(prepare_calls.get(), 1);
assert!(retained_slot_seen.get() >= 3);
assert_eq!(batch.ifds_prepare_scratch().output_slot_count(), 4);
assert_eq!(batch.ifds_scratch().output_slot_count(), 1);
}