solana-ibc: store next sequences in a single object

solana/solana-ibc/programs/solana-ibc/src/execution_context.rs

@@ -377,78 +377,41 @@ impl ExecutionContext for SolanaIbcStorage<'_, '_> {
 
     fn store_next_sequence_send(
         &mut self,
-        seq_send_path: &SeqSendPath,
+        path: &SeqSendPath,
         seq: Sequence,
     ) -> Result {
-        msg!(
-            "store_next_sequence_send: path: {}, seq: {:?}",
-            seq_send_path,
-            seq
-        );
-        let seq_send_key =
-            (seq_send_path.0.to_string(), seq_send_path.1.to_string());
-
-        let next_seq_send_trie_key = TrieKey::from(seq_send_path);
-        let trie = self.trie.as_mut().unwrap();
-        let seq_in_u64: u64 = seq.into();
-        let seq_in_bytes = seq_in_u64.to_be_bytes();
-
-        trie.set(
-            &next_seq_send_trie_key,
-            &lib::hash::CryptoHash::digest(&seq_in_bytes),
+        msg!("store_next_sequence_send: path: {path}, seq: {seq}");
+        self.store_next_sequence(
+            path.into(),
+            super::SequenceTripleIdx::Send,
+            seq,
         )
-        .unwrap();
-
-        self.next_sequence_send.insert(seq_send_key, u64::from(seq));
-        Ok(())
     }
 
     fn store_next_sequence_recv(
         &mut self,
-        seq_recv_path: &SeqRecvPath,
+        path: &SeqRecvPath,
         seq: Sequence,
     ) -> Result {
-        msg!(
-            "store_next_sequence_recv: path: {}, seq: {:?}",
-            seq_recv_path,
-            seq
-        );
-        let seq_recv_key =
-            (seq_recv_path.0.to_string(), seq_recv_path.1.to_string());
-        let next_seq_recv_trie_key = TrieKey::from(seq_recv_path);
-        let trie = self.trie.as_mut().unwrap();
-        let seq_in_u64: u64 = seq.into();
-        let seq_in_bytes = seq_in_u64.to_be_bytes();
-
-        trie.set(
-            &next_seq_recv_trie_key,
-            &lib::hash::CryptoHash::digest(&seq_in_bytes),
+        msg!("store_next_sequence_recv: path: {path}, seq: {seq}");
+        self.store_next_sequence(
+            path.into(),
+            super::SequenceTripleIdx::Recv,
+            seq,
         )
-        .unwrap();
-        self.next_sequence_recv.insert(seq_recv_key, u64::from(seq));
-        Ok(())
     }
 
     fn store_next_sequence_ack(
         &mut self,
-        seq_ack_path: &SeqAckPath,
+        path: &SeqAckPath,
         seq: Sequence,
     ) -> Result {
-        msg!("store_next_sequence_ack: path: {}, seq: {:?}", seq_ack_path, seq);
-        let seq_ack_key =
-            (seq_ack_path.0.to_string(), seq_ack_path.1.to_string());
-        let next_seq_ack_trie_key = TrieKey::from(seq_ack_path);
-        let trie = self.trie.as_mut().unwrap();
-        let seq_in_u64: u64 = seq.into();
-        let seq_in_bytes = seq_in_u64.to_be_bytes();
-
-        trie.set(
-            &next_seq_ack_trie_key,
-            &lib::hash::CryptoHash::digest(&seq_in_bytes),
+        msg!("store_next_sequence_ack: path: {path}, seq: {seq}");
+        self.store_next_sequence(
+            path.into(),
+            super::SequenceTripleIdx::Ack,
+            seq,
         )
-        .unwrap();
-        self.next_sequence_ack.insert(seq_ack_key, u64::from(seq));
-        Ok(())
     }
 
     fn increase_channel_counter(&mut self) -> Result {
@@ -478,6 +441,26 @@ impl ExecutionContext for SolanaIbcStorage<'_, '_> {
     fn get_client_execution_context(&mut self) -> &mut Self::E { self }
 }
 
+impl SolanaIbcStorage<'_, '_> {
+    fn store_next_sequence(
+        &mut self,
+        path: crate::trie_key::SequencePath<'_>,
+        index: super::SequenceTripleIdx,
+        seq: Sequence,
+    ) -> Result {
+        let map_key = (path.port_id.to_string(), path.channel_id.to_string());
+        let triple = self.next_sequence.entry(map_key).or_default();
+        triple.set(index, seq);
+
+        let trie_key = TrieKey::from(path);
+        let trie = self.trie.as_mut().unwrap();
+        trie.set(&trie_key, &triple.to_hash()).unwrap();
+
+        Ok(())
+    }
+}
+
+
 fn record_packet_sequence(
     hash_map: &mut BTreeMap<(InnerPortId, InnerChannelId), Vec<InnerSequence>>,
     port_id: &PortId,

solana/solana-ibc/programs/solana-ibc/src/lib.rs

@@ -6,6 +6,7 @@ use std::collections::BTreeMap;
 
 use anchor_lang::prelude::*;
 use borsh::{BorshDeserialize, BorshSerialize};
+use ibc::core::ics04_channel::packet::Sequence;
 use ibc::core::ics24_host::identifier::PortId;
 use ibc::core::router::{Module, ModuleId, Router};
 use module_holder::ModuleHolder;
@@ -90,9 +91,7 @@ pub mod solana_ibc {
             connection_to_client: solana_ibc_store.connection_to_client.clone(),
             port_channel_id_set: solana_ibc_store.port_channel_id_set.clone(),
             channel_counter: solana_ibc_store.channel_counter,
-            next_sequence_send: solana_ibc_store.next_sequence_send.clone(),
-            next_sequence_recv: solana_ibc_store.next_sequence_recv.clone(),
-            next_sequence_ack: solana_ibc_store.next_sequence_ack.clone(),
+            next_sequence: solana_ibc_store.next_sequence.clone(),
             packet_commitment_sequence_sets: solana_ibc_store
                 .packet_commitment_sequence_sets
                 .clone(),
@@ -129,9 +128,7 @@ pub mod solana_ibc {
             connection_to_client: solana_ibc_store.connection_to_client.clone(),
             port_channel_id_set: solana_ibc_store.port_channel_id_set.clone(),
             channel_counter: solana_ibc_store.channel_counter,
-            next_sequence_send: solana_ibc_store.next_sequence_send.clone(),
-            next_sequence_recv: solana_ibc_store.next_sequence_recv.clone(),
-            next_sequence_ack: solana_ibc_store.next_sequence_ack.clone(),
+            next_sequence: solana_ibc_store.next_sequence.clone(),
             packet_commitment_sequence_sets: solana_ibc_store
                 .packet_commitment_sequence_sets
                 .clone(),
@@ -192,12 +189,8 @@ pub mod solana_ibc {
         solana_ibc_store.port_channel_id_set =
             solana_real_storage.port_channel_id_set.clone();
         solana_ibc_store.channel_counter = solana_real_storage.channel_counter;
-        solana_ibc_store.next_sequence_send =
-            solana_real_storage.next_sequence_send.clone();
-        solana_ibc_store.next_sequence_recv =
-            solana_real_storage.next_sequence_recv.clone();
-        solana_ibc_store.next_sequence_ack =
-            solana_real_storage.next_sequence_ack.clone();
+        solana_ibc_store.next_sequence =
+            solana_real_storage.next_sequence.clone();
         solana_ibc_store.packet_commitment_sequence_sets =
             solana_real_storage.packet_commitment_sequence_sets.clone();
         solana_ibc_store.packet_receipt_sequence_sets =
@@ -252,6 +245,59 @@ pub type InnerConnectionEnd = String; // Serialized
 pub type InnerChannelEnd = String; // Serialized
 pub type InnerConsensusState = String; // Serialized
 
+/// A triple of send, receive and acknowledge sequences.
+#[derive(
+    Clone,
+    Debug,
+    Default,
+    PartialEq,
+    Eq,
+    borsh::BorshSerialize,
+    borsh::BorshDeserialize,
+)]
+pub struct InnerSequenceTriple {
+    sequences: [u64; 3],
+    mask: u8,
+}
+
+#[derive(Clone, Copy)]
+pub enum SequenceTripleIdx {
+    Send = 0,
+    Recv = 1,
+    Ack = 2,
+}
+
+impl InnerSequenceTriple {
+    /// Returns sequence at given index or `None` if it wasn’t set yet.
+    pub fn get(&self, idx: SequenceTripleIdx) -> Option<Sequence> {
+        if self.mask & (1 << (idx as u32)) == 1 {
+            Some(Sequence::from(self.sequences[idx as usize]))
+        } else {
+            None
+        }
+    }
+
+    /// Sets sequence at given index.
+    pub fn set(&mut self, idx: SequenceTripleIdx, seq: Sequence) {
+        self.sequences[idx as usize] = u64::from(seq);
+        self.mask |= 1 << (idx as u32)
+    }
+
+    /// Encodes the object as a `CryptoHash` so it can be stored in the trie
+    /// directly.
+    pub fn to_hash(&self) -> lib::hash::CryptoHash {
+        let mut hash = lib::hash::CryptoHash::default();
+        let (first, tail) = stdx::split_array_mut::<8, 24, 32>(&mut hash.0);
+        let (second, tail) = stdx::split_array_mut::<8, 16, 24>(tail);
+        let (third, tail) = stdx::split_array_mut::<8, 8, 16>(tail);
+        *first = self.sequences[0].to_be_bytes();
+        *second = self.sequences[1].to_be_bytes();
+        *third = self.sequences[2].to_be_bytes();
+        tail[0] = self.mask;
+        hash
+    }
+}
+
 #[account]
 #[derive(Debug)]
 /// All the structs from IBC are stored as String since they dont implement AnchorSerialize and AnchorDeserialize
@@ -283,12 +329,15 @@ pub struct SolanaIbcStorageTemp {
     /// The port and channel id tuples of the channels.
     pub port_channel_id_set: Vec<(InnerPortId, InnerChannelId)>,
     pub channel_counter: u64,
-    pub next_sequence_send:
-        BTreeMap<(InnerPortId, InnerChannelId), InnerSequence>,
-    pub next_sequence_recv:
-        BTreeMap<(InnerPortId, InnerChannelId), InnerSequence>,
-    pub next_sequence_ack:
-        BTreeMap<(InnerPortId, InnerChannelId), InnerSequence>,
+
+    /// Next send, receive and ack sequence for given (port, channel).
+    ///
+    /// We’re storing all three sequences in a single object to reduce amount of
+    /// different maps we need to maintain.  This saves us on the amount of
+    /// trie nodes we need to maintain.
+    pub next_sequence:
+        BTreeMap<(InnerPortId, InnerChannelId), InnerSequenceTriple>,
+
     /// The sequence numbers of the packet commitments.
     pub packet_commitment_sequence_sets:
         BTreeMap<(InnerPortId, InnerChannelId), Vec<InnerSequence>>,
@@ -331,12 +380,15 @@ pub struct SolanaIbcStorage<'a, 'b> {
     /// The port and channel id tuples of the channels.
     pub port_channel_id_set: Vec<(InnerPortId, InnerChannelId)>,
     pub channel_counter: u64,
-    pub next_sequence_send:
-        BTreeMap<(InnerPortId, InnerChannelId), InnerSequence>,
-    pub next_sequence_recv:
-        BTreeMap<(InnerPortId, InnerChannelId), InnerSequence>,
-    pub next_sequence_ack:
-        BTreeMap<(InnerPortId, InnerChannelId), InnerSequence>,
+
+    /// Next send, receive and ack sequence for given (port, channel).
+    ///
+    /// We’re storing all three sequences in a single object to reduce amount of
+    /// different maps we need to maintain.  This saves us on the amount of
+    /// trie nodes we need to maintain.
+    pub next_sequence:
+        BTreeMap<(InnerPortId, InnerChannelId), InnerSequenceTriple>,
+
     /// The sequence numbers of the packet commitments.
     pub packet_commitment_sequence_sets:
         BTreeMap<(InnerPortId, InnerChannelId), Vec<InnerSequence>>,

solana/solana-ibc/programs/solana-ibc/src/trie_key.rs

@@ -40,6 +40,12 @@ use super::{CHANNEL_ID_PREFIX, CONNECTION_ID_PREFIX};
 // to avoid heap allocations.
 pub struct TrieKey(Vec<u8>);
 
+/// A path for next send, receive and ack sequence paths.
+pub struct SequencePath<'a> {
+    pub port_id: &'a PortId,
+    pub channel_id: &'a ChannelId,
+}
+
 /// Constructs a new [`TrieKey`] by concatenating key components.
 ///
 /// The first argument to the macro is a [`Tag`] object.  Remaining must
@@ -114,21 +120,31 @@ impl From<&ChannelEndPath> for TrieKey {
     }
 }
 
-impl From<&SeqSendPath> for TrieKey {
-    fn from(path: &SeqSendPath) -> Self {
-        Self::from_channel_path(Tag::NextSequenceSend, &path.0, &path.1)
+impl<'a> From<&'a SeqSendPath> for SequencePath<'a> {
+    fn from(path: &'a SeqSendPath) -> Self {
+        Self { port_id: &path.0, channel_id: &path.1 }
+    }
+}
+
+impl<'a> From<&'a SeqRecvPath> for SequencePath<'a> {
+    fn from(path: &'a SeqRecvPath) -> Self {
+        Self { port_id: &path.0, channel_id: &path.1 }
     }
 }
 
-impl From<&SeqRecvPath> for TrieKey {
-    fn from(path: &SeqRecvPath) -> Self {
-        Self::from_channel_path(Tag::NextSequenceRecv, &path.0, &path.1)
+impl<'a> From<&'a SeqAckPath> for SequencePath<'a> {
+    fn from(path: &'a SeqAckPath) -> Self {
+        Self { port_id: &path.0, channel_id: &path.1 }
     }
 }
 
-impl From<&SeqAckPath> for TrieKey {
-    fn from(path: &SeqAckPath) -> Self {
-        Self::from_channel_path(Tag::NextSequenceAck, &path.0, &path.1)
+impl From<SequencePath<'_>> for TrieKey {
+    fn from(path: SequencePath<'_>) -> Self {
+        Self::from_channel_path(
+            Tag::NextSequence,
+            path.port_id,
+            path.channel_id,
+        )
     }
 }
 
@@ -169,16 +185,14 @@ impl From<&AckPath> for TrieKey {
 /// for different objects stored in the trie.
 #[repr(u8)]
 enum Tag {
-    ClientState = 1,
-    ConsensusState = 2,
-    Connection = 3,
-    ChannelEnd = 4,
-    NextSequenceSend = 5,
-    NextSequenceRecv = 6,
-    NextSequenceAck = 7,
-    Commitment = 8,
-    Receipt = 9,
-    Ack = 10,
+    ClientState = 0,
+    ConsensusState = 1,
+    Connection = 2,
+    ChannelEnd = 3,
+    NextSequence = 4,
+    Commitment = 5,
+    Receipt = 6,
+    Ack = 8,
 }
 
 /// Component of a [`TrieKey`].