refactor: add AdviceMap wrapper (0xPolygonMiden#1207)

miden/src/cli/data.rs

@@ -8,7 +8,7 @@ use miden::{
 };
 use serde_derive::{Deserialize, Serialize};
 use std::{
-    collections::HashMap,
+    collections::{BTreeMap, HashMap},
     fs,
     io::Write,
     path::{Path, PathBuf},
@@ -157,7 +157,7 @@ impl InputFile {
     }
 
     /// Parse advice map data from the input file.
-    fn parse_advice_map(&self) -> Result<Option<HashMap<[u8; 32], Vec<Felt>>>, String> {
+    fn parse_advice_map(&self) -> Result<Option<BTreeMap<RpoDigest, Vec<Felt>>>, String> {
         let advice_map = match &self.advice_map {
             Some(advice_map) => advice_map,
             None => return Ok(None),
@@ -166,9 +166,8 @@ impl InputFile {
         let map = advice_map
             .iter()
             .map(|(k, v)| {
-                // decode hex key
-                let mut key = [0u8; 32];
-                hex::decode_to_slice(k, &mut key)
+                // Convert key to RpoDigest
+                let mut key = RpoDigest::try_from(k)
                     .map_err(|e| format!("failed to decode advice map key `{k}` - {e}"))?;
 
                 // convert values to Felt
@@ -182,7 +181,7 @@ impl InputFile {
                     .collect::<Result<Vec<_>, _>>()?;
                 Ok((key, values))
             })
-            .collect::<Result<HashMap<[u8; 32], Vec<Felt>>, String>>()?;
+            .collect::<Result<BTreeMap<RpoDigest, Vec<Felt>>, String>>()?;
 
         Ok(Some(map))
     }

miden/tests/integration/operations/decorators/advice.rs

@@ -1,4 +1,9 @@
-use test_utils::{build_test, crypto::MerkleStore, rand::rand_value, Felt};
+use test_utils::{
+    build_test,
+    crypto::{MerkleStore, RpoDigest},
+    rand::rand_value,
+    Felt,
+};
 
 // ADVICE INJECTION
 // ================================================================================================
@@ -191,7 +196,7 @@ fn advice_push_mapval() {
 
     let stack_inputs = [1, 2, 3, 4];
     let adv_map = [(
-        key_to_bytes(stack_inputs),
+        RpoDigest::try_from(stack_inputs).unwrap(),
         vec![Felt::new(8), Felt::new(7), Felt::new(6), Felt::new(5)],
     )];
 
@@ -216,7 +221,7 @@ fn advice_push_mapval() {
 
     let stack_inputs = [1, 2, 3, 4];
     let adv_map = [(
-        key_to_bytes(stack_inputs),
+        RpoDigest::try_from(stack_inputs).unwrap(),
         vec![Felt::new(8), Felt::new(7), Felt::new(6), Felt::new(5)],
     )];
 
@@ -239,7 +244,7 @@ fn advice_push_mapval() {
 
     let stack_inputs = [1, 2, 3, 4];
     let adv_map = [(
-        key_to_bytes(stack_inputs),
+        RpoDigest::try_from(stack_inputs).unwrap(),
         vec![Felt::new(11), Felt::new(12), Felt::new(13), Felt::new(14), Felt::new(15)],
     )];
 
@@ -265,7 +270,7 @@ fn advice_push_mapval() {
 
     let stack_inputs = [1, 2, 3, 4];
     let adv_map = [(
-        key_to_bytes(stack_inputs),
+        RpoDigest::try_from(stack_inputs).unwrap(),
         vec![Felt::new(11), Felt::new(12), Felt::new(13), Felt::new(14), Felt::new(15)],
     )];
 
@@ -324,17 +329,3 @@ fn advice_insert_hdword() {
     let test = build_test!(source, &stack_inputs);
     test.expect_stack(&[1, 2, 3, 4, 5, 6, 7, 8]);
 }
-
-// HELPER FUNCTIONS
-// ================================================================================================
-
-fn key_to_bytes(key: [u64; 4]) -> [u8; 32] {
-    let mut result = [0; 32];
-
-    result[..8].copy_from_slice(&key[0].to_le_bytes());
-    result[8..16].copy_from_slice(&key[1].to_le_bytes());
-    result[16..24].copy_from_slice(&key[2].to_le_bytes());
-    result[24..].copy_from_slice(&key[3].to_le_bytes());
-
-    result
-}

processor/src/host/advice/injectors/smt.rs

@@ -88,7 +88,7 @@ fn get_smt_leaf_preimage<A: AdviceProvider>(
     advice_provider: &A,
     node: Word,
 ) -> Result<Vec<(Word, Word)>, ExecutionError> {
-    let node_bytes = RpoDigest::from(node).as_bytes();
+    let node_bytes = RpoDigest::from(node);
 
     let kv_pairs = advice_provider
         .get_mapped_values(&node_bytes)

processor/src/host/advice/inputs.rs

@@ -1,4 +1,6 @@
-use super::{BTreeMap, Felt, InnerNodeInfo, InputError, MerkleStore, Vec};
+use vm_core::crypto::hash::RpoDigest;
+
+use super::{map::AdviceMap, Felt, InnerNodeInfo, InputError, MerkleStore, Vec};
 
 // ADVICE INPUTS
 // ================================================================================================
@@ -18,7 +20,7 @@ use super::{BTreeMap, Felt, InnerNodeInfo, InputError, MerkleStore, Vec};
 #[derive(Clone, Debug, Default)]
 pub struct AdviceInputs {
     stack: Vec<Felt>,
-    map: BTreeMap<[u8; 32], Vec<Felt>>,
+    map: AdviceMap,
     store: MerkleStore,
 }
 
@@ -53,7 +55,7 @@ impl AdviceInputs {
     /// Extends the map of values with the given argument, replacing previously inserted items.
     pub fn with_map<I>(mut self, iter: I) -> Self
     where
-        I: IntoIterator<Item = ([u8; 32], Vec<Felt>)>,
+        I: IntoIterator<Item = (RpoDigest, Vec<Felt>)>,
     {
         self.map.extend(iter);
         self
@@ -79,7 +81,7 @@ impl AdviceInputs {
     /// Extends the map of values with the given argument, replacing previously inserted items.
     pub fn extend_map<I>(&mut self, iter: I)
     where
-        I: IntoIterator<Item = ([u8; 32], Vec<Felt>)>,
+        I: IntoIterator<Item = (RpoDigest, Vec<Felt>)>,
     {
         self.map.extend(iter);
     }
@@ -108,8 +110,8 @@ impl AdviceInputs {
     }
 
     /// Fetch a values set mapped by the given key.
-    pub fn mapped_values(&self, key: &[u8; 32]) -> Option<&[Felt]> {
-        self.map.get(key).map(Vec::as_slice)
+    pub fn mapped_values(&self, key: &RpoDigest) -> Option<&[Felt]> {
+        self.map.get(key)
     }
 
     /// Returns the underlying [MerkleStore].
@@ -122,7 +124,7 @@ impl AdviceInputs {
 
     /// Decomposes these `[Self]` into their raw components.
     #[allow(clippy::type_complexity)]
-    pub(crate) fn into_parts(self) -> (Vec<Felt>, BTreeMap<[u8; 32], Vec<Felt>>, MerkleStore) {
+    pub(crate) fn into_parts(self) -> (Vec<Felt>, AdviceMap, MerkleStore) {
         let Self { stack, map, store } = self;
         (stack, map, store)
     }
@@ -135,6 +137,6 @@ impl AdviceInputs {
 #[derive(Clone, Debug, Default)]
 pub struct AdviceInputs {
     pub stack: Vec<Felt>,
-    pub map: BTreeMap<[u8; 32], Vec<Felt>>,
+    pub map: AdviceMap,
     pub store: MerkleStore,
 }

processor/src/host/advice/map.rs

@@ -0,0 +1,50 @@
+use super::{BTreeMap, Felt, Vec};
+use vm_core::{crypto::hash::RpoDigest, utils::collections::btree_map::IntoIter};
+
+// ADVICE MAP
+// ================================================================================================
+
+/// Defines a set of non-deterministic (advice) inputs which the VM can access by their keys.
+///
+/// Each key maps to one or more field element. To access the elements, the VM can move the values
+/// associated with a given key onto the advice stack using `adv.push_mapval` instruction. The VM
+/// can also insert new values into the advice map during execution.
+#[derive(Debug, Clone, Default)]
+pub struct AdviceMap(BTreeMap<RpoDigest, Vec<Felt>>);
+
+impl AdviceMap {
+    /// Creates a new advice map.
+    pub fn new() -> Self {
+        Self(BTreeMap::<RpoDigest, Vec<Felt>>::new())
+    }
+
+    /// Returns the values associated with given key.
+    pub fn get(&self, key: &RpoDigest) -> Option<&[Felt]> {
+        self.0.get(key).map(|v| v.as_slice())
+    }
+
+    /// Inserts a key value pair in the advice map and returns the inserted value.
+    pub fn insert(&mut self, key: RpoDigest, value: Vec<Felt>) -> Option<Vec<Felt>> {
+        self.0.insert(key, value)
+    }
+
+    /// Removes the value associated with the key and returns the removed element.
+    pub fn remove(&mut self, key: RpoDigest) -> Option<Vec<Felt>> {
+        self.0.remove(&key)
+    }
+}
+
+impl IntoIterator for AdviceMap {
+    type Item = (RpoDigest, Vec<Felt>);
+    type IntoIter = IntoIter<RpoDigest, Vec<Felt>>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.0.into_iter()
+    }
+}
+
+impl Extend<(RpoDigest, Vec<Felt>)> for AdviceMap {
+    fn extend<T: IntoIterator<Item = (RpoDigest, Vec<Felt>)>>(&mut self, iter: T) {
+        self.0.extend(iter)
+    }
+}

processor/src/host/advice/mod.rs

@@ -6,10 +6,7 @@ use vm_core::{
         hash::RpoDigest,
         merkle::{InnerNodeInfo, MerklePath, MerkleStore, NodeIndex, StoreNode},
     },
-    utils::{
-        collections::{BTreeMap, KvMap, RecordingMap, Vec},
-        IntoBytes,
-    },
+    utils::collections::{BTreeMap, KvMap, RecordingMap, Vec},
     AdviceInjector, SignatureKind,
 };
 
@@ -27,6 +24,8 @@ pub use providers::{MemAdviceProvider, RecAdviceProvider};
 mod source;
 pub use source::AdviceSource;
 
+mod map;
+
 // ADVICE PROVIDER
 // ================================================================================================
 
@@ -528,7 +527,7 @@ pub trait AdviceProvider: Sized {
     // --------------------------------------------------------------------------------------------
 
     /// Returns a reference to the value(s) associated with the specified key in the advice map.
-    fn get_mapped_values(&self, key: &[u8; 32]) -> Option<&[Felt]>;
+    fn get_mapped_values(&self, key: &RpoDigest) -> Option<&[Felt]>;
 
     /// Inserts the provided value into the advice map under the specified key.
     ///
@@ -685,7 +684,7 @@ where
         T::get_signature(self, kind, pub_key, msg)
     }
 
-    fn get_mapped_values(&self, key: &[u8; 32]) -> Option<&[Felt]> {
+    fn get_mapped_values(&self, key: &RpoDigest) -> Option<&[Felt]> {
         T::get_mapped_values(self, key)
     }
 

processor/src/host/advice/providers.rs

@@ -1,7 +1,6 @@
 use super::{
-    injectors, AdviceInputs, AdviceProvider, AdviceSource, BTreeMap, ExecutionError, Felt,
-    IntoBytes, KvMap, MerklePath, MerkleStore, NodeIndex, RecordingMap, RpoDigest, StoreNode, Vec,
-    Word,
+    injectors, AdviceInputs, AdviceProvider, AdviceSource, BTreeMap, ExecutionError, Felt, KvMap,
+    MerklePath, MerkleStore, NodeIndex, RecordingMap, RpoDigest, StoreNode, Vec, Word,
 };
 use crate::ProcessState;
 use vm_core::SignatureKind;
@@ -12,8 +11,8 @@ use vm_core::SignatureKind;
 type SimpleMerkleMap = BTreeMap<RpoDigest, StoreNode>;
 type RecordingMerkleMap = RecordingMap<RpoDigest, StoreNode>;
 
-type SimpleAdviceMap = BTreeMap<[u8; 32], Vec<Felt>>;
-type RecordingAdviceMap = RecordingMap<[u8; 32], Vec<Felt>>;
+type SimpleAdviceMap = BTreeMap<RpoDigest, Vec<Felt>>;
+type RecordingAdviceMap = RecordingMap<RpoDigest, Vec<Felt>>;
 
 // BASE ADVICE PROVIDER
 // ================================================================================================
@@ -23,7 +22,7 @@ type RecordingAdviceMap = RecordingMap<[u8; 32], Vec<Felt>>;
 #[derive(Debug, Clone, Default)]
 pub struct BaseAdviceProvider<M, S>
 where
-    M: KvMap<[u8; 32], Vec<Felt>>,
+    M: KvMap<RpoDigest, Vec<Felt>>,
     S: KvMap<RpoDigest, StoreNode>,
 {
     stack: Vec<Felt>,
@@ -33,7 +32,7 @@ where
 
 impl<M, S> From<AdviceInputs> for BaseAdviceProvider<M, S>
 where
-    M: KvMap<[u8; 32], Vec<Felt>>,
+    M: KvMap<RpoDigest, Vec<Felt>>,
     S: KvMap<RpoDigest, StoreNode>,
 {
     fn from(inputs: AdviceInputs) -> Self {
@@ -49,7 +48,7 @@ where
 
 impl<M, S> AdviceProvider for BaseAdviceProvider<M, S>
 where
-    M: KvMap<[u8; 32], Vec<Felt>>,
+    M: KvMap<RpoDigest, Vec<Felt>>,
     S: KvMap<RpoDigest, StoreNode>,
 {
     // ADVICE STACK
@@ -92,10 +91,8 @@ where
                 self.stack.extend(word.iter().rev());
             }
             AdviceSource::Map { key, include_len } => {
-                let values = self
-                    .map
-                    .get(&key.into_bytes())
-                    .ok_or(ExecutionError::AdviceMapKeyNotFound(key))?;
+                let values =
+                    self.map.get(&key.into()).ok_or(ExecutionError::AdviceMapKeyNotFound(key))?;
 
                 self.stack.extend(values.iter().rev());
                 if include_len {
@@ -116,7 +113,7 @@ where
     ) -> Result<Vec<Felt>, ExecutionError> {
         let pk_sk = self
             .map
-            .get(&pub_key.into_bytes())
+            .get(&pub_key.into())
             .ok_or(ExecutionError::AdviceMapKeyNotFound(pub_key))?;
 
         match kind {
@@ -127,12 +124,12 @@ where
     // ADVICE MAP
     // --------------------------------------------------------------------------------------------
 
-    fn get_mapped_values(&self, key: &[u8; 32]) -> Option<&[Felt]> {
+    fn get_mapped_values(&self, key: &RpoDigest) -> Option<&[Felt]> {
         self.map.get(key).map(|v| v.as_slice())
     }
 
     fn insert_into_map(&mut self, key: Word, values: Vec<Felt>) -> Result<(), ExecutionError> {
-        self.map.insert(key.into_bytes(), values);
+        self.map.insert(key.into(), values);
         Ok(())
     }
 
@@ -304,7 +301,7 @@ impl AdviceProvider for MemAdviceProvider {
         self.provider.get_signature(kind, pub_key, msg)
     }
 
-    fn get_mapped_values(&self, key: &[u8; 32]) -> Option<&[Felt]> {
+    fn get_mapped_values(&self, key: &RpoDigest) -> Option<&[Felt]> {
         self.provider.get_mapped_values(key)
     }
 
@@ -429,7 +426,7 @@ impl AdviceProvider for RecAdviceProvider {
         self.provider.get_signature(kind, pub_key, msg)
     }
 
-    fn get_mapped_values(&self, key: &[u8; 32]) -> Option<&[Felt]> {
+    fn get_mapped_values(&self, key: &RpoDigest) -> Option<&[Felt]> {
         self.provider.get_mapped_values(key)
     }