hpl_toolkit/utils/
long_vec_map.rs

1use anchor_lang::prelude::*;
2use std::collections::HashMap;
3
4#[cfg(feature = "compression")]
5use crate::compression::*;
6
7#[cfg(feature = "schema")]
8use crate::schema::*;
9
10use super::{vec_map::KeyValue, Inner, VecMap};
11
12#[cfg_attr(feature = "debug", derive(Debug))]
13#[derive(AnchorSerialize, AnchorDeserialize, Clone, PartialEq)]
14pub struct LongVecMap<K: Inner, V: Inner> {
15    inner: Vec<KeyValue<K, V>>,
16}
17
18impl<K: Inner + PartialEq, V: Inner> LongVecMap<K, V> {
19    pub fn new() -> Self {
20        Self { inner: vec![] }
21    }
22
23    pub fn intoinner(self) -> Vec<KeyValue<K, V>> {
24        self.inner
25    }
26
27    pub fn insert(&mut self, key: K, value: V) {
28        for KeyValue { key: k, value: v } in &mut self.inner {
29            if *k == key {
30                *v = value;
31                return;
32            }
33        }
34        self.inner.push(KeyValue { key, value });
35    }
36
37    pub fn get(&self, key: &K) -> Option<&V> {
38        self.inner
39            .iter()
40            .find_map(|KeyValue { key: k, value: v }| if k == key { Some(v) } else { None })
41    }
42
43    pub fn get_mut(&mut self, key: &K) -> Option<&mut V> {
44        self.inner
45            .iter_mut()
46            .find_map(|KeyValue { key: k, value: v }| if k == key { Some(v) } else { None })
47    }
48
49    pub fn iter<'i>(&'i self) -> impl Iterator<Item = (&'i K, &'i V)> {
50        let a = self.inner.iter().map(|i| -> (&K, &V) { i.into() });
51        a
52    }
53
54    pub fn iter_mut<'i>(&'i mut self) -> impl Iterator<Item = (&'i K, &'i mut V)> {
55        let a = self.inner.iter_mut().map(|i| -> (&K, &mut V) { i.into() });
56
57        a
58    }
59
60    pub fn remove(&mut self, key: &K) {
61        self.inner.retain(|KeyValue { key: k, .. }| k != key);
62    }
63}
64// Custom iterator for LongVecMap
65pub struct LongVecMapIter<K: Inner + PartialEq, V: Inner> {
66    inner: std::vec::IntoIter<KeyValue<K, V>>,
67}
68
69impl<K: Inner + PartialEq, V: Inner> Iterator for LongVecMapIter<K, V>
70where
71    K: Clone,
72    V: Clone,
73{
74    type Item = (K, V);
75
76    fn next(&mut self) -> Option<Self::Item> {
77        self.inner.next().map(|kv| (kv.key, kv.value))
78    }
79}
80impl<K: Inner + PartialEq, V: Inner> FromIterator<(K, V)> for LongVecMap<K, V> {
81    fn from_iter<I: IntoIterator<Item = (K, V)>>(iter: I) -> Self {
82        let mut vec_map = LongVecMap::new();
83        for (key, value) in iter {
84            vec_map.inner.push(KeyValue { key, value });
85        }
86        vec_map
87    }
88}
89impl<K: Inner + PartialEq, V: Inner> IntoIterator for LongVecMap<K, V>
90where
91    K: Inner + PartialEq,
92    V: Inner,
93{
94    type Item = (K, V);
95    type IntoIter = LongVecMapIter<K, V>;
96
97    fn into_iter(self) -> Self::IntoIter {
98        LongVecMapIter {
99            inner: self.inner.into_iter(),
100        }
101    }
102}
103
104impl<K: Inner + PartialEq, V: Inner> FromIterator<KeyValue<K, V>> for LongVecMap<K, V> {
105    fn from_iter<T: IntoIterator<Item = KeyValue<K, V>>>(iter: T) -> Self {
106        Self {
107            inner: iter.into_iter().collect(),
108        }
109    }
110}
111
112impl<K: Inner + PartialEq, V: Inner> From<&HashMap<K, V>> for LongVecMap<K, V> {
113    fn from(map: &HashMap<K, V>) -> Self {
114        Self {
115            inner: map
116                .iter()
117                .map(|(k, v)| KeyValue {
118                    key: k.clone(),
119                    value: v.clone(),
120                })
121                .collect(),
122        }
123    }
124}
125
126impl<K: Inner + PartialEq, V: Inner> From<HashMap<K, V>> for LongVecMap<K, V> {
127    fn from(map: HashMap<K, V>) -> Self {
128        Self {
129            inner: map
130                .into_iter()
131                .map(|(k, v)| KeyValue { key: k, value: v })
132                .collect(),
133        }
134    }
135}
136
137impl<K: Inner + PartialEq, V: Inner + PartialEq> From<VecMap<K, V>> for LongVecMap<K, V> {
138    fn from(map: VecMap<K, V>) -> Self {
139        Self {
140            inner: map.intoinner().as_vec(),
141        }
142    }
143}
144
145impl<K: Inner + PartialEq, V: Inner + PartialEq> Into<VecMap<K, V>> for LongVecMap<K, V> {
146    fn into(self) -> VecMap<K, V> {
147        VecMap::from(self.inner)
148    }
149}
150
151#[cfg(feature = "compression")]
152impl<K: Inner + ToNode, V: Inner + ToNode> ToNode for LongVecMap<K, V> {
153    fn to_node(&self) -> [u8; 32] {
154        let mut seeds: Vec<[u8; 32]> = vec![];
155        for KeyValue { key, value } in &self.inner {
156            seeds.push(key.to_node());
157            seeds.push(value.to_node());
158        }
159        let seeds_refs: Vec<&[u8]> = seeds.iter().map(|node| &node[..]).collect();
160        anchor_lang::solana_program::keccak::hashv(&seeds_refs[..]).to_bytes()
161    }
162}
163
164#[cfg(feature = "schema")]
165impl<K: Inner + ToSchema + PartialEq, V: Inner + ToSchema> ToSchema for LongVecMap<K, V> {
166    fn schema() -> Schema {
167        Schema::VecMap(Box::new(K::schema()), Box::new(V::schema()))
168    }
169
170    fn schema_value(&self) -> SchemaValue {
171        let mut schema = LongVecMap::<SchemaValue, SchemaValue>::new();
172        self.inner.iter().for_each(|KeyValue { key, value }| {
173            schema.insert(key.schema_value(), value.schema_value());
174        });
175        SchemaValue::VecMap(schema.into())
176    }
177}