Merge pull request #1834 from CosmWasm/1498-key-only-iteration

Key- and value-only iteration

CHANGELOG.md

@@ -14,11 +14,15 @@ and this project adheres to
   `Int{64,128,256,512}`. It is now const for all integer types.
 - cosmwasm-std: Implement `TryFrom<Decimal256>` for `Decimal` ([#1832])
 - cosmwasm-std: Add `StdAck`. ([#1512])
+- cosmwasm-std: Add new imports `db_next_{key, value}` for iterating storage
+  keys / values only and make `Storage::{range_keys, range_values}` more
+  efficient. This requires the `cosmwasm_1_4` feature to be enabled. ([#1834])
 
 [#1512]: https://github.com/CosmWasm/cosmwasm/issues/1512
 [#1799]: https://github.com/CosmWasm/cosmwasm/pull/1799
 [#1806]: https://github.com/CosmWasm/cosmwasm/pull/1806
 [#1832]: https://github.com/CosmWasm/cosmwasm/pull/1832
+[#1834]: https://github.com/CosmWasm/cosmwasm/pull/1834
 
 ### Changed
 

contracts/burner/Cargo.toml

@@ -33,7 +33,7 @@ backtraces = ["cosmwasm-std/backtraces", "cosmwasm-vm/backtraces"]
 
 [dependencies]
 cosmwasm-schema = { path = "../../packages/schema" }
-cosmwasm-std = { path = "../../packages/std", features = ["iterator"] }
+cosmwasm-std = { path = "../../packages/std", features = ["iterator", "cosmwasm_1_4"] }
 schemars = "0.8.3"
 serde = { version = "1.0.103", default-features = false, features = ["derive"] }
 

contracts/burner/src/contract.rs

@@ -52,9 +52,8 @@ pub fn execute_cleanup(
         let take_this_scan = std::cmp::min(PER_SCAN, limit);
         let keys: Vec<_> = deps
             .storage
-            .range(None, None, Order::Ascending)
+            .range_keys(None, None, Order::Ascending)
             .take(take_this_scan)
-            .map(|(k, _)| k)
             .collect();
         let deleted_this_scan = keys.len();
         for k in keys {

packages/std/Cargo.toml

@@ -42,6 +42,10 @@ cosmwasm_1_2 = ["cosmwasm_1_1"]
 # This feature makes `BankQuery::DenomMetadata` available for the contract to call, but requires
 # the host blockchain to run CosmWasm `1.3.0` or higher.
 cosmwasm_1_3 = ["cosmwasm_1_2"]
+# Together with the `iterator` feature this enables additional imports for more
+# efficient iteration over DB keys or values.
+# It requires the host blockchain to run CosmWasm `1.4.0` or higher.
+cosmwasm_1_4 = ["cosmwasm_1_3"]
 
 [dependencies]
 base64 = "0.21.0"

packages/std/src/imports.rs

@@ -37,6 +37,10 @@ extern "C" {
     fn db_scan(start_ptr: u32, end_ptr: u32, order: i32) -> u32;
     #[cfg(feature = "iterator")]
     fn db_next(iterator_id: u32) -> u32;
+    #[cfg(all(feature = "iterator", feature = "cosmwasm_1_4"))]
+    fn db_next_key(iterator_id: u32) -> u32;
+    #[cfg(all(feature = "iterator", feature = "cosmwasm_1_4"))]
+    fn db_next_value(iterator_id: u32) -> u32;
 
     fn addr_validate(source_ptr: u32) -> u32;
     fn addr_canonicalize(source_ptr: u32, destination_ptr: u32) -> u32;
@@ -129,16 +133,89 @@ impl Storage for ExternalStorage {
         end: Option<&[u8]>,
         order: Order,
     ) -> Box<dyn Iterator<Item = Record>> {
-        // There is lots of gotchas on turning options into regions for FFI, thus this design
-        // See: https://github.com/CosmWasm/cosmwasm/pull/509
-        let start_region = start.map(build_region);
-        let end_region = end.map(build_region);
-        let start_region_addr = get_optional_region_address(&start_region.as_ref());
-        let end_region_addr = get_optional_region_address(&end_region.as_ref());
-        let iterator_id = unsafe { db_scan(start_region_addr, end_region_addr, order as i32) };
+        let iterator_id = create_iter(start, end, order);
         let iter = ExternalIterator { iterator_id };
         Box::new(iter)
     }
+
+    #[cfg(all(feature = "cosmwasm_1_4", feature = "iterator"))]
+    fn range_keys<'a>(
+        &'a self,
+        start: Option<&[u8]>,
+        end: Option<&[u8]>,
+        order: Order,
+    ) -> Box<dyn Iterator<Item = Vec<u8>> + 'a> {
+        let iterator_id = create_iter(start, end, order);
+        let iter = ExternalPartialIterator {
+            iterator_id,
+            partial_type: PartialType::Keys,
+        };
+        Box::new(iter)
+    }
+
+    #[cfg(all(feature = "cosmwasm_1_4", feature = "iterator"))]
+    fn range_values<'a>(
+        &'a self,
+        start: Option<&[u8]>,
+        end: Option<&[u8]>,
+        order: Order,
+    ) -> Box<dyn Iterator<Item = Vec<u8>> + 'a> {
+        let iterator_id = create_iter(start, end, order);
+        let iter = ExternalPartialIterator {
+            iterator_id,
+            partial_type: PartialType::Values,
+        };
+        Box::new(iter)
+    }
+}
+
+#[cfg(feature = "iterator")]
+fn create_iter(start: Option<&[u8]>, end: Option<&[u8]>, order: Order) -> u32 {
+    // There is lots of gotchas on turning options into regions for FFI, thus this design
+    // See: https://github.com/CosmWasm/cosmwasm/pull/509
+    let start_region = start.map(build_region);
+    let end_region = end.map(build_region);
+    let start_region_addr = get_optional_region_address(&start_region.as_ref());
+    let end_region_addr = get_optional_region_address(&end_region.as_ref());
+    unsafe { db_scan(start_region_addr, end_region_addr, order as i32) }
+}
+
+#[cfg(all(feature = "cosmwasm_1_4", feature = "iterator"))]
+enum PartialType {
+    Keys,
+    Values,
+}
+
+/// ExternalPartialIterator makes a call out to `next_key` or `next_value`
+/// depending on its `partial_type`.
+/// Compared to `ExternalIterator`, it allows iterating only over the keys or
+/// values instead of both.
+#[cfg(all(feature = "cosmwasm_1_4", feature = "iterator"))]
+struct ExternalPartialIterator {
+    iterator_id: u32,
+    partial_type: PartialType,
+}
+
+#[cfg(all(feature = "cosmwasm_1_4", feature = "iterator"))]
+impl Iterator for ExternalPartialIterator {
+    type Item = Vec<u8>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        // here we differentiate between the two types
+        let next_result = match self.partial_type {
+            PartialType::Keys => unsafe { db_next_key(self.iterator_id) },
+            PartialType::Values => unsafe { db_next_value(self.iterator_id) },
+        };
+
+        if next_result == 0 {
+            // iterator is done
+            return None;
+        }
+
+        let data_region = next_result as *mut Region;
+        let data = unsafe { consume_region(data_region) };
+        Some(data)
+    }
 }
 
 #[cfg(feature = "iterator")]

packages/vm/src/backend.rs

@@ -122,6 +122,34 @@ pub trait Storage {
     #[cfg(feature = "iterator")]
     fn next(&mut self, iterator_id: u32) -> BackendResult<Option<Record>>;
 
+    /// Returns the next value of the iterator with the given ID.
+    /// Since the iterator is incremented, the corresponding key will never be accessible.
+    ///
+    /// If the ID is not found, a BackendError::IteratorDoesNotExist is returned.
+    ///
+    /// The default implementation uses [`Storage::next`] and discards the key.
+    /// More efficient implementations might be possible depending on the storage.
+    #[cfg(feature = "iterator")]
+    fn next_value(&mut self, iterator_id: u32) -> BackendResult<Option<Vec<u8>>> {
+        let (result, gas_info) = self.next(iterator_id);
+        let result = result.map(|record| record.map(|(_, v)| v));
+        (result, gas_info)
+    }
+
+    /// Returns the next key of the iterator with the given ID.
+    /// Since the iterator is incremented, the corresponding value will never be accessible.
+    ///
+    /// If the ID is not found, a BackendError::IteratorDoesNotExist is returned.
+    ///
+    /// The default implementation uses [`Storage::next`] and discards the value.
+    /// More efficient implementations might be possible depending on the storage.
+    #[cfg(feature = "iterator")]
+    fn next_key(&mut self, iterator_id: u32) -> BackendResult<Option<Vec<u8>>> {
+        let (result, gas_info) = self.next(iterator_id);
+        let result = result.map(|record| record.map(|(k, _)| k));
+        (result, gas_info)
+    }
+
     fn set(&mut self, key: &[u8], value: &[u8]) -> BackendResult<()>;
 
     /// Removes a database entry at `key`.

packages/vm/src/compatibility.rs

@@ -30,6 +30,10 @@ const SUPPORTED_IMPORTS: &[&str] = &[
     "env.db_scan",
     #[cfg(feature = "iterator")]
     "env.db_next",
+    #[cfg(feature = "iterator")]
+    "env.db_next_key",
+    #[cfg(feature = "iterator")]
+    "env.db_next_value",
 ];
 
 /// Lists all entry points we expect to be present when calling a contract.

packages/vm/src/environment.rs

@@ -493,6 +493,8 @@ mod tests {
                 "db_remove" => Function::new_typed(&mut store, |_a: u32| {}),
                 "db_scan" => Function::new_typed(&mut store, |_a: u32, _b: u32, _c: i32| -> u32 { 0 }),
                 "db_next" => Function::new_typed(&mut store, |_a: u32| -> u32 { 0 }),
+                "db_next_key" => Function::new_typed(&mut store, |_a: u32| -> u32 { 0 }),
+                "db_next_value" => Function::new_typed(&mut store, |_a: u32| -> u32 { 0 }),
                 "query_chain" => Function::new_typed(&mut store, |_a: u32| -> u32 { 0 }),
                 "addr_validate" => Function::new_typed(&mut store, |_a: u32| -> u32 { 0 }),
                 "addr_canonicalize" => Function::new_typed(&mut store, |_a: u32, _b: u32| -> u32 { 0 }),

packages/vm/src/imports.rs

@@ -533,6 +533,46 @@ pub fn do_db_next<A: BackendApi + 'static, S: Storage + 'static, Q: Querier + 's
     write_to_contract(data, &mut store, &out_data)
 }
 
+#[cfg(feature = "iterator")]
+pub fn do_db_next_key<A: BackendApi + 'static, S: Storage + 'static, Q: Querier + 'static>(
+    mut env: FunctionEnvMut<Environment<A, S, Q>>,
+    iterator_id: u32,
+) -> VmResult<u32> {
+    let (data, mut store) = env.data_and_store_mut();
+
+    let (result, gas_info) =
+        data.with_storage_from_context::<_, _>(|store| Ok(store.next_key(iterator_id)))?;
+
+    process_gas_info(data, &mut store, gas_info)?;
+
+    let key = match result? {
+        Some(key) => key,
+        None => return Ok(0),
+    };
+
+    write_to_contract(data, &mut store, &key)
+}
+
+#[cfg(feature = "iterator")]
+pub fn do_db_next_value<A: BackendApi + 'static, S: Storage + 'static, Q: Querier + 'static>(
+    mut env: FunctionEnvMut<Environment<A, S, Q>>,
+    iterator_id: u32,
+) -> VmResult<u32> {
+    let (data, mut store) = env.data_and_store_mut();
+
+    let (result, gas_info) =
+        data.with_storage_from_context::<_, _>(|store| Ok(store.next_value(iterator_id)))?;
+
+    process_gas_info(data, &mut store, gas_info)?;
+
+    let value = match result? {
+        Some(value) => value,
+        None => return Ok(0),
+    };
+
+    write_to_contract(data, &mut store, &value)
+}
+
 /// Creates a Region in the contract, writes the given data to it and returns the memory location
 fn write_to_contract<A: BackendApi + 'static, S: Storage + 'static, Q: Querier + 'static>(
     data: &Environment<A, S, Q>,
@@ -634,6 +674,8 @@ mod tests {
                 "db_remove" => Function::new_typed(&mut store, |_a: u32| {}),
                 "db_scan" => Function::new_typed(&mut store, |_a: u32, _b: u32, _c: i32| -> u32 { 0 }),
                 "db_next" => Function::new_typed(&mut store, |_a: u32| -> u32 { 0 }),
+                "db_next_key" => Function::new_typed(&mut store, |_a: u32| -> u32 { 0 }),
+                "db_next_value" => Function::new_typed(&mut store, |_a: u32| -> u32 { 0 }),
                 "query_chain" => Function::new_typed(&mut store, |_a: u32| -> u32 { 0 }),
                 "addr_validate" => Function::new_typed(&mut store, |_a: u32| -> u32 { 0 }),
                 "addr_canonicalize" => Function::new_typed(&mut store, |_a: u32, _b: u32| -> u32 { 0 }),
@@ -2173,4 +2215,76 @@ mod tests {
             e => panic!("Unexpected error: {e:?}"),
         }
     }
+
+    #[test]
+    #[cfg(feature = "iterator")]
+    fn do_db_next_key_works() {
+        let api = MockApi::default();
+        let (fe, mut store, _instance) = make_instance(api);
+        let mut fe_mut = fe.into_mut(&mut store);
+
+        leave_default_data(&mut fe_mut);
+
+        let id = do_db_scan(fe_mut.as_mut(), 0, 0, Order::Ascending.into()).unwrap();
+
+        // Entry 1
+        let key_region_ptr = do_db_next_key(fe_mut.as_mut(), id).unwrap();
+        assert_eq!(force_read(&mut fe_mut, key_region_ptr), KEY1);
+
+        // Entry 2
+        let key_region_ptr = do_db_next_key(fe_mut.as_mut(), id).unwrap();
+        assert_eq!(force_read(&mut fe_mut, key_region_ptr), KEY2);
+
+        // End
+        let key_region_ptr: u32 = do_db_next_key(fe_mut.as_mut(), id).unwrap();
+        assert_eq!(key_region_ptr, 0);
+    }
+
+    #[test]
+    #[cfg(feature = "iterator")]
+    fn do_db_next_value_works() {
+        let api = MockApi::default();
+        let (fe, mut store, _instance) = make_instance(api);
+        let mut fe_mut = fe.into_mut(&mut store);
+
+        leave_default_data(&mut fe_mut);
+
+        let id = do_db_scan(fe_mut.as_mut(), 0, 0, Order::Ascending.into()).unwrap();
+
+        // Entry 1
+        let value_region_ptr = do_db_next_value(fe_mut.as_mut(), id).unwrap();
+        assert_eq!(force_read(&mut fe_mut, value_region_ptr), VALUE1);
+
+        // Entry 2
+        let value_region_ptr = do_db_next_value(fe_mut.as_mut(), id).unwrap();
+        assert_eq!(force_read(&mut fe_mut, value_region_ptr), VALUE2);
+
+        // End
+        let value_region_ptr = do_db_next_value(fe_mut.as_mut(), id).unwrap();
+        assert_eq!(value_region_ptr, 0);
+    }
+
+    #[test]
+    #[cfg(feature = "iterator")]
+    fn do_db_next_works_mixed() {
+        let api = MockApi::default();
+        let (fe, mut store, _instance) = make_instance(api);
+        let mut fe_mut = fe.into_mut(&mut store);
+
+        leave_default_data(&mut fe_mut);
+
+        let id = do_db_scan(fe_mut.as_mut(), 0, 0, Order::Ascending.into()).unwrap();
+
+        // Key 1
+        let key_region_ptr = do_db_next_key(fe_mut.as_mut(), id).unwrap();
+        assert_eq!(force_read(&mut fe_mut, key_region_ptr), KEY1);
+
+        // Value 2
+        let value_region_ptr = do_db_next_value(fe_mut.as_mut(), id).unwrap();
+        assert_eq!(force_read(&mut fe_mut, value_region_ptr), VALUE2);
+
+        // End
+        let kv_region_ptr = do_db_next(fe_mut.as_mut(), id).unwrap();
+        assert_eq!(force_read(&mut fe_mut, kv_region_ptr), b"\0\0\0\0\0\0\0\0");
+    }
 }

packages/vm/src/instance.rs

@@ -19,7 +19,7 @@ use crate::imports::{
     do_secp256k1_recover_pubkey, do_secp256k1_verify,
 };
 #[cfg(feature = "iterator")]
-use crate::imports::{do_db_next, do_db_scan};
+use crate::imports::{do_db_next, do_db_next_key, do_db_next_value, do_db_scan};
 use crate::memory::{read_region, write_region};
 use crate::size::Size;
 use crate::wasm_backend::{compile, make_compiling_engine};
@@ -110,7 +110,7 @@ where
         let mut import_obj = Imports::new();
         let mut env_imports = Exports::new();
 
-        // Reads the database entry at the given key into the the value.
+        // Reads the database entry at the given key into the value.
         // Returns 0 if key does not exist and pointer to result region otherwise.
         // Ownership of the key pointer is not transferred to the host.
         // Ownership of the value pointer is transferred to the contract.
@@ -237,6 +237,24 @@ where
             Function::new_typed_with_env(&mut store, &fe, do_db_next),
         );
 
+        // Get next key of iterator with ID `iterator_id`.
+        // Returns 0 if there are no more entries and pointer to result region otherwise.
+        // Ownership of the result region is transferred to the contract.
+        #[cfg(feature = "iterator")]
+        env_imports.insert(
+            "db_next_key",
+            Function::new_typed_with_env(&mut store, &fe, do_db_next_key),
+        );
+
+        // Get next value of iterator with ID `iterator_id`.
+        // Returns 0 if there are no more entries and pointer to result region otherwise.
+        // Ownership of the result region is transferred to the contract.
+        #[cfg(feature = "iterator")]
+        env_imports.insert(
+            "db_next_value",
+            Function::new_typed_with_env(&mut store, &fe, do_db_next_value),
+        );
+
         import_obj.register_namespace("env", env_imports);
 
         if let Some(extra_imports) = extra_imports {