yglnk-core 0.0.2

use crate::trunc_key_at0;
use xxhash_rust::xxh64::xxh64;

#[derive(Clone, Copy, Debug)]
pub struct Settings {
    pub seed: u32,
    pub entsize: u16,
    pub blshift: u8,
}

impl Settings {
    /// NOTE: the key must be truncated first using [`trunc_key_at0`]
    pub fn translate_key(&self, key: &[u8]) -> (u64, u64) {
        let h = xxh64(key, self.seed.into());
        let blmask = (1 << (h % 64)) | (1 << ((h >> self.blshift) % 64));
        (h, blmask)
    }
}

#[derive(Clone, Copy, Debug)]
pub struct Header {
    pub strtab_link: u32,
    pub nbuckets: u32,
    pub nchains: u32,
    pub nblf: u16,
    pub settings: Settings,
}

#[derive(Clone)]
pub struct Ref<'a> {
    strtab: &'a [u8],
    bloom: &'a [u8],
    buckets: &'a [u8],
    chains: &'a [u8],
    settings: Settings,
}

#[derive(Clone)]
pub struct Iter<'a> {
    strtab: &'a [u8],
    chains: core::slice::ChunksExact<'a, u8>,
}

#[derive(Clone, Copy, Debug, PartialEq)]
pub struct Value<R> {
    pub typ: u32,
    pub rest: R,
}

impl Settings {
    #[inline]
    pub const fn chain_entry_size(&self) -> usize {
        match (1 + self.entsize as usize).checked_mul(16) {
            Some(x) => x,
            None => unreachable!(),
        }
    }
}

impl Header {
    pub fn parse(data: &[u8]) -> Option<Self> {
        Some(Self::decode(data.get(0..20)?.try_into().unwrap()))
    }

    pub fn decode(data: [u8; 20]) -> Self {
        Self {
            strtab_link: u32::from_be_bytes(data[0..4].try_into().unwrap()),
            nbuckets: u32::from_be_bytes(data[4..8].try_into().unwrap()),
            nchains: u32::from_be_bytes(data[8..12].try_into().unwrap()),
            nblf: u16::from_be_bytes(data[14..16].try_into().unwrap()),
            settings: Settings {
                entsize: u16::from_be_bytes(data[12..14].try_into().unwrap()),
                blshift: data[16],
                // to make this easier, we first fold the blshift into the seed,
                // and immediately drop that after the conversion
                seed: u32::from_be_bytes(data[16..20].try_into().unwrap()) & 0xffffff,
            },
        }
    }

    pub fn encode(&self) -> [u8; 20] {
        let mut ret = [0u8; 20];
        ret[0..4].copy_from_slice(&u32::to_be_bytes(self.strtab_link));
        ret[4..8].copy_from_slice(&u32::to_be_bytes(self.nbuckets));
        ret[8..12].copy_from_slice(&u32::to_be_bytes(self.nchains));
        ret[12..14].copy_from_slice(&u16::to_be_bytes(self.settings.entsize));
        ret[14..16].copy_from_slice(&u16::to_be_bytes(self.nblf));
        // to make this easier, we first overwrite the blshift field with the seed,
        // then immediately fix it up
        assert_eq!(self.settings.seed & 0xff000000, 0);
        ret[16..20].copy_from_slice(&u32::to_be_bytes(self.settings.seed));
        ret[16] = self.settings.blshift;
        ret
    }

    pub fn tabsize(&self) -> usize {
        let tmp: usize = 5 + usize::try_from(self.nbuckets).unwrap() + 2 * usize::from(self.nblf);
        4 * tmp + self.settings.chain_entry_size() * usize::try_from(self.nchains).unwrap()
    }
}

// hash -> index conversion/transformation helper
pub fn hash_trf(h: u64, items: usize, div: usize) -> usize {
    div * usize::try_from(h % u64::try_from(items / div).unwrap()).unwrap()
}

impl<'a> Value<&'a [u8]> {
    pub fn parse(entry: &'a [u8]) -> Option<Self> {
        if entry.len() < 16 {
            return None;
        }
        Some(Self {
            typ: u32::from_be_bytes(entry[12..16].try_into().unwrap()),
            rest: &entry[16..],
        })
    }
}

impl<'a> Ref<'a> {
    /// `location` should be the offset where the hash table is present (in units of 16 bytes)
    pub fn parse(data: &'a [u8], location: u32) -> Option<Self> {
        let alldata = data;
        let uf = <usize as TryFrom<u32>>::try_from;
        let offset = crate::decode_location(location)?;
        let data = data.get(offset..)?;

        let header = Header::parse(data)?;
        let data = data.get(..header.tabsize())?;

        let bloom_end: usize = 20 + 8 * usize::from(header.nblf);
        let buckets_end = bloom_end + 4 * uf(header.nbuckets).unwrap();
        let chains_end =
            buckets_end + header.settings.chain_entry_size() * uf(header.nchains).unwrap();
        assert_eq!(data.len(), chains_end);

        Some(Ref {
            settings: header.settings,
            strtab: alldata.get(usize::try_from(header.strtab_link).ok()?..)?,
            bloom: &data[20..bloom_end],
            buckets: &data[bloom_end..buckets_end],
            chains: &data[buckets_end..chains_end],
        })
    }

    fn get_e_name<'s>(strtab: &'s [u8], sel: &[u8]) -> Option<&'s [u8]> {
        let e_name_ix = usize::try_from(u32::from_be_bytes(sel[8..12].try_into().unwrap())).ok()?;
        Some(trunc_key_at0(strtab.get(e_name_ix..)?))
    }

    /// NOTE: the key is truncated after the first null byte
    pub fn lookup(&self, key: &[u8]) -> Option<Value<&'a [u8]>> {
        let key = trunc_key_at0(key);
        let (h, blmask) = self.settings.translate_key(key);

        // check bloom filter
        let blsel = hash_trf(h / 64, self.bloom.len(), 8);
        let blword = u64::from_be_bytes(self.bloom[blsel..blsel + 8].try_into().unwrap());
        if (blword & blmask) != blmask {
            return None;
        }

        // retrieve bucket/chain start index
        let bkid = hash_trf(h, self.buckets.len(), 4);
        let chain_start = usize::try_from(u32::from_be_bytes(
            self.buckets[bkid..bkid + 4].try_into().unwrap(),
        ))
        .ok()?;

        for sel in self
            .chains
            .chunks_exact(self.settings.chain_entry_size())
            .skip(chain_start)
        {
            assert!(sel.len() >= 16);
            let e_hash = u64::from_be_bytes(sel[0..8].try_into().unwrap());
            if (h | 1) == (e_hash | 1) {
                let e_name = Self::get_e_name(self.strtab, sel)?;

                if e_name == key {
                    return Some(Value::parse(sel).unwrap());
                }
            }

            if (e_hash & 1) == 0 {
                break;
            }
        }

        None
    }

    pub fn iter(&self) -> Iter<'a> {
        Iter {
            strtab: self.strtab,
            chains: self.chains.chunks_exact(self.settings.chain_entry_size()),
        }
    }
}

impl<'a> Iterator for Iter<'a> {
    type Item = (u64, &'a [u8], Value<&'a [u8]>);

    fn next(&mut self) -> Option<Self::Item> {
        let i = self.chains.next()?;
        let e_name = Ref::get_e_name(self.strtab, i)?;
        Some((
            u64::from_be_bytes(i[0..8].try_into().unwrap()),
            e_name,
            Value::parse(i).unwrap(),
        ))
    }
}

#[derive(Copy, Clone, Debug, PartialEq)]
pub struct PreEntry<R> {
    /// index into the string table corresponding to the entry name
    pub name_ix: u32,
    /// type of the entry
    pub typ: u32,
    /// content of the entry
    pub rest: R,
}

#[cfg(feature = "alloc")]
use alloc::{boxed::Box, vec, vec::Vec};

#[cfg(feature = "alloc")]
pub fn serialize<I, R>(
    strtab: crate::StrtabDescriptorRef<'_>,
    nbuckets: u32,
    nblf: u16,
    settings: Settings,
    data: I,
) -> Option<Vec<u8>>
where
    I: Iterator<Item = PreEntry<R>>,
    R: Into<Box<[u8]>>,
{
    let mut header = Header {
        strtab_link: strtab.location,
        nbuckets,
        // chains get computed later
        nchains: 0,
        nblf,
        settings,
    };

    let mut bloom = vec![0u64; nblf.into()];
    let nbku: usize = nbuckets.try_into().ok()?;
    let (data_lb_count, _) = data.size_hint();
    if usize::try_from(u32::MAX)
        .map(|dub| data_lb_count > dub)
        .unwrap_or(false)
    {
        // too much data, overflow
        return None;
    }
    let mut chains = vec![Vec::<PreEntry<(u64, Box<[u8]>)>>::new(); nbku];
    let actual_entsize = settings.chain_entry_size();
    assert!(actual_entsize >= 16);

    for PreEntry { name_ix, typ, rest } in data {
        let name = &strtab[name_ix];
        let rest: Box<[u8]> = rest.into();
        assert!(actual_entsize >= (16 + rest.len()));
        let (h, blmask) = header.settings.translate_key(name);

        // add to bloom filter; (div=1 because we use u64 entries and convert later)
        let blid = hash_trf(h / 64, bloom.len(), 1);
        bloom[blid] |= blmask;

        // add to chain
        chains[hash_trf(h, nbku, 1)].push(PreEntry {
            name_ix,
            typ,
            rest: (h, rest),
        });
    }

    header.nchains = u32::try_from(chains.iter().map(|i| i.len()).sum::<usize>()).ok()?;
    let mut ret = vec![0u8; header.tabsize()];

    // copy header
    ret[0..20].copy_from_slice(&header.encode());

    // copy bloom filter
    let bloom_end: usize = 20 + 8 * usize::from(header.nblf);
    for (blin, blout) in bloom
        .into_iter()
        .zip(ret[20..bloom_end].chunks_exact_mut(8))
    {
        blout.copy_from_slice(&u64::to_be_bytes(blin));
    }

    // copy buckets and chains
    let (buckets_out, chains_out) =
        ret[bloom_end..].split_at_mut(4 * usize::try_from(header.nbuckets).unwrap());
    assert_eq!(
        chains_out.len(),
        actual_entsize * usize::try_from(header.nchains).unwrap()
    );
    let mut buckets_out = buckets_out.chunks_exact_mut(4);
    let mut chain_offset = 0u32;

    for i in chains {
        buckets_out
            .next()
            .unwrap()
            .copy_from_slice(&u32::to_be_bytes(chain_offset));

        if i.is_empty() {
            continue;
        }
        let ilenm1 = i.len() - 1;

        for (
            n,
            PreEntry {
                name_ix,
                typ,
                rest: (mut h, rest),
            },
        ) in i.into_iter().enumerate()
        {
            let actual_offset = usize::try_from(chain_offset)
                .unwrap()
                .checked_mul(actual_entsize)
                .unwrap();
            let entry = &mut chains_out[actual_offset..actual_offset + actual_entsize];
            chain_offset += 1;

            h |= 1;
            h = if n == ilenm1 {
                // last chain element gets its last hash bit set to 0
                h ^ 1
            } else {
                h
            };

            entry[0..8].copy_from_slice(&u64::to_be_bytes(h));
            entry[8..12].copy_from_slice(&u32::to_be_bytes(name_ix));
            entry[12..16].copy_from_slice(&u32::to_be_bytes(typ));
            if !rest.is_empty() {
                entry[16..16 + rest.len()].copy_from_slice(&rest[..]);
            }
        }
    }

    assert_eq!(buckets_out.next(), None);
    assert_eq!(chain_offset, header.nchains);

    Some(ret)
}

#[cfg(all(test, feature = "alloc"))]
mod tests {
    use super::{Settings, PreEntry};

    #[test]
    fn simple() {
        let strtab = crate::StrtabDescriptorRef {
            data: b"\x00a\x00b\x00c\x00d\x00",
            location: 0,
        };
        let sts = Settings {
            seed: 0xcafe,
            entsize: 0,
            blshift: 20,
        };
        use alloc::{boxed::Box, vec};
        let nullslc: Box<[u8]> = [].to_vec().into_boxed_slice();

        let blob = super::serialize(strtab, 2, 1, sts, [PreEntry {
            name_ix: 1,
            typ: 0,
            rest: nullslc.clone(),
        }, PreEntry {
            name_ix: 3,
            typ: 1,
            rest: nullslc.clone(),
        }, PreEntry {
            name_ix: 5,
            typ: 2,
            rest: nullslc.clone(),
        }, PreEntry {
            name_ix: 7,
            typ: 3,
            rest: nullslc,
        }].into_iter()).expect("unable to serialize hash table");

        let mut conc = vec![0u8; 16];
        conc[..strtab.data.len()].copy_from_slice(strtab.data);
        conc.extend_from_slice(&blob[..]);
        core::mem::drop(blob);

        let crf = super::Ref::parse(&conc, 1).expect("unable to parse hash table again");
        let elem_a = crf.lookup(b"a").expect("unable to retrieve 'b' element");
        assert_eq!(elem_a.typ, 0);
        let elem_b = crf.lookup(b"b").expect("unable to retrieve 'b' element");
        assert_eq!(elem_b.typ, 1);
        let elem_c = crf.lookup(b"c").expect("unable to retrieve 'b' element");
        assert_eq!(elem_c.typ, 2);
        let elem_d = crf.lookup(b"d").expect("unable to retrieve 'b' element");
        assert_eq!(elem_d.typ, 3);
        assert_eq!(crf.lookup(b"e"), None);
    }
}