attempt 4

trash4.rs

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

use arrayref::array_ref;
pub use mpl_token_metadata::state::{
    Collection, CollectionDetails, Data, Key, Metadata, ProgrammableConfig, TokenStandard, Uses,
};
use solana_program::{account_info::AccountInfo, program_error::ProgramError};
pub use spl_token::state::Mint;
pub use state_interface_derive::*;

pub mod metadata;
pub mod mint;

use bytemuck::{Pod, Zeroable};
pub use metadata::{MetadataInterface, MetadataInterfacePack};
pub use mint::{MintInterface, MintInterfacePack};

#[derive(Debug)]
pub enum SolanaInterfaceError {
    SolanaEroror,
    AccountDataError,
    BlobNotFound,
    BorrowError,
    BlobAlreadyExists,
    BlobTooLarge,
    SerializationError,
}

impl From<ProgramError> for SolanaInterfaceError {
    fn from(e: ProgramError) -> Self {
        match e {
            ProgramError::AccountBorrowFailed => return SolanaInterfaceError::AccountDataError,
            _ => SolanaInterfaceError::SolanaEroror,
        }
    }
}
const ACCOUNT_HEADERSIZE: usize = 1;
#[derive(Copy, Clone, Pod, Zeroable)]
#[repr(C)]
pub struct InterfaceAccountHeader {
    number_of_blobs: u8,
}

#[derive(Copy, Clone, PartialEq, Eq)]
pub enum Format {
    Unknown,
    Borsh,
    Bytes,
}

impl From<&u8> for Format {
    fn from(b: &u8) -> Self {
        match b {
            0 => Format::Borsh,
            1 => Format::Bytes,
            _ => Format::Unknown,
        }
    }
}
impl Into<u8> for Format {
    fn into(self) -> u8 {
        match self {
            Format::Borsh => 0,
            Format::Bytes => 1,
            _ => 255,
        }
    }
}

pub trait Blob {
    type ReturnType;
    fn discriminator(&self) -> &'static [u8; 8];
    fn into_bytes(self) -> Result<Vec<u8>, SolanaInterfaceError>;
    fn from_bytes(bytes: &mut [u8]) -> Result<Self::ReturnType, SolanaInterfaceError>;
}

pub struct BlobStore<'a> {
    _ref: Rc<RefCell<&'a mut [u8]>>,
    _vec: Vec<([u8; 8], u32, u32)>,
    _imap: HashMap<[u8; 8], usize>,
}

impl<'a> BlobStore<'a> {
    pub fn new(refm: Rc<RefCell<&'a mut [u8]>>) -> Self {
        let mut b = BlobStore {
            _ref: refm,
            _vec: Vec::new(),
            _imap: HashMap::new(),
        };
        b.hydrate_blobs();
        b
    }

    fn hydrate_blobs(&mut self) -> Result<(), SolanaInterfaceError> {
        let buffer = &mut self._ref.borrow_mut();
        if buffer.len() == 0 {
            return Ok(());
        }
        let (header, tail) = buffer.split_at_mut(ACCOUNT_HEADERSIZE);
        let header = bytemuck::from_bytes::<InterfaceAccountHeader>(header);
        let blobs = header.number_of_blobs;
        let mut data = tail;
        for _ in 0..blobs {
            // [4 bytes lentgth][8 bytes discriminator][1 byte format][N bytes data]
            let (length, tail) = data.split_at_mut(4);
            let (disc_bytes, tail) = tail.split_at_mut(8);

            let disc = array_ref![disc_bytes, 0, 8].clone();
            let data_length = u32::from_le_bytes(*array_ref![length, 0, 4]);
            self._vec.push((disc, data_length, 0));
            self._imap.insert(disc, self._vec.len() - 1);

            data = tail;
        }
        Ok(())
    }

    pub fn from_account_info(ai: &'a AccountInfo<'a>) -> Self {
        Self::new(ai.data.clone())
    }

    pub fn insert(&mut self, blob: impl Blob) -> Result<(), SolanaInterfaceError> {
        let buffer = &mut self._ref.borrow_mut();
        let disc = blob.discriminator();
        let entry = self._imap.get(disc).and_then(|i| self._vec.get(*i));
        match entry {
            Some(_) => Err(SolanaInterfaceError::BlobAlreadyExists),
            None => {
                let (mut offset, size) = self
                    ._vec
                    .iter()
                    .fold((0, 0), |(oo, os), (d, o, s)| (oo + o, os + s));
                offset += ACCOUNT_HEADERSIZE as u32;
                let serialized_data = blob.into_bytes()?;
                let offset = self._ref.borrow().len() as u32;
                let length = serialized_data.len() as u32;
                if size + length > buffer.len() as u32 {
                    return Err(SolanaInterfaceError::BlobTooLarge);
                }
                self._vec.push((*disc, offset, length));
                self._imap.insert(*disc, self._vec.len() - 1);
                let mut final_offset = offset;
                buffer[final_offset as usize..final_offset as usize + 8].copy_from_slice(disc);
                final_offset += 8;
                buffer[final_offset as usize..final_offset as usize + 4].copy_from_slice(disc);
                final_offset += 4;
                buffer[final_offset as usize..final_offset as usize + length as usize]
                    .copy_from_slice(&serialized_data);
                Ok(())
            }
        }
    }

    pub fn contains(&self, disc: &[u8; 8]) -> bool {
        self._imap.contains_key(disc)
    }

    pub fn remove(&mut self, disc: &[u8; 8]) -> Result<(), SolanaInterfaceError> {
        let index = match self._imap.get(disc) {
            Some(index) => *index,
            None => return Err(SolanaInterfaceError::BlobNotFound),
        };

        let (_, removed_offset, removed_size) = self._vec.swap_remove(index);
        self._imap.remove(disc);
        if let Some((last_disc, _, _)) = self._vec.get(index) {
            self._imap.insert(*last_disc, index);
        }

        // Shift data in the buffer to fill the removed gap
        let buffer = &mut self._ref.borrow_mut();
        let remaining_data = &mut buffer[(removed_offset as usize + removed_size as usize)..];
        remaining_data.copy_within(removed_size as usize.., removed_offset as usize);
        Ok(())
    }

    pub fn get_mut<T: Blob>(
        &'a mut self,
        disc: &[u8; 8],
    ) -> Result<RefMut<T>, SolanaInterfaceError> {
        let entry = self._imap.get(disc).and_then(|i| self._vec.get(*i));
        let buffer = self._ref.borrow_mut();

        match entry {
            Some((_, offset, size)) => {
                let mut blob =
                    RefMut::map(buffer, |slice| &mut slice[*offset as usize..*size as usize]);
                let ret = RefMut::map(blob, |mut slice| T::from_bytes(slice).unwrap());
                Ok(ret)
            }
            None => Err(SolanaInterfaceError::BlobNotFound),
        }
    }

    pub fn get<T: Blob>(&self, disc: &[u8; 8]) -> Result<Ref<T>, SolanaInterfaceError> {
        let entry = self._imap.get(disc).and_then(|i| self._vec.get(*i));
        let buffer = self._ref.borrow();

        match entry {
            Some((_, offset, size)) => {
                let blob = Ref::map(buffer, |slice| &slice[*offset as usize..*size as usize]);

                let ret = Ref::map(blob, |slice| T::from_bytes(slice).unwrap());
                Ok(ret)
            }
            None => Err(SolanaInterfaceError::BlobNotFound),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::cell::RefCell;
    use std::rc::Rc;

    use borsh::{BorshDeserialize, BorshSerialize};
    use bytemuck::{Pod, Zeroable};

    // Struct that uses Borsh
    #[derive(BorshSerialize, BorshDeserialize, Debug, PartialEq, Eq, Clone)]
    pub struct BorshBlob {
        pub field1: u32,
        pub field2: String,
    }

    impl Blob for BorshBlob {
        type ReturnType = BorshBlob;
        fn discriminator(&self) -> &'static [u8; 8] {
            b"BORSHBLO"
        }

        fn into_bytes(self) -> Result<Vec<u8>, SolanaInterfaceError> {
            self.try_to_vec()
                .map_err(|_| SolanaInterfaceError::SerializationError)
        }

        fn from_bytes(bytes: &mut [u8]) -> Result<Self::ReturnType, SolanaInterfaceError> {
            BorshBlob::try_from_slice(bytes).map_err(|e| SolanaInterfaceError::BlobTooLarge)
        }
    }

    // Struct that uses bytemuck
    #[derive(Pod, Zeroable, Debug, PartialEq, Clone, Copy)]
    #[repr(C)]
    pub struct BytemuckBlob {
        pub field1: u64,
        pub field2: u64,
    }

    impl Blob for BytemuckBlob {
        fn discriminator(&self) -> &'static [u8; 8] {
            b"BYTEMUCK"
        }

        fn into_bytes(self) -> Result<Vec<u8>, SolanaInterfaceError> {
            Ok(bytemuck::bytes_of(&self).to_vec())
        }
    }

    #[test]
    fn test_borsh_blob() {
        let mut data = vec![0; 128];
        let refm = Rc::new(RefCell::new(&mut data[..]));
        let mut store = BlobStore::new(refm.clone());

        let blob = BorshBlob {
            field1: 42,
            field2: "Test".to_string(),
        };

        store.insert(blob.clone());

        let read_blob = store.get::<BorshBlob>(blob.discriminator()).unwrap();
        assert_eq!(*read_blob, blob);

        store.remove(blob.discriminator()).unwrap();
        assert!(!store.contains(blob.discriminator()));
    }

    #[test]
    fn test_bytemuck_blob() {
        let mut data = vec![0; 128];
        let refm = Rc::new(RefCell::new(&mut data[..]));
        let mut store = BlobStore::new(refm.clone());

        let blob = BytemuckBlob {
            field1: 42,
            field2: 123456789,
        };

        store.insert(blob).unwrap();

        let read_blob = store.get::<BytemuckBlob>(blob.discriminator()).unwrap();
        assert_eq!(*read_blob, blob);

        store.remove(blob.discriminator()).unwrap();
        assert!(!store.contains(blob.discriminator()));
    }

    #[test]
    fn test_store_with_multiple_blobs() {
        let mut data = vec![0; 1024];
        let refm = Rc::new(RefCell::new(&mut data[..]));
        let mut store = BlobStore::new(refm.clone());

        let borsh_blob = BorshBlob {
            field1: 42,
            field2: "Test".to_string(),
        };

        let bytemuck_blob = BytemuckBlob {
            field1: 42,
            field2: 123456789,
        };

        store.insert(borsh_blob.clone()).unwrap();
        store.insert(bytemuck_blob).unwrap();

        let read_borsh_blob = store.get::<BorshBlob>(borsh_blob.discriminator()).unwrap();
        let read_bytemuck_blob = store
            .get::<BytemuckBlob>(bytemuck_blob.discriminator())
            .unwrap();
        assert_eq!(*read_borsh_blob, borsh_blob);
        assert_eq!(*read_bytemuck_blob, bytemuck_blob);

        store.remove(borsh_blob.discriminator()).unwrap();
        store.remove(bytemuck_blob.discriminator()).unwrap();

        assert!(!store.contains(borsh_blob.discriminator()));
        assert!(!store.contains(bytemuck_blob.discriminator()));
    }
}