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 fn total_element_count(&self) -> Option<usize> {
83 self.0.values().try_fold(0usize, |total, values| {
84 WORD_SIZE
85 .checked_add(values.len())
86 .and_then(|entry_elements| total.checked_add(entry_elements))
87 })
88 }
89
90 pub fn entry(&mut self, key: Word) -> Entry<'_, Word, Arc<[Felt]>> {
92 self.0.entry(key)
93 }
94
95 pub fn merge(&mut self, other: &Self) -> Result<(), (MapEntry, Arc<[Felt]>)> {
101 if let Some(conflict) = self.find_conflicting_entry(other) {
102 Err(conflict)
103 } else {
104 self.merge_new(other);
105 Ok(())
106 }
107 }
108
109 fn merge_new(&mut self, other: &Self) {
111 for (key, value) in other.iter() {
112 self.0.entry(*key).or_insert_with(|| value.clone());
113 }
114 }
115
116 fn find_conflicting_entry(&self, other: &Self) -> Option<(MapEntry, Arc<[Felt]>)> {
122 for (key, new_value) in other.iter() {
123 if let Some(existing_value) = self.get(key)
124 && existing_value != new_value
125 {
126 return Some(((*key, existing_value.clone()), new_value.clone()));
128 }
129 }
130 None
132 }
133}
134
135impl From<BTreeMap<Word, Arc<[Felt]>>> for AdviceMap {
136 fn from(value: BTreeMap<Word, Arc<[Felt]>>) -> Self {
137 Self(value)
138 }
139}
140
141impl From<BTreeMap<Word, Vec<Felt>>> for AdviceMap {
142 fn from(value: BTreeMap<Word, Vec<Felt>>) -> Self {
143 value.into_iter().collect()
144 }
145}
146
147impl IntoIterator for AdviceMap {
148 type Item = (Word, Arc<[Felt]>);
149 type IntoIter = IntoIter<Word, Arc<[Felt]>>;
150
151 fn into_iter(self) -> Self::IntoIter {
152 self.0.into_iter()
153 }
154}
155
156impl<V> FromIterator<(Word, V)> for AdviceMap
157where
158 V: Into<Arc<[Felt]>>,
159{
160 fn from_iter<I>(iter: I) -> Self
161 where
162 I: IntoIterator<Item = (Word, V)>,
163 {
164 iter.into_iter()
165 .map(|(key, value)| (key, value.into()))
166 .collect::<BTreeMap<Word, Arc<[Felt]>>>()
167 .into()
168 }
169}
170
171impl<V> Extend<(Word, V)> for AdviceMap
172where
173 V: Into<Arc<[Felt]>>,
174{
175 fn extend<I>(&mut self, iter: I)
176 where
177 I: IntoIterator<Item = (Word, V)>,
178 {
179 self.0.extend(iter.into_iter().map(|(key, value)| (key, value.into())))
180 }
181}
182
183impl Serializable for AdviceMap {
184 fn write_into<W: ByteWriter>(&self, target: &mut W) {
185 target.write_usize(self.0.len());
186 for (key, values) in self.0.iter() {
187 target.write((key, values.to_vec()));
188 }
189 }
190}
191
192impl Deserializable for AdviceMap {
193 fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
194 let mut map = BTreeMap::new();
195 let count = source.read_usize()?;
196 for _ in 0..count {
197 let (key, values): (Word, Vec<Felt>) = source.read()?;
198 map.insert(key, Arc::from(values));
199 }
200 Ok(Self(map))
201 }
202}
203
204#[cfg(test)]
205mod tests {
206 use super::*;
207
208 #[test]
209 fn test_advice_map_serialization() {
210 let mut map1 = AdviceMap::default();
211 map1.insert(Word::default(), vec![Felt::from_u32(1), Felt::from_u32(2)]);
212
213 let bytes = map1.to_bytes();
214
215 let map2 = AdviceMap::read_from_bytes(&bytes).unwrap();
216
217 assert_eq!(map1, map2);
218 }
219}