hpl-toolkit 0.0.5

use anchor_lang::prelude::*;
use std::collections::HashMap;

#[cfg(feature = "compression")]
use crate::compression::*;

#[cfg(feature = "schema")]
use crate::schema::*;

use super::{vec_map::KeyValue, Inner, VecMap};

#[cfg_attr(feature = "debug", derive(Debug))]
#[derive(AnchorSerialize, AnchorDeserialize, Clone, PartialEq)]
pub struct LongVecMap<K: Inner, V: Inner> {
    inner: Vec<KeyValue<K, V>>,
}

impl<K: Inner + PartialEq, V: Inner> LongVecMap<K, V> {
    pub fn new() -> Self {
        Self { inner: vec![] }
    }

    pub fn intoinner(self) -> Vec<KeyValue<K, V>> {
        self.inner
    }

    pub fn insert(&mut self, key: K, value: V) {
        for KeyValue { key: k, value: v } in &mut self.inner {
            if *k == key {
                *v = value;
                return;
            }
        }
        self.inner.push(KeyValue { key, value });
    }

    pub fn get(&self, key: &K) -> Option<&V> {
        self.inner
            .iter()
            .find_map(|KeyValue { key: k, value: v }| if k == key { Some(v) } else { None })
    }

    pub fn get_mut(&mut self, key: &K) -> Option<&mut V> {
        self.inner
            .iter_mut()
            .find_map(|KeyValue { key: k, value: v }| if k == key { Some(v) } else { None })
    }

    pub fn iter<'i>(&'i self) -> impl Iterator<Item = (&'i K, &'i V)> {
        let a = self.inner.iter().map(|i| -> (&K, &V) { i.into() });
        a
    }

    pub fn iter_mut<'i>(&'i mut self) -> impl Iterator<Item = (&'i K, &'i mut V)> {
        let a = self.inner.iter_mut().map(|i| -> (&K, &mut V) { i.into() });

        a
    }

    pub fn remove(&mut self, key: &K) {
        self.inner.retain(|KeyValue { key: k, .. }| k != key);
    }
}
// Custom iterator for LongVecMap
pub struct LongVecMapIter<K: Inner + PartialEq, V: Inner> {
    inner: std::vec::IntoIter<KeyValue<K, V>>,
}

impl<K: Inner + PartialEq, V: Inner> Iterator for LongVecMapIter<K, V>
where
    K: Clone,
    V: Clone,
{
    type Item = (K, V);

    fn next(&mut self) -> Option<Self::Item> {
        self.inner.next().map(|kv| (kv.key, kv.value))
    }
}
impl<K: Inner + PartialEq, V: Inner> FromIterator<(K, V)> for LongVecMap<K, V> {
    fn from_iter<I: IntoIterator<Item = (K, V)>>(iter: I) -> Self {
        let mut vec_map = LongVecMap::new();
        for (key, value) in iter {
            vec_map.inner.push(KeyValue { key, value });
        }
        vec_map
    }
}
impl<K: Inner + PartialEq, V: Inner> IntoIterator for LongVecMap<K, V>
where
    K: Inner + PartialEq,
    V: Inner,
{
    type Item = (K, V);
    type IntoIter = LongVecMapIter<K, V>;

    fn into_iter(self) -> Self::IntoIter {
        LongVecMapIter {
            inner: self.inner.into_iter(),
        }
    }
}

impl<K: Inner + PartialEq, V: Inner> FromIterator<KeyValue<K, V>> for LongVecMap<K, V> {
    fn from_iter<T: IntoIterator<Item = KeyValue<K, V>>>(iter: T) -> Self {
        Self {
            inner: iter.into_iter().collect(),
        }
    }
}

impl<K: Inner + PartialEq, V: Inner> From<&HashMap<K, V>> for LongVecMap<K, V> {
    fn from(map: &HashMap<K, V>) -> Self {
        Self {
            inner: map
                .iter()
                .map(|(k, v)| KeyValue {
                    key: k.clone(),
                    value: v.clone(),
                })
                .collect(),
        }
    }
}

impl<K: Inner + PartialEq, V: Inner> From<HashMap<K, V>> for LongVecMap<K, V> {
    fn from(map: HashMap<K, V>) -> Self {
        Self {
            inner: map
                .into_iter()
                .map(|(k, v)| KeyValue { key: k, value: v })
                .collect(),
        }
    }
}

impl<K: Inner + PartialEq, V: Inner + PartialEq> From<VecMap<K, V>> for LongVecMap<K, V> {
    fn from(map: VecMap<K, V>) -> Self {
        Self {
            inner: map.intoinner().as_vec(),
        }
    }
}

impl<K: Inner + PartialEq, V: Inner + PartialEq> Into<VecMap<K, V>> for LongVecMap<K, V> {
    fn into(self) -> VecMap<K, V> {
        VecMap::from(self.inner)
    }
}

#[cfg(feature = "compression")]
impl<K: Inner + ToNode, V: Inner + ToNode> ToNode for LongVecMap<K, V> {
    fn to_node(&self) -> [u8; 32] {
        let mut seeds: Vec<[u8; 32]> = vec![];
        for KeyValue { key, value } in &self.inner {
            seeds.push(key.to_node());
            seeds.push(value.to_node());
        }
        let seeds_refs: Vec<&[u8]> = seeds.iter().map(|node| &node[..]).collect();
        anchor_lang::solana_program::keccak::hashv(&seeds_refs[..]).to_bytes()
    }
}

#[cfg(feature = "schema")]
impl<K: Inner + ToSchema + PartialEq, V: Inner + ToSchema> ToSchema for LongVecMap<K, V> {
    fn schema() -> Schema {
        Schema::VecMap(Box::new(K::schema()), Box::new(V::schema()))
    }

    fn schema_value(&self) -> SchemaValue {
        let mut schema = LongVecMap::<SchemaValue, SchemaValue>::new();
        self.inner.iter().for_each(|KeyValue { key, value }| {
            schema.insert(key.schema_value(), value.schema_value());
        });
        SchemaValue::VecMap(schema.into())
    }
}