fix(cw-context): encoding, decoding and iteration of ConsensusState heights

Cargo.toml

@@ -119,6 +119,7 @@ tendermint-testgen               = { version = "0.34.0", default-features = fals
 cosmwasm-schema    = { version = "1.5.2" }
 cosmwasm-std       = { version = "1.5.2" }
 cosmwasm-vm        = { version = "1.5.2" }
+cw-storage-plus    = { version = "1.2.0" }
 
 # parity dependencies
 parity-scale-codec = { version = "3.6.5", default-features = false, features = ["full"] }

ibc-clients/cw-context/Cargo.toml

@@ -28,6 +28,7 @@ ibc-client-wasm-types = { workspace = true, features = ["cosmwasm"] }
 # cosmwasm dependencies
 cosmwasm-schema       = { workspace = true }
 cosmwasm-std          = { workspace = true }
+cw-storage-plus       = { workspace = true }
 
 [features]
 default = ["std"]

ibc-clients/cw-context/src/context/client_ctx.rs

@@ -7,11 +7,11 @@
 use ibc_core::client::types::Height;
 use ibc_core::handler::types::error::ContextError;
 use ibc_core::host::types::identifiers::ClientId;
-use ibc_core::host::types::path::{iteration_key, ClientConsensusStatePath, ClientStatePath};
+use ibc_core::host::types::path::{ClientConsensusStatePath, ClientStatePath};
 use ibc_core::primitives::proto::{Any, Protobuf};
 use ibc_core::primitives::Timestamp;
 
-use super::Context;
+use super::{Context, StorageMut};
 use crate::api::ClientType;
 use crate::utils::AnyCodec;
 
@@ -154,11 +154,15 @@
 
         self.insert(prefixed_height_key, revision_height_vec);
 
-        let iteration_key = iteration_key(height.revision_number(), height.revision_height());
-
-        let height_vec = height.to_string().into_bytes();
-
-        self.insert(iteration_key, height_vec);
+        super::CONSENSUS_STATE_HEIGHT_MAP
+            .save(
+                self.storage_mut(),
+                (height.revision_number(), height.revision_height()),
+                &Default::default(),
+            )
+            .map_err(|e| ClientError::Other {
+                description: e.to_string(),
+            })?;
 
         Ok(())
     }
@@ -180,9 +184,10 @@
 
         self.remove(prefixed_height_key);
 
-        let iteration_key = iteration_key(height.revision_number(), height.revision_height());
-
-        self.remove(iteration_key);
+        super::CONSENSUS_STATE_HEIGHT_MAP.remove(
+            self.storage_mut(),
+            (height.revision_number(), height.revision_height()),
+        );
 
         Ok(())
     }

ibc-clients/cw-context/src/context/mod.rs

@@ -18,10 +18,15 @@
 
 use crate::api::ClientType;
 use crate::types::{ContractError, GenesisMetadata, HeightTravel, MigrationPrefix};
-use crate::utils::{parse_height, AnyCodec};
+use crate::utils::AnyCodec;
+use cosmwasm_std::Empty;
+use cw_storage_plus::{Bound, Map};
 
 type Checksum = Vec<u8>;
 
+pub const CONSENSUS_STATE_HEIGHT_MAP: Map<'_, (u64, u64), Empty> =
+    Map::new(ITERATE_CONSENSUS_STATE_PREFIX);
+
 /// Context is a wrapper around the deps and env that gives access to the
 /// methods under the ibc-rs Validation and Execution traits.
 pub struct Context<'a, C: ClientType<'a>> {
@@ -130,13 +135,11 @@
 
     /// Returns the storage of the context.
     pub fn get_heights(&self) -> Result<Vec<Height>, ClientError> {
-        let iterator = self.storage_ref().range(None, None, Order::Ascending);
-
-        let heights: Vec<_> = iterator
-            .filter_map(|(_, value)| parse_height(value).transpose())
-            .collect::<Result<_, _>>()?;
-
-        Ok(heights)
+        CONSENSUS_STATE_HEIGHT_MAP
+            .keys(self.storage_ref(), None, None, Order::Ascending)
+            .flatten()
+            .map(|(rev_number, rev_height)| Height::new(rev_number, rev_height))
+            .collect()
     }
 
     /// Searches for either the earliest next or latest previous height based on
@@ -146,22 +149,32 @@
         height: &Height,
         travel: HeightTravel,
     ) -> Result<Option<Height>, ClientError> {
-        let iteration_key = iteration_key(height.revision_number(), height.revision_height());
-
-        let mut iterator = match travel {
-            HeightTravel::Prev => {
-                self.storage_ref()
-                    .range(None, Some(&iteration_key), Order::Descending)
-            }
-            HeightTravel::Next => {
-                self.storage_ref()
-                    .range(Some(&iteration_key), None, Order::Ascending)
-            }
+        let iterator = match travel {
+            HeightTravel::Prev => CONSENSUS_STATE_HEIGHT_MAP.range(
+                self.storage_ref(),
+                None,
+                Some(Bound::exclusive((
+                    height.revision_number(),
+                    height.revision_height(),
+                ))),
+                Order::Descending,
+            ),
+            HeightTravel::Next => CONSENSUS_STATE_HEIGHT_MAP.range(
+                self.storage_ref(),
+                Some(Bound::exclusive((
+                    height.revision_number(),
+                    height.revision_height(),
+                ))),
+                None,
+                Order::Ascending,
+            ),
         };
 
         iterator
+            .flatten()
+            .map(|((rev_number, rev_height), _)| Height::new(rev_number, rev_height))
             .next()
-            .map_or(Ok(None), |(_, height)| parse_height(height))
+            .transpose()
     }
 
     /// Returns the key for the client update time.
@@ -191,38 +204,36 @@
     pub fn get_metadata(&self) -> Result<Option<Vec<GenesisMetadata>>, ContractError> {
         let mut metadata = Vec::<GenesisMetadata>::new();
 
-        let start_key = ITERATE_CONSENSUS_STATE_PREFIX.to_string().into_bytes();
-
-        let iterator = self
-            .storage_ref()
-            .range(Some(&start_key), None, Order::Ascending);
-
-        for (_, encoded_height) in iterator {
-            let height = parse_height(encoded_height)?;
-
-            match height {
-                Some(height) => {
-                    let processed_height_key = self.client_update_height_key(&height);
-                    metadata.push(GenesisMetadata {
-                        key: processed_height_key.clone(),
-                        value: self.retrieve(&processed_height_key)?,
-                    });
-                    let processed_time_key = self.client_update_time_key(&height);
-                    metadata.push(GenesisMetadata {
-                        key: processed_time_key.clone(),
-                        value: self.retrieve(&processed_time_key)?,
-                    });
-                }
-                None => continue,
-            }
+        let iterator = CONSENSUS_STATE_HEIGHT_MAP
+            .keys(self.storage_ref(), None, None, Order::Ascending)
+            .flatten();
+
+        for (rev_number, rev_height) in iterator {
+            let height = Height::new(rev_number, rev_height)?;
+
+            let processed_height_key = self.client_update_height_key(&height);
+            metadata.push(GenesisMetadata {
+                key: processed_height_key.clone(),
+                value: self.retrieve(&processed_height_key)?,
+            });
+            let processed_time_key = self.client_update_time_key(&height);
+            metadata.push(GenesisMetadata {
+                key: processed_time_key.clone(),
+                value: self.retrieve(&processed_time_key)?,
+            });
         }
 
-        let iterator = self
-            .storage_ref()
-            .range(Some(&start_key), None, Order::Ascending);
+        let iterator = CONSENSUS_STATE_HEIGHT_MAP
+            .keys(self.storage_ref(), None, None, Order::Ascending)
+            .flatten();
+
+        for (rev_number, rev_height) in iterator {
+            let height = Height::new(rev_number, rev_height)?;
 
-        for (key, height) in iterator {
-            metadata.push(GenesisMetadata { key, value: height });
+            metadata.push(GenesisMetadata {
+                key: iteration_key(rev_number, rev_height),
+                value: height.encode_vec(),
+            });
         }
 
         Ok(Some(metadata))

ibc-clients/cw-context/src/lib.rs

@@ -18,8 +18,6 @@
 )]
 #![forbid(unsafe_code)]
 
-extern crate alloc;
-
 pub mod api;
 pub mod context;
 pub mod handlers;

ibc-clients/cw-context/src/utils/mod.rs

@@ -1,25 +1,3 @@
 mod codec;
 
 pub use codec::*;
-use ibc_core::client::types::error::ClientError;
-use ibc_core::client::types::{Height, HeightError};
-
-/// Decodes a `Height` from a UTF-8 encoded byte array.
-pub fn parse_height(encoded_height: Vec<u8>) -> Result<Option<Height>, ClientError> {
-    let height_str = match alloc::str::from_utf8(encoded_height.as_slice()) {
-        Ok(s) => s,
-        // In cases where the height is unavailable, the encoded representation
-        // might not be valid UTF-8, resulting in an invalid string. In such
-        // instances, we return None.
-        Err(_) => return Ok(None),
-    };
-    match Height::try_from(height_str) {
-        Ok(height) => Ok(Some(height)),
-        // This is a valid case, as the key may contain other data. We just skip
-        // it.
-        Err(HeightError::InvalidFormat { .. }) => Ok(None),
-        Err(e) => Err(ClientError::Other {
-            description: e.to_string(),
-        }),
-    }
-}