tablestg 0.4.13

use crate::*;
use std::hash::{Hash,Hasher};

/// Store Data, returned by [Store::get].
pub enum SData {
    Small ( PData, usize, usize ),
    Large ( LVec<u8> ),
}

impl SData
{
    /// Get a reference to the stored data slice.
    pub fn data(&mut self) -> &[u8]
    {
       match self
       {
           SData::Small( pdata, off, len ) => &pdata.data[*off..*off+*len],
           SData::Large( v ) => v,
       }
    }

    /// Must be called before SData drops, see [Store::get].
    pub fn drop(&mut self, ps: &mut PageSet)
    {
        match self
        {
           SData::Small( pdata, _, _ ) => ps.note( std::mem::take(pdata) ),
           SData::Large( _v ) => {},
       }
    }
}

/// Store is similar to [VBuckMap], but allows records of any length.
///
/// Small records ( < 255 bytes ) are fast to store and access, large records are slower.
#[derive(Debug)]
pub struct Store
{
   main: VBuckMapInfo,
   extra: VBuckMapInfo,
   extra_next_id: u64,
}

impl Store
{
    /// Start a new store.
    pub fn new( ps: &mut PageSet ) -> Self
    {
       let main = VBuckMap::new(1, ps).save();
       let extra = VBuckMap::new(1, ps).save();
       Self{ main, extra, extra_next_id: 0 }
    }

    /// Insert user data, must not be a duplicate key ( but this is not checked ).
    pub fn insert<K: VKey>(&mut self, key: &K, user_data: &[u8], ps: &mut PageSet) {
        let len = user_data.len();
        let mut x = LVec::<u8>::new();
        
        if len < 255
        {
            x.push( 1 ); // Small record.
            x.extend_from_slice( user_data );
        }
        else // Large record
        {
            let start_id = self.extra_next_id;

            // Store user_data as chunks.
            let mut done = 0;
            let mut todo = len;
            let mut em = VBuckMap::restore(self.extra, ps);
            let mut id = self.extra_next_id;
            
            while todo > 0 // Save user_data as chunks
            {
               let mut amount = 247; // Chunk size, considering id takes 8 bytes.
               if amount > todo { amount = todo; }

               let chunk = &user_data[done..done+amount];

               let mut t = LVec::new();
               t.extend_from_slice( &id.to_le_bytes() );
               t.extend_from_slice( chunk );

               let key = IdVKey{ id };
               em.insert( &key, &t);

               id += 1;
               todo -= amount;
               done += amount;
            }
            self.extra = em.save();
            self.extra_next_id = id;

            // Write code, start_id and length to x.
            let code : u8 = if len as u64 <= u16::MAX as u64 { 2 } 
               else if len as u64 <= u32::MAX as u64 { 3 }
               else {4};
               
            x.push( code );
            // Could save key hash here, makes rehash (and get) faster.
            x.extend_from_slice( &start_id.to_le_bytes() );
            match code {
               2 => x.extend_from_slice( &(len as u16).to_le_bytes() ),
               3 => x.extend_from_slice( &(len as u32).to_le_bytes() ),
               4 => x.extend_from_slice( &(len as u64).to_le_bytes() ),
               _ => panic!()
            }
        }
        let mut m = VBuckMap::restore(self.main, ps);
        let key = StoreKey{ key, store: self };
        m.insert( &key, &x );
        self.main = m.save();
    }

    /// Get data for specified key, returns SData, or None if key not found.
    /// [SData::drop] must be called on result SData before it is dropped.
    pub fn get<K: VKey>(&mut self, key: &K, ps: &mut PageSet) -> Option<SData> {
        let mut m = VBuckMap::restore(self.main, ps);
        let key = StoreKey{ key, store: self };
        if let Some( (pdata, off, len) ) = m.get( &key )
        {
            if pdata.data[off] == 1 {
               return Some(SData::Small(pdata,off+1,len-1));
            } else {
               let v = self.fetch_chunks( &pdata.data[off..off+len], ps);
               ps.note(pdata);
               return Some(SData::Large(v));
            }
        }
        None
    }

    /// Fetch chunks.
    fn fetch_chunks(&self, x: &[u8], ps: &mut PageSet) -> LVec<u8>
    {
        let code = x[0];
        let mut id = u64::from_le_bytes( x[1..9].try_into().unwrap() );
        let len = match code {
           2  => u16::from_le_bytes( x[9..11].try_into().unwrap() ) as usize,
           3  => u32::from_le_bytes( x[9..13].try_into().unwrap() ) as usize,
           4  => u64::from_le_bytes( x[9..17].try_into().unwrap() ) as usize,
           _ => panic!(),
        };

        let mut done = 0;
        let mut em = VBuckMap::restore(self.extra, ps);
        let mut result = LVec::new();
        while done < len
        {
            let key = IdVKey{ id };
            let (rdata, off, amt) = em.get( &key ).unwrap();
            let (off, amt) = (off+8, amt-8); // Skip the id.
            result.extend_from_slice( &rdata.data[off..off+amt] );
            em.ps.note( rdata );
            done += amt;
            id += 1;
        }
        result
    }
}

struct StoreKey<'a, K: VKey>
{
   key: &'a K,
   store: &'a Store,
}

impl <'a, K: VKey> VKey for StoreKey<'a,K>
{
    fn ok(&self, bytes: &[u8], ps: &mut PageSet) -> bool {
        if bytes[0] == 1 {
            self.key.ok( &bytes[1..], ps )
        } else {
            let v = self.store.fetch_chunks( bytes, ps );
            let result = self.key.ok( &v, ps );
            result
        }
    }
    fn rehash<H: Hasher>(&self, bytes: &[u8], h: &mut H, ps: &mut PageSet) {
        if bytes[0] == 1 {
            self.key.rehash( &bytes[1..], h, ps  );
        } else {
            let v = self.store.fetch_chunks( bytes, ps );
            self.key.rehash( &v, h, ps );
        }
    }
}

impl <'a, K: VKey> Hash for StoreKey<'a,K>
{
    fn hash<H>(&self, h: &mut H) where H: Hasher {
        self.key.hash( h );
    }
}

// ###################################### test test test test ################################

#[cfg(test)]
fn test_insert(m: &mut Store, dt: &DataType, v: &Value, id: u64, ps: &mut PageSet) {

    let td = dt.value_to_bytes( v );
    let mut x = Vec::new();
    x.extend_from_slice(&id.to_le_bytes());
    x.extend_from_slice(&td);

    // println!("adding {:?} id={}", &x, id );

    let key = IdVKey { id };
    m.insert(&key, &x, ps);

    // println!("checking id={} m={:?}", id, m);

    if let Some(mut sd) = m.get(&key, ps)
    {
       assert!( sd.data() == &x );
       sd.drop(ps);
    }
    else { panic!() }   
}

#[cfg(test)]
fn test_get(m: &mut Store, dt: &DataType, v: &Value, id: u64, ps: &mut PageSet) {

    let td = dt.value_to_bytes( v );
    let mut x = Vec::new();
    x.extend_from_slice(&id.to_le_bytes());
    x.extend_from_slice(&td);

    // println!("getting {:?} id={}", &x, id );

    let key = IdVKey { id };

    if let Some(mut sd) = m.get(&key, ps)
    {
       assert!( sd.data() == &x );
       sd.drop(ps);
    }
    else { panic!() }
    
}

#[cfg(test)]
fn cust_dt() -> DataType {
    use pstd::veca;

    DataType::Struct(veca![
        (LString::from("Name"), DataType::String),
        (LString::from("Email"), DataType::String),
        (LString::from("Postal"), DataType::String),
    ])
}

#[cfg(test)]
pub fn test_store(ps: &mut PageSet) {
    use pstd::veca;


    let mut m = Store::new(ps);

    let n = 256 * 256;
    // let n = 1_000_000;
    // let n = 100;

    println!("testing store n={}", n);

    let _big = vec![ b'a'; 256 ];

    let v = Value::List( veca![
            Value::String(LString::from("maz")),
            Value::String(LString::from("maz@gmail.com")),
            Value::String(LString::from(tos(&_big))),
            // Value::String(LString::from("33 Sandpiper")),
        ]);

    println!("test value={:?}", v);

    let dt = cust_dt();

    println!("testing insert store={:?}", m);

    for i in 0..n {
        test_insert(&mut m, &dt, &v, i, ps);
    }

    println!("testing get store={:?}", m);

    for i in 0..n {
        test_get(&mut m, &dt, &v, i, ps);
    }

/*
    println!("testing store went ok");

    let start = std::time::Instant::now();
    let mut iter = m.iter();
    let mut count = 0;
    while let Some(_r) = iter.next(ps) {
        // println!( "iterating... r = {}", tos(r) );
        count += 1;
    }

    assert_eq!(count, n);

    println!(
        "time to iterate over {} records = {} micro-sec",
        n,
        start.elapsed().as_micros()
    );

    println!("iter map={:?}", &iter.map);

    iter.drop(ps);
*/

    println!("testing store - everything is ok");
}