tablestg 0.4.22

/// Store is similar to [VBuckMap], but allows records of any length.
///
/// [`insert`], `get`, `remove` and `iter` give access to keyed records.
/// Small keyed records ( < 255 bytes ) are fast to store and access, large records are slower.
/// For example iterating over 64K small records could take 1 milli-sec, but 90 milli-sec for large records.
/// Insert performance is improved if key length is known and small ( [VKey::len] ).
///
/// Keyed records are mutable in the sense that they can be removed then re-inserted with the same key.
///
/// [`store`], `fetch` and `delete` give access to data where the key is an u64 id chosen by the Store.
/// This can be used to reduce the size of a keyed record by storing large parts separately.
///
/// [`insert`]: Store::insert
/// [`store`]: Store::store
#[derive(Debug, serde::Serialize, serde::Deserialize)]
pub struct Store {
    /// Stores keyed records.
    main: VBuckMapInfo,
    /// Stores records keyed by extra_next_id.
    extra: VBuckMapInfo,
    /// Next available id.
    extra_next_id: u64,
    /// Number of records in main.
    record_count: u64,
    /// Records number of main removals, but subsequent insertions reduce this.
    remove_balance: u64,
    /// Number of records in extra
    extra_record_count: u64,
    /// Changed
    changed: bool,
}

use crate::{
    DataType, IdVKey, LVec, LazyItem, MSPX, PData, PageSet, SPX, VBuckMap, VBuckMapInfo,
    VBuckMapIter, VKey, Value,
};
use std::hash::{Hash, Hasher};

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,
            record_count: 0,
            remove_balance: 0,
            extra_record_count: 0,
            changed: false,
        }
    }

    /// 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::with_capacity(256);

        if len < 255 {
            x.push(1); // Small record.
            x.extend_from_slice(user_data);
        } else
        // Large record
        {
            // Store user_data as chunks.
            // In future, could use a Page or a TreeVec for large sizes.
            // A reason for NOT allowing records larger than 255 (to be stored directly ) is it could cause premature bucket expansion.

            let mut done = 0;

            // Code is number of bytes required to store length of user_data (2, 4 or 8 ).
            let code: u8 = if len as u64 <= u16::MAX as u64 {
                2
            } else if len as u64 <= u32::MAX as u64 {
                4
            } else {
                8
            };

            // Add code, id and len to x.
            x.push(code);
            x.extend_from_slice(&self.extra_next_id.to_le_bytes());
            match code {
                2 => x.extend_from_slice(&(len as u16).to_le_bytes()),
                4 => x.extend_from_slice(&(len as u32).to_le_bytes()),
                8 => x.extend_from_slice(&(len as u64).to_le_bytes()),
                _ => panic!(),
            }

            if let Some(klen) = key.len() {
                assert!(klen <= len);
                if 1 + 8 + code as usize + klen < 256 {
                    // Store key bytes in x, makes rehash and ok more efficient.
                    x.extend_from_slice(&user_data[0..klen]);
                    assert!(x.len() < 256);
                    done += klen;
                }
            }
            self.store(&user_data[done..], ps);
        }
        let mut m = VBuckMap::restore(self.main, ps);
        let key = StoreKey { key, store: self };
        m.insert(&key, &x);
        self.main = m.save();
        self.record_count += 1;
        self.changed = true;
        if self.remove_balance > 0 {
            self.remove_balance -= 1;
        }
    }

    /// Get data for specified key, returns SData, or None if key not found.
    pub fn get<K: VKey>(&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.borrow().data[off] == 1 {
                return Some(SData::Small(pdata, off + 1, len - 1));
            } else {
                let v = self.fetch_chunks(&pdata.borrow().data[off..off + len], ps);
                return Some(SData::Large(v));
            }
        }
        None
    }

    /// Remove specified key.
    pub fn remove<K: VKey>(&mut self, key: &K, ps: &mut PageSet) -> bool {
        let mut m = VBuckMap::restore(self.main, ps);
        let got = {
            let key = StoreKey { key, store: self };
            m.get(&key)
        };
        if let Some((pdata, off, len)) = got {
            if pdata.borrow().data[off] != 1 {
                self.remove_chunks(&pdata.borrow().data[off..off + len], m.ps);
            }
            let key = StoreKey { key, store: self };
            m.remove(&key);
            self.main = m.save(); // Not necessary yet, but maybe one day remove might change root/buckets.
            self.record_count -= 1;
            self.remove_balance += 1;
            self.changed = true;
            true
        } else {
            false
        }
    }

    /// Get iterator that returns all records (rows).
    pub fn iter<'a>(&'a self, ps: &mut PageSet) -> StoreIter<'a> {
        let inner = VBuckMap::restore(self.main, ps).iter();
        StoreIter {
            inner,
            store: self,
            v: LVec::new(),
        }
    }

    /// Store arbitrary size data, returns id. Use fetch to get it back.
    pub fn store(&mut self, user_data: &[u8], ps: &mut PageSet) -> u64 {
        let result = self.extra_next_id;
        let mut id = result;
        let mut em = VBuckMap::restore(self.extra, ps);
        let mut done = 0;
        let mut todo = user_data.len();
        while todo > 0 {
            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::with_capacity(8 + chunk.len());
            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_next_id = id;
        self.extra = em.save();
        self.extra_record_count += 1;
        self.changed = true;
        result
    }

    /// Fetch stored data, len must not exceed len of stored user_data.
    pub fn fetch(&self, id: u64, len: usize, ps: &mut PageSet) -> LVec<u8> {
        let mut result = LVec::with_capacity(len);
        self.do_fetch_chunks(id, len, &mut result, ps);
        result
    }

    /// Delete stored data, len must be equal to original len of stored user_data.
    pub fn delete(&mut self, mut id: u64, len: usize, ps: &mut PageSet) {
        let mut em = VBuckMap::restore(self.extra, ps);
        let mut done = 0;
        while done < len {
            let key = IdVKey { id };
            let amt = em.remove(&key);
            assert!(amt > 0);
            done += amt;
            id += 1;
        }
        self.extra = em.save(); // Not necessary currently.
        self.extra_record_count -= 1;
        self.changed = true;
    }

    /// Delete everything, Store is no longer usable.
    pub fn delete_all(&mut self, ps: &mut PageSet) {
        VBuckMap::restore(self.main, ps).delete_all();
        VBuckMap::restore(self.extra, ps).delete_all();
    }

    pub fn changed(&self) -> bool {
        self.changed
    }

    /// Removes chunks using info in x.
    fn remove_chunks(&mut self, x: &[u8], ps: &mut PageSet) {
        let (id, len, local) = self.parse_x(x);
        self.delete(id, len - local, ps);
    }

    /// Returns chunk start id, length of user data and non-chunk length.
    fn parse_x(&self, x: &[u8]) -> (u64, usize, usize) {
        let code = x[0];
        let 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,
            4 => u32::from_le_bytes(x[9..13].try_into().unwrap()) as usize,
            8 => u64::from_le_bytes(x[9..17].try_into().unwrap()) as usize,
            _ => panic!(),
        };
        let over = 1 + 8 + code as usize; // Number of non user-data bytes stored in x.
        let local = x.len() - over; // Number of user_data bytes stored in x.
        (id, len, local)
    }

    /// Fetch all chunk data.
    fn fetch_chunks(&self, x: &[u8], ps: &mut PageSet) -> LVec<u8> {
        let (_, len, _) = self.parse_x(x);
        self.chunks(x, len, ps)
    }

    /// Fetch some or all bytes of chunk data.
    fn some_chunks(&self, x: &[u8], len: Option<usize>, ps: &mut PageSet) -> LVec<u8> {
        if let Some(len) = len {
            self.chunks(x, len, ps)
        } else {
            self.fetch_chunks(x, ps)
        }
    }

    /// Fetch len bytes of chunk data ( including any data stored in x ).
    fn chunks(&self, x: &[u8], len: usize, ps: &mut PageSet) -> LVec<u8> {
        let mut done = 0;
        let mut result = LVec::with_capacity(len);

        let code = x[0];
        let over = 1 + 8 + code as usize; // Number of non user-data bytes stored in x.
        let local = x.len() - over; // Number of user_data bytes stored in x.
        if local > 0 {
            result.extend_from_slice(&x[over..over + local]);
            done += local;
        }

        if done < len {
            let id = u64::from_le_bytes(x[1..9].try_into().unwrap());
            self.do_fetch_chunks(id, len, &mut result, ps);
        }
        result
    }

    fn do_fetch_chunks(&self, mut id: u64, len: usize, to: &mut LVec<u8>, ps: &mut PageSet) {
        let mut em = VBuckMap::restore(self.extra, ps);
        let mut done = 0;
        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.
            to.extend_from_slice(&rdata.borrow().data[off..off + amt]);
            done += amt;
            id += 1;
        }
    }
}

/// Value Iterator - result of [Store::iter] returns all records (rows).
pub struct StoreIter<'a> {
    inner: VBuckMapIter,
    pub store: &'a Store,
    v: LVec<u8>,
}

impl<'a> StoreIter<'a> {
    /// Get reference to next row.
    pub fn next(&mut self, ps: &mut PageSet) -> Option<&[u8]> {
        if let Some(data) = self.inner.next(ps) {
            let result = if data[0] == 1 {
                &data[1..]
            } else {
                self.v = self.store.fetch_chunks(data, ps);
                &self.v
            };
            Some(result)
        } else {
            None
        }
    }
}

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.some_chunks(bytes, self.key.len(), ps);
            self.key.ok(&v, ps)
        }
    }
    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.some_chunks(bytes, self.key.len(), 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);
    }
}

/// Store Data, returned by [Store::get].
pub enum SData {
    /// Data is a single section of PData. PData, off, len. Maybe len is not needed.
    Small(PData, usize, usize),
    /// Data put together from Extra.
    Large(LVec<u8>),
}

impl SData {
    /// Decode data to Value.
    pub fn decode(&self, dt: &DataType, spx: &mut SPX) -> Value {
        self.decode_at(dt, 0, spx)
    }

    /// Decode data at offset to Value.
    pub fn decode_at(&self, dt: &DataType, at: usize, spx: &mut SPX) -> Value {
        match self {
            SData::Small(pdata, off, _len) => {
                let off = at + *off;
                let buf = &pdata.borrow().data[off..];
                dt.bytes_to_value(buf, spx)
            }
            SData::Large(v) => dt.bytes_to_value(&v[at..], spx),
        }
    }

    /// Decode data, deleting any indirectly stored values.
    pub fn decode_del(&self, dt: &DataType, spx: &mut MSPX) -> Value {
        match self {
            SData::Small(pdata, off, _len) => {
                let buf = &pdata.borrow().data[*off..];
                dt.bytes_to_value_del(buf, spx)
            }
            SData::Large(v) => dt.bytes_to_value_del(v, spx),
        }
    }

    /// Get lazy row items.
    pub fn lazy_row_items(&self, dt: &DataType) -> LVec<LazyItem> {
        match self {
            SData::Small(pdata, off, _len) => {
                let data = &pdata.borrow().data[*off..];
                let mut ix = 0;
                dt.lazy_row_items(data, &mut ix)
            }
            SData::Large(v) => {
                let mut ix = 0;
                dt.lazy_row_items(v, &mut ix)
            }
        }
    }
}

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

#[cfg(test)]
fn test_insert(m: &mut Store, td: &[u8], id: u64, ps: &mut PageSet) {
    let mut x = LVec::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(_sd) = m.get(&key, ps) {
        // assert_eq!(sd.data(), &*x);
    } else {
        panic!()
    }
}

#[cfg(test)]
fn test_get(m: &mut Store, td: &[u8], id: u64, ps: &mut PageSet) {
    let mut x = LVec::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(_sd) = m.get(&key, ps) {
        // assert!(sd.data() == &*x);
    } else {
        panic!()
    }
}

#[cfg(test)]
pub fn test_store(ps: &mut PageSet) {
    let mut m = Store::new(ps);

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

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

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

    let td = b"Hello George";

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

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

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

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

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

    println!("testing store iter");
    let start = std::time::Instant::now();

    for _ in 0..10 {
        let mut iter = m.iter(ps);
        while let Some(_r) = iter.next(ps) {}
    }
    println!(
        "Time to iterate over {} rows 10 times = {} micro-sec",
        n,
        start.elapsed().as_micros()
    );

    println!("testing remove");

    for id in 0..n {
        let key = IdVKey { id };
        m.remove(&key, ps);
    }

    println!("testing store - everything is ok, m={:?}", m);
}