tablestg 0.4.15

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

/// Key for [VBuckMap].
#[allow(clippy::len_without_is_empty)]
pub trait VKey: Hash {
    /// Does key match specified bytes?
    fn ok(&self, bytes: &[u8], ps: &mut PageSet) -> bool;

    /// Compute hash from record bytes ( used when number of buckets is increased ).
    fn rehash<H: Hasher>(&self, bytes: &[u8], h: &mut H, ps: &mut PageSet);

    /// Number of bytes of user data needed for ok or rehash, if known.
    fn len(&self) -> Option<usize> {
        None
    }
}

/// Key for records where key is 64 bits, the first 8 bytes of the record.
#[derive(Hash)]
pub struct IdVKey {
    pub id: u64,
}

impl VKey for IdVKey {
    fn ok(&self, bytes: &[u8], _ps: &mut PageSet) -> bool {
        let loc = &bytes[0..8];
        let id = u64::from_le_bytes(loc.try_into().unwrap());
        self.id == id
    }
    fn rehash<H: Hasher>(&self, bytes: &[u8], h: &mut H, _ps: &mut PageSet) {
        let loc = &bytes[0..8];
        let id = u64::from_le_bytes(loc.try_into().unwrap());
        h.write_u64(id)
    }

    fn len(&self) -> Option<usize> {
        Some(8)
    }
}

/// VBuckMap root and buckets, returned by [VBuckMap::save].
#[derive(
    Copy,
    Clone,
    Debug,
    Default,
    Hash,
    PartialEq,
    Eq,
    PartialOrd,
    Ord,
    serde::Serialize,
    serde::Deserialize,
)]
pub struct VBuckMapInfo {
    root: u64,
    buckets: u64,
}

/// Hash Map implemented as list of buckets, stores small variable size rows.
pub struct VBuckMap<'a> {
    /// Number of buckets.
    buckets: u64,
    /// Root page.
    root: u64,
    /// Pageset for fetching/saving pages.
    pub ps: &'a mut PageSet,
    /// Has root / buckets changed?
    new_root: bool,
}

impl<'a> VBuckMap<'a> {
    /// Start a new map with specified number of buckets ( which must be > 0 ).
    pub fn new(buckets: u64, ps: &'a mut PageSet) -> Self {
        debug_assert!(buckets > 0);
        let mut pt = PageTree::new(ps.new_page(), 1, ps);
        pt.resize(buckets);
        Self {
            buckets,
            root: pt.root,
            ps,
            new_root: true,
        }
    }

    /// Insert user data, must not be a duplicate key ( but this is not checked ).
    /// Length of user data must be less than 256.
    pub fn insert<K: VKey>(&mut self, key: &K, user_data: &[u8]) {
        debug_assert!(user_data.len() < 256);
        self.do_insert(user_data, Self::hash(key), key)
    }

    /// Get data for specified key, returns PData, offset and length, or None if key not found.
    /// PData must be noted before it is dropped.
    pub fn get<K: VKey>(&mut self, key: &K) -> Option<(PData, usize, usize)> {
        let hash = Self::hash(key);
        let pnum = self.get_page_num_from_hash(hash, false);
        if pnum == 0 {
            None
        } else {
            let pdata = self.ps.load(pnum);
            if let Some((off, len)) = vbucket::Reader::new(&pdata.data).get(key, hash, self.ps) {
                Some((pdata, off, len))
            } else {
                self.ps.note(pdata);
                None
            }
        }
    }

    /// Remove a key. Returns the amount of user data deleted.
    pub fn remove<K: VKey>(&mut self, key: &K) -> usize {
        let hash = Self::hash(key);
        let pnum = self.get_page_num_from_hash(hash, false);
        if pnum != 0 {
            let mut pdata = self.ps.load(pnum);
            let md = pdata.make_mut();
            if md.is_empty() {
                let size = self.ps.compute_size(1000); // Start page size
                md.resize(size, 0);
            }
            let mut w = vbucket::Writer::new(md);
            let result = w.remove(key, hash, self.ps);
            if result > 0 {
                pdata.changed = true;
            }
            self.ps.note(pdata);
            return result;
        }
        0
    }

    /// Get iterator that returns all records (rows). Note that [VBuckMapIter::drop] must be called before VBuckMapIter drops.
    pub fn iter(&mut self) -> VBuckMapIter {
        let pnum = self.get_page_num_from_pix(0, false);
        let pdata = self.ps.load(pnum);
        VBuckMapIter {
            pix: 0,
            pdata,
            map: self.save(),
            pos: vbucket::Pos::start(),
        }
    }

    /// Has root and buckets changed ( so needs to be saved )?
    pub fn root_changed(&self) -> bool {
        self.new_root
    }

    /// Get the root and number of buckets. These can change on any insert.
    pub fn save(&self) -> VBuckMapInfo {
        VBuckMapInfo {
            root: self.root,
            buckets: self.buckets,
        }
    }

    /// Restore from saved root and buckets.
    pub fn restore(info: VBuckMapInfo, ps: &'a mut PageSet) -> Self {
        Self {
            root: info.root,
            buckets: info.buckets,
            ps,
            new_root: false,
        }
    }

    /// Delete everything. Map is no longer usable.
    pub fn delete(&mut self) {
        let mut pt = PageTree::new(self.root, self.buckets, self.ps);
        pt.drop_pages();
        self.root = 0;
    }

    /// Calculate hash
    fn rehash<K: VKey>(key: &K, user_data: &[u8], ps: &mut PageSet) -> u64 {
        let mut h = fxhash::FxHasher::default();
        key.rehash(user_data, &mut h, ps);
        h.finish()
    }

    /// Calculate hash
    fn hash<K: VKey>(key: &K) -> u64 {
        let mut h = fxhash::FxHasher::default();
        key.hash(&mut h);
        h.finish()
    }

    /// Get page number for hash.
    fn get_page_num_from_hash(&mut self, hash: u64, create: bool) -> u64 {
        let pix = hash % self.buckets;
        self.get_page_num_from_pix(pix, create)
    }

    /// Get page number for specified index (pix).
    fn get_page_num_from_pix(&mut self, pix: u64, create: bool) -> u64 {
        debug_assert!(pix < self.buckets);
        let mut pt = PageTree::new(self.root, self.buckets, self.ps);
        pt.get(pix, create)
    }

    /// Insert user_data.
    fn do_insert<K: VKey>(&mut self, user_data: &[u8], hash: u64, key: &K) {
        while self.try_insert(user_data, hash).is_err() {
            self.expand(key);
        }
    }

    /// Attempt to insert user_data. If this fails due to page size limit being reached returns Err.
    fn try_insert(&mut self, user_data: &[u8], hash: u64) -> Result<(), ()> {
        let pnum = self.get_page_num_from_hash(hash, true);
        let mut pdata = self.ps.load(pnum);

        pdata.changed = true;

        let mut md = pdata.make_mut();
        if md.is_empty() {
            let size = self.ps.compute_size(1000); // Start page size
            md.resize(size, 0);
        }
        assert!(md.len() >= 1000);
        let mut w = vbucket::Writer::new(md);

        if w.entry_full() {
            self.ps.note(pdata);
            return Err(());
        }

        let space = w.space(user_data.len()); // Extra space needed.

        if space > 0 {
            let old = vbucket::Reader::new(md);
            pdata.data = Arc::new(old.rebuild(md.len()));
            md = pdata.make_mut();
            w = vbucket::Writer::new(md);
            let space = w.space(user_data.len());

            if space > 0
            // Simple compact didn't work, need a bigger page size.
            {
                let old = vbucket::Reader::new(md);
                let new_size = self.ps.compute_size(md.len() + space);
                if new_size == 0 {
                    self.ps.note(pdata);
                    return Err(()); // Page size limit reached, expand number of buckets.
                }
                pdata.data = Arc::new(old.rebuild(new_size));
                md = pdata.make_mut();
                w = vbucket::Writer::new(md);
            }
        }
        w.insert(user_data, hash);
        self.ps.note(pdata);
        Ok(())
    }

    /// Increase the number of buckets, due to page size limit being reached for some bucket.
    fn expand<K: VKey>(&mut self, key: &K) {
        let buckets = 1 + self.buckets * 2;

        println!("expand buckets={} new buckets={}", self.buckets, buckets);

        // We cannot mutably borrow ps multiple times, so use save().
        let mut new = VBuckMap::new(buckets, self.ps).save();

        let mut iter = VBuckMap::restore(self.save(), self.ps).iter();
        while let Some(r) = iter.next(self.ps) {
            let h = Self::rehash(key, r, self.ps);
            let mut m = VBuckMap::restore(new, self.ps);
            m.do_insert(r, h, key);
            new = m.save();
        }
        iter.drop(self.ps);

        self.delete(); // Delete the old pages, use pages from new map.
        self.root = new.root;
        self.buckets = new.buckets;
        self.new_root = true;
    }
}

/// Iterator - returns all records (rows). Note that drop must be called before VBuckMapIter drops.
pub struct VBuckMapIter {
    pix: u64,
    pdata: PData,
    map: VBuckMapInfo,
    pos: vbucket::Pos,
}

impl VBuckMapIter {
    /// Get reference to next record.
    pub fn next(&mut self, ps: &mut PageSet) -> Option<&[u8]> {
        if let Some((off, len)) = self.off_and_len(ps) {
            Some(&self.pdata.data[off..off + len])
        } else {
            None
        }
    }

    /// Get mutable reference to next record. If the record is changed, changed must be called.
    /// This does not allow the record length to be changed. For those updates, record ids should
    /// first be collected, then the records to be updated should be removed and re-inserted in the map.
    pub fn next_mut(&mut self, ps: &mut PageSet) -> Option<&mut [u8]> {
        if let Some((off, len)) = self.off_and_len(ps) {
            let md = self.pdata.make_mut();
            Some(&mut md[off..off + len])
        } else {
            None
        }
    }

    /// Mark current page as changed ( see [Self::next_mut] ).
    pub fn changed(&mut self) {
        self.pdata.changed = true;
    }

    /// Drop the iterator ( cannot use Drop trait as this needs ps reference ).
    pub fn drop(&mut self, ps: &mut PageSet) {
        let pdata = std::mem::take(&mut self.pdata);
        ps.note(pdata);
    }

    /// Get offset and length of next record.
    fn off_and_len(&mut self, ps: &mut PageSet) -> Option<(usize, usize)> {
        loop {
            if let Some(x) = self.look() {
                return Some(x);
            } else if self.pix + 1 == self.map.buckets {
                return None;
            } else {
                self.pix += 1;
                ps.note(std::mem::take(&mut self.pdata));
                let mut m = VBuckMap::restore(self.map, ps);
                let pnum = m.get_page_num_from_pix(self.pix, false);
                self.pdata = ps.load(pnum);
                self.pos = vbucket::Pos::start();
            }
        }
    }

    /// Look in current pdata for offset and length of next record.
    fn look(&mut self) -> Option<(usize, usize)> {
        vbucket::Reader::new(&self.pdata.data).iter_next(&mut self.pos)
    }
}

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

#[cfg(test)]
fn test_insert(m: &mut VBuckMap, dt: &DataType, v: &Value, id: u64) {
    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 {}", &x );

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

    let (d, off, len) = m.get(&key).unwrap();

    assert_eq!(&d.data[off..off + len], x);
    m.ps.note(d);
}

#[cfg(test)]
fn test_get(m: &mut VBuckMap, dt: &DataType, v: &Value, id: u64) {
    let td = dt.value_to_bytes(v);

    let key = IdVKey { id };

    let (d, off, len) = m.get(&key).unwrap();

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

    assert_eq!(&d.data[off..off + len], x);
    m.ps.note(d);
}

#[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_vbuckmap(ps: &mut PageSet) {
    use pstd::veca;

    println!("testing vbuckmap");

    let save = VBuckMap::new(1, ps).save();

    // Test iterating over empty map.
    let mut iter = VBuckMap::restore(save, ps).iter();
    while let Some(_r) = iter.next(ps) {
        panic!()
    }
    iter.drop(ps);

    let mut m = VBuckMap::restore(save, ps);

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

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

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

    let dt = cust_dt();

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

    for i in 0..n {
        test_get(&mut m, &dt, &v, i);
    }
    println!("testing vbuckmap test_get 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);
    }

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

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