compactly 0.1.8

use super::{Encode, EncodingStrategy, LowCardinality};
use std::{collections::HashMap, hash::Hash, sync::Arc};

#[cfg(test)]
use expect_test::expect;

#[derive(Clone)]
pub struct CacheContext<T: Encode + Clone + Hash + PartialEq + Eq> {
    cached: HashMap<T, usize>,
    cache: Vec<T>,
    is_cached: <bool as Encode>::Context,
    context: T::Context,
    index: <usize as Encode>::Context,
}

impl<T: Encode + Clone + Hash + PartialEq + Eq> Default for CacheContext<T> {
    #[inline]
    fn default() -> Self {
        Self {
            cached: HashMap::new(),
            cache: Vec::new(),
            is_cached: Default::default(),
            context: Default::default(),
            index: Default::default(),
        }
    }
}

macro_rules! impl_low_cardinality {
    ($t:ty, $mod:ident) => {
        mod $mod {
            use super::{CacheContext, Encode, EncodingStrategy, LowCardinality};
            impl EncodingStrategy<$t> for LowCardinality {
                type Context = CacheContext<$t>;
                #[inline]
                fn encode<W: std::io::Write>(
                    value: &$t,
                    writer: &mut super::super::Writer<W>,
                    ctx: &mut Self::Context,
                ) -> Result<(), std::io::Error> {
                    let looked_up = ctx.cached.get(value).copied();
                    looked_up.is_some().encode(writer, &mut ctx.is_cached)?;
                    if let Some(idx) = looked_up {
                        idx.encode(writer, &mut ctx.index)
                    } else {
                        ctx.cached.insert(value.clone(), ctx.cached.len());
                        value.encode(writer, &mut ctx.context)
                    }
                }
                #[inline]
                fn decode<R: std::io::Read>(
                    reader: &mut super::super::Reader<R>,
                    ctx: &mut Self::Context,
                ) -> Result<$t, std::io::Error> {
                    let is_cached = bool::decode(reader, &mut ctx.is_cached)?;
                    if is_cached {
                        let idx = usize::decode(reader, &mut ctx.index)?;
                        ctx.cache
                            .get(idx)
                            .cloned()
                            .ok_or_else(|| std::io::Error::other("bad low_cardinality index"))
                    } else {
                        let value = <$t>::decode(reader, &mut ctx.context)?;
                        ctx.cache.push(value.clone());
                        Ok(value)
                    }
                }
            }
        }
    };
}

impl_low_cardinality!(String, string);
impl_low_cardinality!(Vec<u8>, bytes);
impl_low_cardinality!(u64, mod_u64);

// `Arc<str>` needs its own context: the encode-side HashMap uses `Arc<str>` as
// keys (deduplication by content) and the decode-side cache holds `Arc<str>`
// values so that a cache hit is a cheap refcount bump rather than a fresh
// allocation. This is the recommended alternative to `LowCardinality<String>`,
// which clones (reallocates) the String on every repeated value.
#[derive(Default, Clone)]
pub struct ArcStrCacheContext {
    cached: HashMap<Arc<str>, usize>,
    cache: Vec<Arc<str>>,
    is_cached: <bool as Encode>::Context,
    string_ctx: <String as Encode>::Context,
    index: <usize as Encode>::Context,
}

impl EncodingStrategy<Arc<str>> for LowCardinality {
    type Context = ArcStrCacheContext;
    #[inline]
    fn encode<W: std::io::Write>(
        value: &Arc<str>,
        writer: &mut super::Writer<W>,
        ctx: &mut Self::Context,
    ) -> Result<(), std::io::Error> {
        let looked_up = ctx.cached.get(value.as_ref()).copied();
        looked_up.is_some().encode(writer, &mut ctx.is_cached)?;
        if let Some(idx) = looked_up {
            idx.encode(writer, &mut ctx.index)
        } else {
            ctx.cached.insert(value.clone(), ctx.cached.len());
            value.to_string().encode(writer, &mut ctx.string_ctx)
        }
    }
    #[inline]
    fn decode<R: std::io::Read>(
        reader: &mut super::Reader<R>,
        ctx: &mut Self::Context,
    ) -> Result<Arc<str>, std::io::Error> {
        let is_cached = bool::decode(reader, &mut ctx.is_cached)?;
        if is_cached {
            let idx = usize::decode(reader, &mut ctx.index)?;
            ctx.cache
                .get(idx)
                .cloned()
                .ok_or_else(|| std::io::Error::other("bad low_cardinality index"))
        } else {
            let s = String::decode(reader, &mut ctx.string_ctx)?;
            let value: Arc<str> = Arc::from(s.as_str());
            ctx.cache.push(value.clone());
            Ok(value)
        }
    }
}

impl<const N: usize, T: Encode + Clone + Hash + Eq> EncodingStrategy<[T; N]> for LowCardinality {
    type Context = CacheContext<[T; N]>;
    fn encode<W: std::io::Write>(
        value: &[T; N],
        writer: &mut super::Writer<W>,
        ctx: &mut Self::Context,
    ) -> Result<(), std::io::Error> {
        let looked_up = ctx.cached.get(value).copied();
        looked_up.is_some().encode(writer, &mut ctx.is_cached)?;
        if let Some(idx) = looked_up {
            idx.encode(writer, &mut ctx.index)
        } else {
            ctx.cached.insert(value.clone(), ctx.cached.len());
            value.encode(writer, &mut ctx.context)
        }
    }
    fn decode<R: std::io::Read>(
        reader: &mut super::Reader<R>,
        ctx: &mut Self::Context,
    ) -> Result<[T; N], std::io::Error> {
        let is_cached = bool::decode(reader, &mut ctx.is_cached)?;
        if is_cached {
            let idx = usize::decode(reader, &mut ctx.index)?;
            ctx.cache
                .get(idx)
                .cloned()
                .ok_or_else(|| std::io::Error::other("bad low_cardinality index"))
        } else {
            let value = <[T; N]>::decode(reader, &mut ctx.context)?;
            ctx.cache.push(value.clone());
            Ok(value)
        }
    }
}

impl<T> EncodingStrategy<Vec<T>> for LowCardinality
where
    T: Encode,
    LowCardinality: EncodingStrategy<T>,
{
    type Context = (
        <usize as Encode>::Context,
        <LowCardinality as EncodingStrategy<T>>::Context,
    );
    fn encode<W: std::io::Write>(
        value: &Vec<T>,
        writer: &mut super::Writer<W>,
        ctx: &mut Self::Context,
    ) -> Result<(), std::io::Error> {
        value.len().encode(writer, &mut ctx.0)?;
        for v in value {
            LowCardinality::encode(v, writer, &mut ctx.1)?;
        }
        Ok(())
    }
    fn decode<R: std::io::Read>(
        reader: &mut super::Reader<R>,
        ctx: &mut Self::Context,
    ) -> Result<Vec<T>, std::io::Error> {
        let n = usize::decode(reader, &mut ctx.0)?;
        let mut x = Vec::with_capacity(n);
        for _ in 0..n {
            x.push(LowCardinality::decode(reader, &mut ctx.1)?);
        }
        Ok(x)
    }
}

#[test]
fn low_cardinality() {
    use super::encoded_bits;
    use crate::Encoded;

    let strings = [
        b"hello world, this is the very first string".to_vec(),
        b"This is a second string, which is like unto the first, and yet quite different".to_vec(),
    ];
    let mut v = Vec::new();
    for i in 0..1024 {
        v.push(if i % 3 == 0 {
            strings[0].clone()
        } else {
            strings[1].clone()
        });
    }
    let low = v
        .iter()
        .cloned()
        .map(Encoded::<_, LowCardinality>::new)
        .collect::<Vec<_>>();

    expect!["284470"].assert_eq(&encoded_bits!(v.clone()));
    expect!["1677"].assert_eq(&encoded_bits!(low.clone()));
    expect!["613"].assert_eq(&encoded_bits!(strings.clone().to_vec()));
    expect!["615"].assert_eq(&encoded_bits!(strings
        .iter()
        .cloned()
        .map(Encoded::<_, LowCardinality>::new)
        .collect::<Vec<_>>()));
}