use alloc::vec::Vec;
use concinnity_memory::{Arena, ArenaVec};
#[derive(Clone, Copy)]
pub struct FrameContext<'a> {
pub scratch: &'a Arena,
}
impl<'a> FrameContext<'a> {
pub fn new(scratch: &'a Arena) -> Self {
Self { scratch }
}
pub fn collect<T, I>(&self, items: I) -> FrameVec<'a, T>
where
T: Copy,
I: IntoIterator<Item = T>,
{
let items = items.into_iter();
let reservation = items
.size_hint()
.1
.and_then(|upper| self.scratch.vec::<T>(upper));
match reservation {
Some(mut out) => {
out.extend(items);
FrameVec::Scratch(out)
}
None => FrameVec::Heap(items.collect()),
}
}
pub fn filled<T: Copy>(&self, len: usize, value: T) -> FrameVec<'a, T> {
match self.scratch.vec::<T>(len) {
Some(mut out) => {
out.extend(core::iter::repeat_n(value, len));
FrameVec::Scratch(out)
}
None => FrameVec::Heap(alloc::vec![value; len]),
}
}
pub fn vec<T: Copy>(&self, capacity: usize) -> FrameVec<'a, T> {
match self.scratch.vec::<T>(capacity) {
Some(out) => FrameVec::Scratch(out),
None => FrameVec::Heap(Vec::new()),
}
}
}
pub enum FrameVec<'a, T: Copy> {
Scratch(ArenaVec<'a, T>),
Heap(Vec<T>),
}
impl<T: Copy> FrameVec<'_, T> {
pub fn push(&mut self, value: T) {
match self {
FrameVec::Scratch(v) => {
if !v.push(value) {
let mut heap = Vec::with_capacity(v.len() + 1);
heap.extend_from_slice(v);
heap.push(value);
*self = FrameVec::Heap(heap);
}
}
FrameVec::Heap(v) => v.push(value),
}
}
}
impl<T: Copy> core::ops::Deref for FrameVec<'_, T> {
type Target = [T];
fn deref(&self) -> &[T] {
match self {
FrameVec::Scratch(v) => v,
FrameVec::Heap(v) => v,
}
}
}
impl<T: Copy> core::ops::DerefMut for FrameVec<'_, T> {
fn deref_mut(&mut self) -> &mut [T] {
match self {
FrameVec::Scratch(v) => v,
FrameVec::Heap(v) => v,
}
}
}
impl<'v, T: Copy> IntoIterator for &'v FrameVec<'_, T> {
type Item = &'v T;
type IntoIter = core::slice::Iter<'v, T>;
fn into_iter(self) -> Self::IntoIter {
self.iter()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn a_gather_that_fits_lands_in_scratch() {
let arena = Arena::with_capacity(4096);
let frame = FrameContext::new(&arena);
let out = frame.collect([1u32, 2, 3]);
assert!(matches!(out, FrameVec::Scratch(_)));
assert_eq!(&*out, &[1, 2, 3]);
assert_eq!(arena.overflows(), 0);
assert!(arena.used() > 0, "it came out of the reserve");
}
#[test]
fn a_gather_too_large_falls_back_to_the_heap_and_is_recorded() {
let arena = Arena::with_capacity(8);
let frame = FrameContext::new(&arena);
let out = frame.collect([1u64, 2, 3, 4]);
assert!(matches!(out, FrameVec::Heap(_)));
assert_eq!(&*out, &[1, 2, 3, 4], "the fallback holds the same values");
assert_eq!(arena.overflows(), 1);
}
#[test]
fn both_arms_read_the_same_way() {
let roomy = Arena::with_capacity(4096);
let tight = Arena::with_capacity(0);
let items = [7u16, 8, 9];
let from_scratch = FrameContext::new(&roomy).collect(items);
let from_heap = FrameContext::new(&tight).collect(items);
assert_eq!(&*from_scratch, &*from_heap);
assert_eq!(from_scratch.len(), 3);
assert_eq!(from_heap.iter().copied().sum::<u16>(), 24);
for (a, b) in (&from_scratch).into_iter().zip(&from_heap) {
assert_eq!(a, b);
}
}
#[test]
fn a_filled_frame_is_writable_in_place() {
let arena = Arena::with_capacity(4096);
let mut frame = FrameContext::new(&arena).filled(4, None::<u32>);
assert!(matches!(frame, FrameVec::Scratch(_)));
assert_eq!(&*frame, &[None, None, None, None]);
frame[2] = Some(9);
assert_eq!(&*frame, &[None, None, Some(9), None]);
}
#[test]
fn a_filled_frame_that_overflowed_is_still_writable() {
let arena = Arena::with_capacity(0);
let mut frame = FrameContext::new(&arena).filled(3, 0u32);
assert!(matches!(frame, FrameVec::Heap(_)));
frame[1] = 5;
assert_eq!(&*frame, &[0, 5, 0]);
assert_eq!(arena.overflows(), 1);
}
#[test]
fn a_reserved_frame_takes_pushes_in_scratch() {
let arena = Arena::with_capacity(4096);
let mut out = FrameContext::new(&arena).vec::<u32>(3);
out.push(1);
out.push(2);
assert!(matches!(out, FrameVec::Scratch(_)));
assert_eq!(&*out, &[1, 2]);
assert_eq!(arena.overflows(), 0);
}
#[test]
fn a_push_past_the_reservation_moves_to_the_heap() {
let arena = Arena::with_capacity(4096);
let mut out = FrameContext::new(&arena).vec::<u32>(2);
out.push(1);
out.push(2);
out.push(3);
assert!(matches!(out, FrameVec::Heap(_)));
assert_eq!(&*out, &[1, 2, 3]);
}
#[test]
fn a_reservation_the_reserve_cannot_hold_starts_on_the_heap() {
let arena = Arena::with_capacity(8);
let mut out = FrameContext::new(&arena).vec::<u64>(64);
assert!(matches!(out, FrameVec::Heap(_)));
assert_eq!(arena.overflows(), 1, "the decline is recorded");
out.push(7);
assert_eq!(&*out, &[7]);
}
#[test]
fn an_empty_gather_costs_nothing() {
let arena = Arena::with_capacity(4096);
let out = FrameContext::new(&arena).collect([0u8; 0]);
assert!(out.is_empty());
assert_eq!(arena.overflows(), 0);
}
#[test]
fn a_copied_context_shares_one_reserve() {
let arena = Arena::with_capacity(4096);
let frame = FrameContext::new(&arena);
let copy = frame;
let _a = frame.collect([1u32; 4]);
let used = arena.used();
let _b = copy.collect([2u32; 4]);
assert!(arena.used() > used, "the copy drew from the same reserve");
}
}