1use alloc::{
2 collections::{
3 BTreeMap,
4 btree_map::{Entry, IntoIter},
5 },
6 sync::Arc,
7 vec::Vec,
8};
9
10use crate::{
11 Felt, WORD_SIZE, Word,
12 crypto::hash::Poseidon2,
13 serde::{ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable},
14};
15
16#[derive(Debug, Clone, Default, PartialEq, Eq)]
29#[cfg_attr(
30 all(feature = "arbitrary", test),
31 miden_test_serialization_macros::serialization_test
32)]
33pub struct AdviceMap(BTreeMap<Word, Arc<[Felt]>>);
34
35type MapEntry = (Word, Arc<[Felt]>);
37
38impl AdviceMap {
39 pub fn get(&self, key: &Word) -> Option<&Arc<[Felt]>> {
41 self.0.get(key)
42 }
43
44 pub fn contains_key(&self, key: &Word) -> bool {
46 self.0.contains_key(key)
47 }
48
49 pub fn insert(&mut self, key: Word, value: impl Into<Arc<[Felt]>>) -> Option<Arc<[Felt]>> {
51 self.0.insert(key, value.into())
52 }
53
54 pub fn remove(&mut self, key: &Word) -> Option<Arc<[Felt]>> {
56 self.0.remove(key)
57 }
58
59 pub fn iter(&self) -> impl Iterator<Item = (&Word, &Arc<[Felt]>)> {
61 self.0.iter()
62 }
63
64 pub fn len(&self) -> usize {
66 self.0.len()
67 }
68
69 pub fn is_empty(&self) -> bool {
71 self.0.is_empty()
72 }
73
74 pub fn total_element_count(&self) -> Option<usize> {
79 self.0.values().try_fold(0usize, |total, values| {
80 WORD_SIZE
81 .checked_add(values.len())
82 .and_then(|entry_elements| total.checked_add(entry_elements))
83 })
84 }
85
86 pub fn commitment(&self) -> Word {
93 let entry_commitments = self
94 .iter()
95 .map(|(key, values)| {
96 let mut elements = Vec::with_capacity(WORD_SIZE + values.len());
97 elements.extend_from_slice(key.as_elements());
98 elements.extend_from_slice(values);
99 Poseidon2::hash_elements(&elements)
100 })
101 .collect::<Vec<_>>();
102
103 Poseidon2::merge_many(&entry_commitments)
104 }
105
106 pub fn entry(&mut self, key: Word) -> Entry<'_, Word, Arc<[Felt]>> {
108 self.0.entry(key)
109 }
110
111 pub fn merge(&mut self, other: &Self) -> Result<(), (MapEntry, Arc<[Felt]>)> {
117 if let Some(conflict) = self.find_conflicting_entry(other) {
118 Err(conflict)
119 } else {
120 self.merge_new(other);
121 Ok(())
122 }
123 }
124
125 fn merge_new(&mut self, other: &Self) {
127 for (key, value) in other.iter() {
128 self.0.entry(*key).or_insert_with(|| value.clone());
129 }
130 }
131
132 fn find_conflicting_entry(&self, other: &Self) -> Option<(MapEntry, Arc<[Felt]>)> {
138 for (key, new_value) in other.iter() {
139 if let Some(existing_value) = self.get(key)
140 && existing_value != new_value
141 {
142 return Some(((*key, existing_value.clone()), new_value.clone()));
144 }
145 }
146 None
148 }
149}
150
151impl From<BTreeMap<Word, Arc<[Felt]>>> for AdviceMap {
152 fn from(value: BTreeMap<Word, Arc<[Felt]>>) -> Self {
153 Self(value)
154 }
155}
156
157impl From<BTreeMap<Word, Vec<Felt>>> for AdviceMap {
158 fn from(value: BTreeMap<Word, Vec<Felt>>) -> Self {
159 value.into_iter().collect()
160 }
161}
162
163impl IntoIterator for AdviceMap {
164 type Item = (Word, Arc<[Felt]>);
165 type IntoIter = IntoIter<Word, Arc<[Felt]>>;
166
167 fn into_iter(self) -> Self::IntoIter {
168 self.0.into_iter()
169 }
170}
171
172impl<V> FromIterator<(Word, V)> for AdviceMap
173where
174 V: Into<Arc<[Felt]>>,
175{
176 fn from_iter<I>(iter: I) -> Self
177 where
178 I: IntoIterator<Item = (Word, V)>,
179 {
180 iter.into_iter()
181 .map(|(key, value)| (key, value.into()))
182 .collect::<BTreeMap<Word, Arc<[Felt]>>>()
183 .into()
184 }
185}
186
187impl<V> Extend<(Word, V)> for AdviceMap
188where
189 V: Into<Arc<[Felt]>>,
190{
191 fn extend<I>(&mut self, iter: I)
192 where
193 I: IntoIterator<Item = (Word, V)>,
194 {
195 self.0.extend(iter.into_iter().map(|(key, value)| (key, value.into())))
196 }
197}
198
199impl Serializable for AdviceMap {
200 fn write_into<W: ByteWriter>(&self, target: &mut W) {
201 target.write_usize(self.0.len());
202 for (key, values) in self.0.iter() {
203 target.write((key, values.to_vec()));
204 }
205 }
206}
207
208impl Deserializable for AdviceMap {
209 fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
210 let mut map = BTreeMap::new();
211 let count = source.read_usize()?;
212 for _ in 0..count {
213 let (key, values): (Word, Vec<Felt>) = source.read()?;
214 map.insert(key, Arc::from(values));
215 }
216 Ok(Self(map))
217 }
218}
219
220#[cfg(test)]
221mod tests {
222 use super::*;
223
224 #[test]
225 fn test_advice_map_serialization() {
226 let mut map1 = AdviceMap::default();
227 map1.insert(Word::default(), vec![Felt::from_u32(1), Felt::from_u32(2)]);
228
229 let bytes = map1.to_bytes();
230
231 let map2 = AdviceMap::read_from_bytes(&bytes).unwrap();
232
233 assert_eq!(map1, map2);
234 }
235}