Merge pull request #987 from 0xPolygonMiden/frisitano-tx-result

feat: Extend `AdviceProvider` trait
0xPolygonMiden · Jul 4, 2023 · 7cfaaec · 7cfaaec
2 parents 3aa346f + 764e8be
commit 7cfaaec
Show file tree

Hide file tree

Showing 2 changed files with 63 additions and 0 deletions.
diff --git a/processor/src/advice/mod.rs b/processor/src/advice/mod.rs
@@ -1,4 +1,5 @@
 use super::{ExecutionError, Felt, InputError, StarkField, Word};
+use core::borrow::Borrow;
 use vm_core::{
     crypto::{
         hash::RpoDigest,
@@ -100,6 +101,11 @@ pub trait AdviceProvider {
     /// are replaced with the specified values.
     fn insert_into_map(&mut self, key: Word, values: Vec<Felt>) -> Result<(), ExecutionError>;
 
+    // ADVICE MAP
+    // --------------------------------------------------------------------------------------------
+    /// 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]>;
+
     // ADVISE SETS
     // --------------------------------------------------------------------------------------------
 
@@ -177,6 +183,15 @@ pub trait AdviceProvider {
     /// advice provider.
     fn merge_roots(&mut self, lhs: Word, rhs: Word) -> Result<Word, ExecutionError>;
 
+    /// Returns a subset of this Merkle store such that the returned Merkle store contains all
+    /// nodes which are descendants of the specified roots.
+    ///
+    /// The roots for which no descendants exist in this Merkle store are ignored.
+    fn get_store_subset<I, R>(&self, roots: I) -> MerkleStore
+    where
+        I: Iterator<Item = R>,
+        R: Borrow<RpoDigest>;
+
     // CONTEXT MANAGEMENT
     // --------------------------------------------------------------------------------------------
 
@@ -214,6 +229,10 @@ where
         T::insert_into_map(self, key, values)
     }
 
+    fn get_mapped_values(&self, key: &[u8; 32]) -> Option<&[Felt]> {
+        T::get_mapped_values(self, key)
+    }
+
     fn get_tree_node(
         &self,
         root: Word,
@@ -255,6 +274,14 @@ where
         T::merge_roots(self, lhs, rhs)
     }
 
+    fn get_store_subset<I, R>(&self, roots: I) -> MerkleStore
+    where
+        I: Iterator<Item = R>,
+        R: Borrow<RpoDigest>,
+    {
+        T::get_store_subset(self, roots)
+    }
+
     fn advance_clock(&mut self) {
         T::advance_clock(self)
     }

diff --git a/processor/src/advice/providers.rs b/processor/src/advice/providers.rs
@@ -106,6 +106,12 @@ where
         Ok(())
     }
 
+    // ADVICE MAP
+    // --------------------------------------------------------------------------------------------
+    fn get_mapped_values(&self, key: &[u8; 32]) -> Option<&[Felt]> {
+        self.map.get(key).map(|v| v.as_slice())
+    }
+
     // ADVISE SETS
     // --------------------------------------------------------------------------------------------
 
@@ -184,6 +190,14 @@ where
             .map_err(ExecutionError::MerkleStoreMergeFailed)
     }
 
+    fn get_store_subset<I, R>(&self, roots: I) -> MerkleStore
+    where
+        I: Iterator<Item = R>,
+        R: core::borrow::Borrow<RpoDigest>,
+    {
+        self.store.subset(roots).into_inner().into_iter().collect()
+    }
+
     // CONTEXT MANAGEMENT
     // --------------------------------------------------------------------------------------------
 
@@ -257,6 +271,10 @@ impl AdviceProvider for MemAdviceProvider {
         self.provider.insert_into_map(key, values)
     }
 
+    fn get_mapped_values(&self, key: &[u8; 32]) -> Option<&[Felt]> {
+        self.provider.get_mapped_values(key)
+    }
+
     fn get_tree_node(&self, root: Word, depth: &Felt, index: &Felt) -> Result<Word, ExecutionError> {
         self.provider.get_tree_node(root, depth, index)
     }
@@ -277,6 +295,13 @@ impl AdviceProvider for MemAdviceProvider {
         self.provider.merge_roots(lhs, rhs)
     }
 
+    fn get_store_subset<I, R>(&self, roots: I) -> MerkleStore
+        where
+            I: Iterator<Item = R>,
+            R: core::borrow::Borrow<RpoDigest> {
+        self.provider.get_store_subset(roots)
+    }
+
     fn advance_clock(&mut self) {
         self.provider.advance_clock()
     }
@@ -383,6 +408,10 @@ impl AdviceProvider for RecAdviceProvider {
         self.provider.insert_into_map(key, values)
     }
 
+    fn get_mapped_values(&self, key: &[u8; 32]) -> Option<&[Felt]> {
+        self.provider.get_mapped_values(key)
+    }
+
     fn get_tree_node(&self, root: Word, depth: &Felt, index: &Felt) -> Result<Word, ExecutionError> {
         self.provider.get_tree_node(root, depth, index)
     }
@@ -403,6 +432,13 @@ impl AdviceProvider for RecAdviceProvider {
         self.provider.merge_roots(lhs, rhs)
     }
 
+    fn get_store_subset<I, R>(&self, roots: I) -> MerkleStore
+        where
+            I: Iterator<Item = R>,
+            R: core::borrow::Borrow<RpoDigest> {
+        self.provider.get_store_subset(roots)
+    }
+
     fn advance_clock(&mut self) {
         self.provider.advance_clock()
     }