tlv

current.rs

use std::{
    cell::{Ref, RefCell, RefMut},
    collections::HashMap,
    rc::Rc,
};

use arrayref::array_ref;

#[derive(Debug)]
pub enum ValueError {
    GenericError,
}

type Disc = [u8; 8];
type LocLen = (u32, u32);

pub trait Held<'a> {
    type Inner: Value<'a>;
    fn get(b: RefMut<'a, [u8]>) -> Result<Self::Inner, ValueError>;
}

pub trait Value<'a> {
    const DISC: Disc;
    fn len(&self) -> u32;
    fn owned(&self) -> bool;
    fn bytes(&self) -> Vec<u8>;
    fn bytes_ref(&self) -> &[u8];
    //Includes Disc and Length
    fn consume(self) -> Vec<u8>;
}

#[derive(Debug)]
pub struct Holder<'a> {
    layout: RefCell<HashMap<Disc, LocLen>>,
    universe: Rc<RefCell<&'a mut [u8]>>,
    unused_space: RefCell<u32>,
    empty_at: RefCell<u32>,
}

impl<'a> Holder<'a> {
    pub fn new(container: Rc<RefCell<&'a mut [u8]>>) -> Result<Self, ValueError> {
        let s = Self {
            layout: RefCell::new(HashMap::new()),
            universe: container,
            unused_space: RefCell::new(0),
            empty_at: RefCell::new(0),
        };
        s.hydrate_blobs()?;
        Ok(s)
    }

    pub fn realloc(&mut self, container: Rc<RefCell<&'a mut [u8]>>) -> Result<(), ValueError> {
        self.universe = container;
        self.hydrate_blobs()?;
        Ok(())
    }

    fn hydrate_blobs(&self) -> Result<(), ValueError> {
        let buffer = &mut self.universe.borrow_mut();
        if buffer.len() == 0 {
            return Ok(());
        }
        let mut offset = 1;
        let (number_of_values, tail) = buffer.split_at_mut(1);
        let number_of_values = number_of_values[0];
        let mut data = tail;
        for _ in 0..number_of_values {
            //[8 bytes discriminator][4 bytes lentgth][N bytes data]

            let (disc, tail) = data.split_at_mut(8);
            let disc: &Disc = array_ref![disc, 0, 8];
            let (length_bytes, tail) = tail.split_at_mut(4);
            let data_length = u32::from_le_bytes(*array_ref![length_bytes, 0, 4]);
            let len = data_length as usize;
            let (_, tail) = tail.split_at_mut(len);
            self.layout
                .borrow_mut()
                .insert(*disc, (offset, data_length));
            offset += 12 + data_length;
            data = tail;
        }
        self.unused_space.replace(buffer.len() as u32 - offset);
        self.empty_at.replace(offset);
        Ok(())
    }

    pub fn insert<F>(&self, blob: F) -> Result<(), ValueError>
    where
        F: Value<'a>,
    {
        if *self.unused_space.borrow() < blob.len() {
            return Err(ValueError::GenericError);
        } else {
            let mut layout = self.layout.borrow_mut();
            let mut buffer = self.universe.borrow_mut();
            let disc = F::DISC;
            let len = blob.len();
            let offset = *self.empty_at.borrow();
            let off_u = offset as usize;
            let end = offset + 12 + len;
            let end_u = end as usize;
            buffer[off_u..off_u + 8].copy_from_slice(&disc);
            buffer[off_u + 8..off_u + 12].copy_from_slice(&len.to_le_bytes());
            if blob.owned() {
                let bytes = blob.bytes();
                buffer[off_u + 12..end_u].copy_from_slice(&bytes);
            } else {
                let bytes = blob.bytes_ref();
                buffer[off_u + 12..end_u].copy_from_slice(bytes);
            }
            layout.insert(disc, (offset, len));
            self.empty_at.replace(end);
            println!("{}. {}, {}", buffer.len(), end, offset);
            self.unused_space.replace(buffer.len() as u32 - end);
            Ok(())
        }
    }

    pub fn remove<F>(&self) -> Result<(), ValueError>
    where
        F: Held<'a> + Value<'a>,
    {
        {
            let mut layout = self.layout.borrow_mut();
            let mut buffer = self.universe.borrow_mut();
            let (offset, len) = layout.remove(&F::DISC).ok_or(ValueError::GenericError)?;
            let offset_u = offset as usize;
            let end = offset + 12 + len;
            let end_u = end as usize;
            buffer[offset_u..end_u].fill(0);
            buffer[0] = buffer[0].checked_sub(1).unwrap_or(0);
            //get next greatest offset
            let empty_at = *self.empty_at.borrow();
            if empty_at - 12 - len > offset {
                let mut move_start = offset + 12 + len;
                let mut move_len = 1;
                for (i, v) in layout.values().enumerate() {
                    if v.0 > offset {
                        if v.0 < move_start {
                            move_start = v.0;
                        }
                        move_len += 12 + v.1;
                        break;
                    }
                }
                buffer.copy_within(move_start as usize..move_len as usize, offset_u);
                *self.empty_at.borrow_mut() = move_len;
                *self.unused_space.borrow_mut() = buffer.len() as u32 - move_len;
            }
        }
        Ok(())
    }

    pub fn get<F>(&'a self) -> Result<F::Inner, ValueError>
    where
        F: Held<'a> + Value<'a>,
    {
        let layout = self.layout.borrow();
        let (offset, len) = layout.get(&F::DISC).ok_or(ValueError::GenericError)?;
        let offset_u = *offset as usize;
        let end = offset + 12 + len;
        let data_u = offset_u + 12;
        let end_u = end as usize;
        let f = RefMut::map(self.universe.borrow_mut(), |s| s[data_u..end_u].as_mut());
        F::get(f)
    }

    pub fn free_space(&self) -> u32 {
        self.unused_space.borrow().clone()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use borsh::{BorshDeserialize, BorshSerialize};
    use bytemuck::{Pod, Zeroable};
    use std::mem;

    #[derive(BorshSerialize, BorshDeserialize, PartialEq, Debug)]
    struct SampleBorsh {
        x: i32,
        y: i32,
    }

    #[derive(Clone, Copy, Debug, PartialEq, Pod, Zeroable)]
    #[repr(C)]
    struct SampleBytemuck {
        a: u64,
        b: u64,
    }

    impl<'a> Value<'a> for SampleBorsh {
        const DISC: Disc = *b"SAMPLE_B";

        fn len(&self) -> u32 {
            (mem::size_of::<SampleBorsh>()) as u32
        }

        fn owned(&self) -> bool {
            true
        }

        fn bytes(&self) -> Vec<u8> {
            self.try_to_vec().unwrap()
        }

        fn bytes_ref(&self) -> &[u8] {
            &[]
        }

        fn consume(self) -> Vec<u8> {
            [
                Self::DISC.to_vec(),
                self.len().to_le_bytes().to_vec(),
                self.try_to_vec().unwrap(),
            ]
            .concat()
        }
    }

    impl<'a> Value<'a> for SampleBytemuck {
        const DISC: Disc = *b"SAMPLE_M";

        fn len(&self) -> u32 {
            (mem::size_of::<SampleBytemuck>()) as u32
        }

        fn owned(&self) -> bool {
            false
        }

        fn bytes(&self) -> Vec<u8> {
            bytemuck::bytes_of(self).to_vec()
        }

        fn bytes_ref(&self) -> &[u8] {
            bytemuck::bytes_of(self)
        }

        fn consume(self) -> Vec<u8> {
            [
                Self::DISC.to_vec(),
                self.len().to_le_bytes().to_vec(),
                bytemuck::bytes_of(&self).to_vec(),
            ]
            .concat()
        }
    }

    impl<'a> Held<'a> for SampleBorsh {
        type Inner = SampleBorsh;

        fn get(b: RefMut<'a, [u8]>) -> Result<Self::Inner, ValueError> {
            SampleBorsh::try_from_slice(&b).map_err(|_| ValueError::GenericError)
        }
    }

    impl<'a> Held<'a> for SampleBytemuck {
        type Inner = SampleBytemuck;

        fn get(b: RefMut<'a, [u8]>) -> Result<Self::Inner, ValueError> {
            let b: &[u8] = b.as_ref();
            bytemuck::try_pod_read_unaligned(b).map_err(|_| ValueError::GenericError)
        }
    }

    #[test]
    fn test_hydration() {
        let sample = SampleBorsh { x: 1, y: 2 };
        let sample2 = SampleBytemuck {
            a: 3423482374237498270,
            b: 4059680459866,
        };
        let sample_vec = sample.consume();
        println!("sample_vec len: {:?}", sample_vec);
        let sample2_vec = sample2.consume();
        println!("sample_vec2 len: {:?}", sample2_vec);
        let mut buffer = vec![0; 1 + sample_vec.len() + sample2_vec.len()];
        println!("buffer len: {}", buffer.len());
        let mut offset = 1;
        buffer[0] = 2;
        buffer[offset..offset + sample_vec.len()].copy_from_slice(&sample_vec);
        offset = offset + sample_vec.len();
        buffer[offset..offset + sample2_vec.len()].copy_from_slice(&sample2_vec);

        let h = Holder::new(Rc::new(RefCell::new(&mut buffer))).unwrap();
        let s = h.get::<SampleBorsh>().unwrap();
        println!("{:?}", s);
        assert_eq!(s.x, 1);
        assert_eq!(s.y, 2);

        let s = h.get::<SampleBytemuck>().unwrap();
        println!("{:?}", s);
        assert_eq!(s.a, 3423482374237498270);
        assert_eq!(s.b, 4059680459866);
    }

    #[test]
    fn insert_to_empty() {
        let mut buffer = vec![0; 100];
        let h = Holder::new(Rc::new(RefCell::new(&mut buffer))).unwrap();
        let sample = SampleBorsh { x: 1, y: 2 };
        h.insert(sample).unwrap();
        let sample2 = SampleBytemuck {
            a: 3423482374237498270,
            b: 4059680459866,
        };
        h.insert(sample2).unwrap();
        let s = h.get::<SampleBorsh>().unwrap();
        println!("{:?}", s);
        assert_eq!(s.x, 1);
        assert_eq!(s.y, 2);

        let s = h.get::<SampleBytemuck>().unwrap();
        println!("{:?}", s);
        assert_eq!(s.a, 3423482374237498270);
        assert_eq!(s.b, 4059680459866);
    }

    #[test]
    pub fn insert_to_empty_then_remove() {
        let mut buffer = vec![0; 100];
        let h = Holder::new(Rc::new(RefCell::new(&mut buffer))).unwrap();
        assert_eq!(h.free_space(), 99);
        let sample = SampleBorsh { x: 1, y: 2 };
        let sample2 = SampleBytemuck {
            a: 3423482374237498270,
            b: 4059680459866,
        };
        let sample_len = sample.len();
        let dec_len = 99 - (sample_len + 12 + sample2.len() + 12);

        h.insert(sample).unwrap();
        h.insert(sample2).unwrap();

        assert_eq!(h.free_space(), dec_len);

        let s = h.get::<SampleBorsh>().unwrap();
        println!("{:?}", s);
        assert_eq!(s.x, 1);
        assert_eq!(s.y, 2);

        let s = h.get::<SampleBytemuck>().unwrap();
        println!("{:?}", s);
        assert_eq!(s.a, 3423482374237498270);
        assert_eq!(s.b, 4059680459866);

        h.remove::<SampleBorsh>().unwrap();

        assert_eq!(h.free_space(), dec_len + sample_len + 12);

        let s = h.get::<SampleBorsh>();
        assert!(s.is_err());

        let s = h.get::<SampleBytemuck>().unwrap();
        println!("{:?}", s);
        assert_eq!(s.a, 3423482374237498270);
        assert_eq!(s.b, 4059680459866);
    }
}