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::{short_vec::ShortVecIter, Inner, ShortVec};
11
12#[cfg_attr(feature = "debug", derive(Debug))]
13#[derive(AnchorSerialize, AnchorDeserialize, Clone, PartialEq)]
14pub struct KeyValue<K: Inner, V: Inner> {
15 pub key: K,
16 pub value: V,
17}
18
19impl<K: Inner, V: Inner> Into<(K, V)> for KeyValue<K, V> {
20 fn into(self) -> (K, V) {
21 (self.key, self.value)
22 }
23}
24impl<'i, K: Inner, V: Inner> Into<(&'i K, &'i V)> for &'i KeyValue<K, V> {
25 fn into(self) -> (&'i K, &'i V) {
26 (&self.key, &self.value)
27 }
28}
29impl<'i, K: Inner, V: Inner> Into<(&'i K, &'i mut V)> for &'i mut KeyValue<K, V> {
30 fn into(self) -> (&'i K, &'i mut V) {
31 (&self.key, &mut self.value)
32 }
33}
34impl<K: Inner, V: Inner> From<(&K, &V)> for KeyValue<K, V> {
35 fn from((key, value): (&K, &V)) -> Self {
36 KeyValue {
37 key: key.to_owned(),
38 value: value.to_owned(),
39 }
40 }
41}
42impl<K: Inner, V: Inner> From<(K, V)> for KeyValue<K, V> {
43 fn from((key, value): (K, V)) -> Self {
44 Self { key, value }
45 }
46}
47
48#[cfg_attr(feature = "debug", derive(Debug))]
49#[derive(AnchorSerialize, AnchorDeserialize, Clone, PartialEq)]
50pub struct VecMap<K: Inner + PartialEq, V: Inner + PartialEq> {
51 inner: ShortVec<KeyValue<K, V>>,
52}
53
54impl<K: Inner + PartialEq, V: Inner + PartialEq> VecMap<K, V> {
55 pub fn new() -> Self {
56 Self {
57 inner: ShortVec::new(),
58 }
59 }
60
61 pub fn intoinner(self) -> ShortVec<KeyValue<K, V>> {
62 self.inner
63 }
64
65 pub fn insert(&mut self, key: K, value: V) {
66 for KeyValue { key: k, value: v } in self.inner.inner_mut() {
67 if k == &key {
68 *v = value;
69 return;
70 }
71 }
72 self.inner.push(KeyValue { key, value });
73 }
74
75 pub fn get(&self, key: &K) -> Option<&V> {
76 self.inner
77 .inner()
78 .iter()
79 .find_map(|KeyValue { key: k, value: v }| if k == key { Some(v) } else { None })
80 }
81
82 pub fn get_mut(&mut self, key: &K) -> Option<&mut V> {
83 self.inner
84 .inner_mut()
85 .iter_mut()
86 .find_map(|KeyValue { key: k, value: v }| if k == key { Some(v) } else { None })
87 }
88
89 pub fn iter<'i>(&'i self) -> impl Iterator<Item = (&'i K, &'i V)> {
90 let a = self.inner.inner().iter().map(|i| -> (&K, &V) { i.into() });
91 a
92 }
93
94 pub fn iter_mut<'i>(&'i mut self) -> impl Iterator<Item = (&'i K, &'i mut V)> {
95 let a = self
96 .inner
97 .inner_mut()
98 .iter_mut()
99 .map(|i| -> (&K, &mut V) { i.into() });
100
101 a
102 }
103
104 pub fn remove(&mut self, key: &K) {
105 self.inner
106 .inner_mut()
107 .retain(|KeyValue { key: k, .. }| k != key);
108 }
109}
110
111pub struct VecMapIter<K: Inner + PartialEq, V: Inner + PartialEq> {
113 inner: ShortVecIter<KeyValue<K, V>>,
114}
115
116impl<K: Inner + PartialEq, V: Inner + PartialEq> Iterator for VecMapIter<K, V>
117where
118 K: Clone,
119 V: Clone,
120{
121 type Item = (K, V);
122
123 fn next(&mut self) -> Option<Self::Item> {
124 self.inner.next().map(|kv| (kv.key, kv.value))
125 }
126}
127impl<K: Inner + PartialEq, V: Inner + PartialEq> FromIterator<(K, V)> for VecMap<K, V> {
128 fn from_iter<I: IntoIterator<Item = (K, V)>>(iter: I) -> Self {
129 let mut vec_map = VecMap::new();
130 for (key, value) in iter {
131 vec_map.inner.push(KeyValue { key, value });
132 }
133 vec_map
134 }
135}
136impl<K: Inner + PartialEq, V: Inner + PartialEq> IntoIterator for VecMap<K, V>
137where
138 K: Inner + PartialEq,
139 V: Inner + PartialEq,
140{
141 type Item = (K, V);
142 type IntoIter = VecMapIter<K, V>;
143
144 fn into_iter(self) -> Self::IntoIter {
145 VecMapIter {
146 inner: self.inner.into_iter(),
147 }
148 }
149}
150
151impl<K: Inner + PartialEq, V: Inner + PartialEq> FromIterator<KeyValue<K, V>> for VecMap<K, V> {
152 fn from_iter<T: IntoIterator<Item = KeyValue<K, V>>>(iter: T) -> Self {
153 Self {
154 inner: iter.into_iter().collect(),
155 }
156 }
157}
158
159impl<K: Inner + PartialEq, V: Inner + PartialEq> From<&HashMap<K, V>> for VecMap<K, V> {
160 fn from(map: &HashMap<K, V>) -> Self {
161 Self {
162 inner: map
163 .iter()
164 .map(|(k, v)| KeyValue {
165 key: k.clone(),
166 value: v.clone(),
167 })
168 .collect(),
169 }
170 }
171}
172
173impl<K: Inner + PartialEq, V: Inner + PartialEq> From<HashMap<K, V>> for VecMap<K, V> {
174 fn from(map: HashMap<K, V>) -> Self {
175 Self {
176 inner: map
177 .into_iter()
178 .map(|(k, v)| KeyValue { key: k, value: v })
179 .collect(),
180 }
181 }
182}
183
184impl<K: Inner + PartialEq, V: Inner + PartialEq> From<Vec<KeyValue<K, V>>> for VecMap<K, V> {
185 fn from(map: Vec<KeyValue<K, V>>) -> Self {
186 Self { inner: map.into() }
187 }
188}
189
190#[cfg(feature = "compression")]
191impl<K: Inner + PartialEq + ToNode, V: Inner + PartialEq + ToNode> ToNode for VecMap<K, V> {
192 fn to_node(&self) -> [u8; 32] {
193 let mut seeds: Vec<[u8; 32]> = vec![];
194 for (key, value) in self.iter() {
195 seeds.push(key.to_node());
196 seeds.push(value.to_node());
197 }
198 let seeds_refs: Vec<&[u8]> = seeds.iter().map(|node| &node[..]).collect();
199 anchor_lang::solana_program::keccak::hashv(&seeds_refs[..]).to_bytes()
200 }
201}
202
203#[cfg(feature = "schema")]
204impl<K: Inner + PartialEq + ToSchema, V: Inner + PartialEq + ToSchema> ToSchema for VecMap<K, V> {
205 fn schema() -> Schema {
206 Schema::VecMap(Box::new(K::schema()), Box::new(V::schema()))
207 }
208
209 fn schema_value(&self) -> SchemaValue {
210 let mut schema = VecMap::<SchemaValue, SchemaValue>::new();
211 self.inner
212 .inner()
213 .iter()
214 .for_each(|KeyValue { key, value }| {
215 schema.insert(key.schema_value(), value.schema_value());
216 });
217 SchemaValue::VecMap(schema)
218 }
219}