use storage PrefixIter instead of DB iterator that respects write log

shared/src/ledger/native_vp/mod.rs

@@ -176,7 +176,7 @@ where
     H: 'static + StorageHasher,
     CA: 'static + WasmCacheAccess,
 {
-    type PrefixIter<'iter> = <DB as storage::DBIter<'iter>>::PrefixIter where Self: 'iter;
+    type PrefixIter<'iter> = storage::PrefixIter<'iter, DB> where Self: 'iter;
 
     fn read_bytes(
         &self,
@@ -198,30 +198,34 @@ where
         vp_host_fns::has_key_pre(
             &mut self.ctx.gas_meter.borrow_mut(),
             self.ctx.storage,
+            self.ctx.write_log,
             key,
         )
         .into_storage_result()
     }
 
-    fn iter_next<'iter>(
+    fn iter_prefix<'iter>(
         &'iter self,
-        iter: &mut Self::PrefixIter<'iter>,
-    ) -> Result<Option<(String, Vec<u8>)>, storage_api::Error> {
-        vp_host_fns::iter_pre_next::<DB>(
+        prefix: &crate::types::storage::Key,
+    ) -> Result<Self::PrefixIter<'iter>, storage_api::Error> {
+        vp_host_fns::iter_prefix_pre(
             &mut self.ctx.gas_meter.borrow_mut(),
-            iter,
+            self.ctx.write_log,
+            self.ctx.storage,
+            prefix,
         )
         .into_storage_result()
     }
 
     // ---- Methods below are implemented in `self.ctx`, because they are
     //      the same in `pre/post` ----
 
-    fn iter_prefix<'iter>(
+    fn iter_next<'iter>(
         &'iter self,
-        prefix: &crate::types::storage::Key,
-    ) -> Result<Self::PrefixIter<'iter>, storage_api::Error> {
-        self.ctx.iter_prefix(prefix)
+        iter: &mut Self::PrefixIter<'iter>,
+    ) -> Result<Option<(String, Vec<u8>)>, storage_api::Error> {
+        vp_host_fns::iter_next::<DB>(&mut self.ctx.gas_meter.borrow_mut(), iter)
+            .into_storage_result()
     }
 
     fn get_chain_id(&self) -> Result<String, storage_api::Error> {
@@ -256,7 +260,7 @@ where
     H: 'static + StorageHasher,
     CA: 'static + WasmCacheAccess,
 {
-    type PrefixIter<'iter> = <DB as storage::DBIter<'iter>>::PrefixIter where Self:'iter;
+    type PrefixIter<'iter> = storage::PrefixIter<'iter, DB> where Self: 'iter;
 
     fn read_bytes(
         &self,
@@ -284,26 +288,28 @@ where
         .into_storage_result()
     }
 
-    fn iter_next<'iter>(
+    fn iter_prefix<'iter>(
         &'iter self,
-        iter: &mut Self::PrefixIter<'iter>,
-    ) -> Result<Option<(String, Vec<u8>)>, storage_api::Error> {
-        vp_host_fns::iter_post_next::<DB>(
+        prefix: &crate::types::storage::Key,
+    ) -> Result<Self::PrefixIter<'iter>, storage_api::Error> {
+        vp_host_fns::iter_prefix_post(
             &mut self.ctx.gas_meter.borrow_mut(),
             self.ctx.write_log,
-            iter,
+            self.ctx.storage,
+            prefix,
         )
         .into_storage_result()
     }
 
     // ---- Methods below are implemented in `self.ctx`, because they are
     //      the same in `pre/post` ----
 
-    fn iter_prefix<'iter>(
+    fn iter_next<'iter>(
         &'iter self,
-        prefix: &crate::types::storage::Key,
-    ) -> Result<Self::PrefixIter<'iter>, storage_api::Error> {
-        self.ctx.iter_prefix(prefix)
+        iter: &mut Self::PrefixIter<'iter>,
+    ) -> Result<Option<(String, Vec<u8>)>, storage_api::Error> {
+        vp_host_fns::iter_next::<DB>(&mut self.ctx.gas_meter.borrow_mut(), iter)
+            .into_storage_result()
     }
 
     fn get_chain_id(&self) -> Result<String, storage_api::Error> {
@@ -339,7 +345,7 @@ where
 {
     type Post = CtxPostStorageRead<'view, 'a, DB, H, CA>;
     type Pre = CtxPreStorageRead<'view, 'a, DB, H, CA>;
-    type PrefixIter<'iter> = <DB as storage::DBIter<'iter>>::PrefixIter where Self: 'iter;
+    type PrefixIter<'iter> = storage::PrefixIter<'iter, DB> where Self: 'iter;
 
     fn pre(&'view self) -> Self::Pre {
         CtxPreStorageRead { ctx: self }
@@ -426,8 +432,9 @@ where
         &'iter self,
         prefix: &Key,
     ) -> Result<Self::PrefixIter<'iter>, storage_api::Error> {
-        vp_host_fns::iter_prefix(
+        vp_host_fns::iter_prefix_pre(
             &mut self.gas_meter.borrow_mut(),
+            self.write_log,
             self.storage,
             prefix,
         )

shared/src/ledger/vp_host_fns.rs

@@ -311,71 +311,51 @@ where
     Ok(storage.native_token.clone())
 }
 
-/// Storage prefix iterator, ordered by storage keys. It will try to get an
-/// iterator from the storage.
-pub fn iter_prefix<'a, DB, H>(
+/// Storage prefix iterator for prior state (before tx execution), ordered by
+/// storage keys. It will try to get an iterator from the storage.
+pub fn iter_prefix_pre<'a, DB, H>(
     gas_meter: &mut VpGasMeter,
+    write_log: &'a WriteLog,
     storage: &'a Storage<DB, H>,
     prefix: &Key,
-) -> EnvResult<<DB as storage::DBIter<'a>>::PrefixIter>
+) -> EnvResult<storage::PrefixIter<'a, DB>>
 where
     DB: storage::DB + for<'iter> storage::DBIter<'iter>,
     H: StorageHasher,
 {
-    let (iter, gas) = storage.iter_prefix(prefix);
+    let (iter, gas) = storage::iter_prefix_pre(write_log, storage, prefix);
     add_gas(gas_meter, gas)?;
     Ok(iter)
 }
 
-/// Storage prefix iterator for prior state (before tx execution). It will try
-/// to read from the storage.
-pub fn iter_pre_next<DB>(
+/// Storage prefix iterator for posterior state (after tx execution), ordered by
+/// storage keys. It will try to get an iterator from the storage.
+pub fn iter_prefix_post<'a, DB, H>(
     gas_meter: &mut VpGasMeter,
-    iter: &mut <DB as storage::DBIter<'_>>::PrefixIter,
-) -> EnvResult<Option<(String, Vec<u8>)>>
+    write_log: &'a WriteLog,
+    storage: &'a Storage<DB, H>,
+    prefix: &Key,
+) -> EnvResult<storage::PrefixIter<'a, DB>>
 where
     DB: storage::DB + for<'iter> storage::DBIter<'iter>,
+    H: StorageHasher,
 {
-    if let Some((key, val, gas)) = iter.next() {
-        add_gas(gas_meter, gas)?;
-        return Ok(Some((key, val)));
-    }
-    Ok(None)
+    let (iter, gas) = storage::iter_prefix_post(write_log, storage, prefix);
+    add_gas(gas_meter, gas)?;
+    Ok(iter)
 }
 
-/// Storage prefix iterator next for posterior state (after tx execution). It
-/// will try to read from the write log first and if no entry found then from
-/// the storage.
-pub fn iter_post_next<DB>(
+/// Get the next item in a storage prefix iterator (pre or post).
+pub fn iter_next<DB>(
     gas_meter: &mut VpGasMeter,
-    write_log: &WriteLog,
-    iter: &mut <DB as storage::DBIter<'_>>::PrefixIter,
+    iter: &mut storage::PrefixIter<DB>,
 ) -> EnvResult<Option<(String, Vec<u8>)>>
 where
     DB: storage::DB + for<'iter> storage::DBIter<'iter>,
 {
-    for (key, val, iter_gas) in iter {
-        let (log_val, log_gas) = write_log.read(
-            &Key::parse(key.clone()).map_err(RuntimeError::StorageDataError)?,
-        );
-        add_gas(gas_meter, iter_gas + log_gas)?;
-        match log_val {
-            Some(&write_log::StorageModification::Write { ref value }) => {
-                return Ok(Some((key, value.clone())));
-            }
-            Some(&write_log::StorageModification::Delete) => {
-                // check the next because the key has already deleted
-                continue;
-            }
-            Some(&write_log::StorageModification::InitAccount { .. }) => {
-                // a VP of a new account doesn't need to be iterated
-                continue;
-            }
-            Some(&write_log::StorageModification::Temp { .. }) => {
-                return Err(RuntimeError::ReadTemporaryValueError);
-            }
-            None => return Ok(Some((key, val))),
-        }
+    if let Some((key, val, gas)) = iter.next() {
+        add_gas(gas_meter, gas)?;
+        return Ok(Some((key, val)));
     }
     Ok(None)
 }

shared/src/vm/host_env.rs

@@ -696,7 +696,7 @@ pub fn tx_iter_prefix<MEM, DB, H, CA>(
 ) -> TxResult<u64>
 where
     MEM: VmMemory,
-    DB: storage::DB + for<'iter> storage::DBIter<'iter>,
+    DB: 'static + storage::DB + for<'iter> storage::DBIter<'iter>,
     H: StorageHasher,
     CA: WasmCacheAccess,
 {
@@ -706,15 +706,17 @@ where
         .map_err(|e| TxRuntimeError::MemoryError(Box::new(e)))?;
     tx_add_gas(env, gas)?;
 
-    tracing::debug!("tx_iter_prefix {}, prefix {}", prefix, prefix_ptr);
+    tracing::debug!("tx_iter_prefix {}", prefix);
 
     let prefix =
         Key::parse(prefix).map_err(TxRuntimeError::StorageDataError)?;
 
+    let write_log = unsafe { env.ctx.write_log.get() };
     let storage = unsafe { env.ctx.storage.get() };
-    let iterators = unsafe { env.ctx.iterators.get() };
-    let (iter, gas) = storage.iter_prefix(&prefix);
+    let (iter, gas) = storage::iter_prefix_post(write_log, storage, &prefix);
     tx_add_gas(env, gas)?;
+
+    let iterators = unsafe { env.ctx.iterators.get() };
     Ok(iterators.insert(iter).id())
 }
 
@@ -1183,7 +1185,9 @@ where
     let key =
         Key::parse(key).map_err(vp_host_fns::RuntimeError::StorageDataError)?;
     let storage = unsafe { env.ctx.storage.get() };
-    let present = vp_host_fns::has_key_pre(gas_meter, storage, &key)?;
+    let write_log = unsafe { env.ctx.write_log.get() };
+    let present =
+        vp_host_fns::has_key_pre(gas_meter, storage, write_log, &key)?;
     Ok(HostEnvResult::from(present).to_i64())
 }
 
@@ -1220,17 +1224,18 @@ where
     Ok(HostEnvResult::from(present).to_i64())
 }
 
-/// Storage prefix iterator function exposed to the wasm VM VP environment.
-/// It will try to get an iterator from the storage and return the corresponding
-/// ID of the iterator, ordered by storage keys.
-pub fn vp_iter_prefix<MEM, DB, H, EVAL, CA>(
+/// Storage prefix iterator function for prior state (before tx execution)
+/// exposed to the wasm VM VP environment. It will try to get an iterator from
+/// the storage and return the corresponding ID of the iterator, ordered by
+/// storage keys.
+pub fn vp_iter_prefix_pre<MEM, DB, H, EVAL, CA>(
     env: &VpVmEnv<MEM, DB, H, EVAL, CA>,
     prefix_ptr: u64,
     prefix_len: u64,
 ) -> vp_host_fns::EnvResult<u64>
 where
     MEM: VmMemory,
-    DB: storage::DB + for<'iter> storage::DBIter<'iter>,
+    DB: 'static + storage::DB + for<'iter> storage::DBIter<'iter>,
     H: StorageHasher,
     EVAL: VpEvaluator,
     CA: WasmCacheAccess,
@@ -1242,63 +1247,63 @@ where
     let gas_meter = unsafe { env.ctx.gas_meter.get() };
     vp_host_fns::add_gas(gas_meter, gas)?;
 
+    tracing::debug!("vp_iter_prefix_pre {}", prefix);
+
     let prefix = Key::parse(prefix)
         .map_err(vp_host_fns::RuntimeError::StorageDataError)?;
-    tracing::debug!("vp_iter_prefix {}", prefix);
 
+    let write_log = unsafe { env.ctx.write_log.get() };
     let storage = unsafe { env.ctx.storage.get() };
-    let iter = vp_host_fns::iter_prefix(gas_meter, storage, &prefix)?;
+    let iter =
+        vp_host_fns::iter_prefix_pre(gas_meter, write_log, storage, &prefix)?;
+
     let iterators = unsafe { env.ctx.iterators.get() };
     Ok(iterators.insert(iter).id())
 }
 
-/// Storage prefix iterator for prior state (before tx execution) function
-/// exposed to the wasm VM VP environment. It will try to read from the storage.
-///
-/// Returns `-1` when the key is not present, or the length of the data when
-/// the key is present (the length may be `0`).
-pub fn vp_iter_pre_next<MEM, DB, H, EVAL, CA>(
+/// Storage prefix iterator function for posterior state (after tx execution)
+/// exposed to the wasm VM VP environment. It will try to get an iterator from
+/// the storage and return the corresponding ID of the iterator, ordered by
+/// storage keys.
+pub fn vp_iter_prefix_post<MEM, DB, H, EVAL, CA>(
     env: &VpVmEnv<MEM, DB, H, EVAL, CA>,
-    iter_id: u64,
-) -> vp_host_fns::EnvResult<i64>
+    prefix_ptr: u64,
+    prefix_len: u64,
+) -> vp_host_fns::EnvResult<u64>
 where
     MEM: VmMemory,
-    DB: storage::DB + for<'iter> storage::DBIter<'iter>,
+    DB: 'static + storage::DB + for<'iter> storage::DBIter<'iter>,
     H: StorageHasher,
     EVAL: VpEvaluator,
     CA: WasmCacheAccess,
 {
-    tracing::debug!("vp_iter_pre_next iter_id {}", iter_id);
+    let (prefix, gas) = env
+        .memory
+        .read_string(prefix_ptr, prefix_len as _)
+        .map_err(|e| vp_host_fns::RuntimeError::MemoryError(Box::new(e)))?;
+    let gas_meter = unsafe { env.ctx.gas_meter.get() };
+    vp_host_fns::add_gas(gas_meter, gas)?;
+
+    tracing::debug!("vp_iter_prefix_post {}", prefix);
+
+    let prefix = Key::parse(prefix)
+        .map_err(vp_host_fns::RuntimeError::StorageDataError)?;
+
+    let write_log = unsafe { env.ctx.write_log.get() };
+    let storage = unsafe { env.ctx.storage.get() };
+    let iter =
+        vp_host_fns::iter_prefix_post(gas_meter, write_log, storage, &prefix)?;
 
     let iterators = unsafe { env.ctx.iterators.get() };
-    let iter_id = PrefixIteratorId::new(iter_id);
-    if let Some(iter) = iterators.get_mut(iter_id) {
-        let gas_meter = unsafe { env.ctx.gas_meter.get() };
-        if let Some((key, val)) =
-            vp_host_fns::iter_pre_next::<DB>(gas_meter, iter)?
-        {
-            let key_val = KeyVal { key, val }
-                .try_to_vec()
-                .map_err(vp_host_fns::RuntimeError::EncodingError)?;
-            let len: i64 = key_val
-                .len()
-                .try_into()
-                .map_err(vp_host_fns::RuntimeError::NumConversionError)?;
-            let result_buffer = unsafe { env.ctx.result_buffer.get() };
-            result_buffer.replace(key_val);
-            return Ok(len);
-        }
-    }
-    Ok(HostEnvResult::Fail.to_i64())
+    Ok(iterators.insert(iter).id())
 }
 
-/// Storage prefix iterator next for posterior state (after tx execution)
-/// function exposed to the wasm VM VP environment. It will try to read from the
-/// write log first and if no entry found then from the storage.
+/// Storage prefix iterator for prior or posterior state function
+/// exposed to the wasm VM VP environment.
 ///
 /// Returns `-1` when the key is not present, or the length of the data when
 /// the key is present (the length may be `0`).
-pub fn vp_iter_post_next<MEM, DB, H, EVAL, CA>(
+pub fn vp_iter_next<MEM, DB, H, EVAL, CA>(
     env: &VpVmEnv<MEM, DB, H, EVAL, CA>,
     iter_id: u64,
 ) -> vp_host_fns::EnvResult<i64>
@@ -1309,16 +1314,13 @@ where
     EVAL: VpEvaluator,
     CA: WasmCacheAccess,
 {
-    tracing::debug!("vp_iter_post_next iter_id {}", iter_id);
+    tracing::debug!("vp_iter_next iter_id {}", iter_id);
 
     let iterators = unsafe { env.ctx.iterators.get() };
     let iter_id = PrefixIteratorId::new(iter_id);
     if let Some(iter) = iterators.get_mut(iter_id) {
         let gas_meter = unsafe { env.ctx.gas_meter.get() };
-        let write_log = unsafe { env.ctx.write_log.get() };
-        if let Some((key, val)) =
-            vp_host_fns::iter_post_next::<DB>(gas_meter, write_log, iter)?
-        {
+        if let Some((key, val)) = vp_host_fns::iter_next(gas_meter, iter)? {
             let key_val = KeyVal { key, val }
                 .try_to_vec()
                 .map_err(vp_host_fns::RuntimeError::EncodingError)?;