tablestg 0.4.15

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.
/// 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] ).
#[derive(Debug)]
pub struct Store {
    main: VBuckMapInfo,
    extra: VBuckMapInfo,
    extra_next_id: u64,
    record_count: u64,
    delete_balance: u64, // Records number of deletions, but subsequent insertions reduce this.
}

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,
            delete_balance: 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
        {
            // Store user_data as chunks.
            let mut done = 0;
            let mut id = self.extra_next_id;

            // 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(&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;
                }
            }

            let mut em = VBuckMap::restore(self.extra, ps);
            let mut todo = len - done;
            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;
        }
        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;
        if self.delete_balance > 0 { self.delete_balance -= 1; }
    }

    /// 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
    }

    /// 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.data[off] != 1 {
                self.remove_chunks(&pdata.data[off..off + len], m.ps);
            }
            m.ps.note(pdata);
            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.delete_balance += 1;
            true
        } else {
            false
        }
    }

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

    /// Removes chunks using info in x.
    fn remove_chunks(&mut self, x: &[u8], ps: &mut PageSet) {
        let (mut id, len, local) = self.parse_x(x);
        let mut em = VBuckMap::restore(self.extra, ps);
        let mut done = 0;
        while done < len - local {
            let key = IdVKey { id };
            let amt = em.remove(&key);
            assert!(amt > 0);
            done += amt;
            id += 1;
        }
        self.extra = em.save(); // Not necessary currently.
    }

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

        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 mut id = u64::from_le_bytes(x[1..9].try_into().unwrap());
            let mut em = VBuckMap::restore(self.extra, ps);
            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
    }
}

/// Iterator - result of [Store::iter] returns all records (rows). Note that drop must be called before StoreIter drops.
pub struct StoreIter<'a> {
    inner: VBuckMapIter,
    store: &'a Store,
    v: LVec<u8>,
}

impl<'a> StoreIter<'a> {
    /// Get reference to next record.
    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
        }
    }

    /// Drop the iterator ( cannot use Drop trait as this needs ps reference ).
    pub fn drop(&mut self, ps: &mut PageSet) {
        self.inner.drop(ps);
    }
}

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);
    }
}

// ###################################### 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_eq!(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;
    // let n = 1;

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

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

    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(ps);
    let mut count = 0;
    while let Some(_r) = iter.next(ps) {
        // println!( "iterating... r = {}", tos(r) );
        count += 1;
    }
    iter.drop(ps);

    assert_eq!(count, n);

    println!(
        "time to iterate over {} records = {} 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);
}