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