Make compatible with risc0.

Introduce the CompactHeader type, which allows us to circumvent
informalsystems/tendermint-rs#1309.
Update to Tendermint 0.32, which allows removal of unnecessary data
copies.

Cargo.toml

@@ -7,7 +7,8 @@ edition = "2021"
 
 [dependencies]
 sovereign-sdk =  { git = "https://github.com/Sovereign-Labs/sovereign.git", rev = "5e43c3ee9b5785abdca33b21c86fd38dbd9285e0" }
-tendermint = "0.27"
+tendermint = "0.32"
+tendermint-proto = "0.32"
 
 prost = "0.11"
 prost-types = "0.11"
@@ -26,8 +27,6 @@ anyhow = "1.0.62"
 jsonrpsee = { version = "0.16.2", features = ["http-client"], optional = true }
 tracing = "0.1.37"
 
-
-#nmt-rs = { path = "../nmt-rs", features = ["serde", "borsh"] }
 nmt-rs = { git = "https://github.com/Sovereign-Labs/nmt-rs.git", rev = "aec2dcdc279b381162537f5b20ce43d1d46dc42f", features = ["serde", "borsh"] }
 
 [dev-dependencies]
@@ -37,8 +36,10 @@ postcard = { version = "1", features = ["use-std"]}
 prost-build = { version = "0.11" }
 
 
-#[patch.'https://github.com/Sovereign-Labs/sovereign.git']
-#sovereign-sdk = { path = "../sovereign/sdk" }
+[patch.crates-io]
+# Patch tendermint until the "0.32" release lands. We need 0.32 to avoid a lot of unnecessary data copying during header verification
+tendermint = { git = "https://github.com/informalsystems/tendermint-rs.git", rev = "e014de927abed7c5fcbf8186780a61b5c9c1e775" }
+tendermint-proto = { git = "https://github.com/informalsystems/tendermint-rs.git", rev = "e014de927abed7c5fcbf8186780a61b5c9c1e775" }
 
 [features]
 default = ["native"]

src/celestia.rs

@@ -1,4 +1,4 @@
-use std::ops::Range;
+use std::{cell::RefCell, ops::Range};
 
 use borsh::{BorshDeserialize, BorshSerialize};
 use nmt_rs::NamespacedHash;
@@ -8,8 +8,12 @@ use sovereign_sdk::core::traits::{
     AddressTrait as Address, BlockHeaderTrait as BlockHeader, CanonicalHash,
 };
 pub use tendermint::block::Header as TendermintHeader;
+use tendermint::{crypto::default::Sha256, merkle::simple_hash_from_byte_vectors, Hash};
+use tendermint_proto::Protobuf;
 use tracing::debug;
 
+pub use tendermint_proto::v0_34 as celestia_tm_version;
+
 const NAMESPACED_HASH_LEN: usize = 48;
 
 use crate::{
@@ -26,6 +30,158 @@ pub struct MarshalledDataAvailabilityHeader {
     pub column_roots: Vec<String>,
 }
 
+#[derive(PartialEq, Eq, Debug, Clone, Serialize, Deserialize, BorshDeserialize, BorshSerialize)]
+pub struct PartialBlockId {
+    pub hash: ProtobufHash,
+    pub part_set_header: Vec<u8>,
+}
+
+/// A partially serialized tendermint header. Only fields which are actually inspected by
+/// Jupiter are included in their raw form. Other fields are pre-encoded as protobufs.
+///
+/// This type was first introduced as a way to circumvent a bug in tendermint-rs which prevents
+/// a tendermint::block::Header from being deserialized in most formats except JSON. However
+/// it also provides a significant efficiency benefit over the standard tendermint type, which
+/// performs a complete protobuf serialization every time `.hash()` is called.
+#[derive(PartialEq, Eq, Debug, Clone, Serialize, Deserialize)]
+pub struct CompactHeader {
+    /// Header version
+    pub version: Vec<u8>,
+
+    /// Chain ID
+    pub chain_id: Vec<u8>,
+
+    /// Current block height
+    pub height: Vec<u8>,
+
+    /// Current timestamp
+    pub time: Vec<u8>,
+
+    /// Previous block info
+    pub last_block_id: Vec<u8>,
+
+    /// Commit from validators from the last block
+    pub last_commit_hash: Vec<u8>,
+
+    /// Merkle root of transaction hashes
+    pub data_hash: Option<ProtobufHash>,
+
+    /// Validators for the current block
+    pub validators_hash: Vec<u8>,
+
+    /// Validators for the next block
+    pub next_validators_hash: Vec<u8>,
+
+    /// Consensus params for the current block
+    pub consensus_hash: Vec<u8>,
+
+    /// State after txs from the previous block
+    pub app_hash: Vec<u8>,
+
+    /// Root hash of all results from the txs from the previous block
+    pub last_results_hash: Vec<u8>,
+
+    /// Hash of evidence included in the block
+    pub evidence_hash: Vec<u8>,
+
+    /// Original proposer of the block
+    pub proposer_address: Vec<u8>,
+}
+
+trait EncodeTm34 {
+    fn encode_to_tm34_protobuf(&self) -> Result<Vec<u8>, BoxError>;
+}
+
+impl From<TendermintHeader> for CompactHeader {
+    fn from(value: TendermintHeader) -> Self {
+        let data_hash = if let Some(h) = value.data_hash {
+            match h {
+                Hash::Sha256(value) => Some(ProtobufHash(value)),
+                Hash::None => None,
+            }
+        } else {
+            None
+        };
+        Self {
+            version: Protobuf::<celestia_tm_version::version::Consensus>::encode_vec(
+                &value.version,
+            )
+            .unwrap(),
+            chain_id: value.chain_id.encode_vec().unwrap(),
+            height: value.height.encode_vec().unwrap(),
+            time: value.time.encode_vec().unwrap(),
+            last_block_id: Protobuf::<celestia_tm_version::types::BlockId>::encode_vec(
+                &value.last_block_id.unwrap_or_default(),
+            )
+            .unwrap(),
+            last_commit_hash: value
+                .last_commit_hash
+                .unwrap_or_default()
+                .encode_vec()
+                .unwrap(),
+            data_hash,
+            validators_hash: value.validators_hash.encode_vec().unwrap(),
+            next_validators_hash: value.next_validators_hash.encode_vec().unwrap(),
+            consensus_hash: value.consensus_hash.encode_vec().unwrap(),
+            app_hash: value.app_hash.encode_vec().unwrap(),
+            last_results_hash: value
+                .last_results_hash
+                .unwrap_or_default()
+                .encode_vec()
+                .unwrap(),
+            evidence_hash: value
+                .evidence_hash
+                .unwrap_or_default()
+                .encode_vec()
+                .unwrap(),
+            proposer_address: value.proposer_address.encode_vec().unwrap(),
+        }
+    }
+}
+
+#[derive(Clone, PartialEq, Eq, Debug, Serialize, Deserialize, BorshDeserialize, BorshSerialize)]
+pub struct ProtobufHash(pub [u8; 32]);
+
+pub fn protobuf_encode(hash: &Option<ProtobufHash>) -> Vec<u8> {
+    match hash {
+        Some(ProtobufHash(value)) => prost::Message::encode_to_vec(&value.to_vec()),
+        None => prost::Message::encode_to_vec(&vec![]),
+    }
+}
+
+impl CompactHeader {
+    /// Hash this header
+    // TODO: this function can be made even more efficient. Rather than computing the block hash,
+    // we could provide the hash as a non-deterministic input and simply verify the correctness of the
+    // fields that we care about.
+    pub fn hash(&self) -> Hash {
+        // Note that if there is an encoding problem this will
+        // panic (as the golang code would):
+        // https://github.com/tendermint/tendermint/blob/134fe2896275bb926b49743c1e25493f6b24cc31/types/block.go#L393
+        // https://github.com/tendermint/tendermint/blob/134fe2896275bb926b49743c1e25493f6b24cc31/types/encoding_helper.go#L9:6
+
+        let encoded_data_hash = protobuf_encode(&self.data_hash);
+        let fields_bytes = vec![
+            &self.version,
+            &self.chain_id,
+            &self.height,
+            &self.time,
+            &self.last_block_id,
+            &self.last_commit_hash,
+            &encoded_data_hash,
+            &self.validators_hash,
+            &self.next_validators_hash,
+            &self.consensus_hash,
+            &self.app_hash,
+            &self.last_results_hash,
+            &self.evidence_hash,
+            &self.proposer_address,
+        ];
+
+        Hash::Sha256(simple_hash_from_byte_vectors::<Sha256>(&fields_bytes))
+    }
+}
+
 #[derive(
     PartialEq, Debug, Clone, Deserialize, serde::Serialize, BorshDeserialize, BorshSerialize,
 )]
@@ -60,6 +216,8 @@ impl TryFrom<MarshalledDataAvailabilityHeader> for DataAvailabilityHeader {
     }
 }
 
+/// The response from the celestia `/header` endpoint. Must be converted to a
+/// [`CelestiaHeader`] before use.
 #[derive(Deserialize, Serialize, PartialEq, Debug, Clone)]
 pub struct CelestiaHeaderResponse {
     pub header: tendermint::block::Header,
@@ -75,10 +233,20 @@ pub struct NamespacedSharesResponse {
 #[derive(Debug, PartialEq, Clone, Deserialize, serde::Serialize)]
 pub struct CelestiaHeader {
     pub dah: DataAvailabilityHeader,
-    pub header: tendermint::block::Header,
+    pub header: CompactHeader,
+    #[serde(skip)]
+    cached_prev_hash: RefCell<Option<TmHash>>,
 }
 
 impl CelestiaHeader {
+    pub fn new(dah: DataAvailabilityHeader, header: CompactHeader) -> Self {
+        Self {
+            dah,
+            header,
+            cached_prev_hash: RefCell::new(None),
+        }
+    }
+
     pub fn square_size(&self) -> usize {
         self.dah.row_roots.len()
     }
@@ -101,13 +269,22 @@ pub struct BlobWithSender {
 impl BlockHeader for CelestiaHeader {
     type Hash = TmHash;
 
-    fn prev_hash(&self) -> &Self::Hash {
-        self.header
-            .last_block_id
-            .as_ref()
+    fn prev_hash(&self) -> Self::Hash {
+        // Try to return the cached value
+        if let Some(hash) = self.cached_prev_hash.borrow().as_ref() {
+            return hash.clone();
+        }
+        // If we reach this point, we know that the cach is empty - so there can't be any outstanding references to its value.
+        // That means its safe to borrow the cache mutably and populate it.
+        let mut cached_hash = self.cached_prev_hash.borrow_mut();
+        let hash =
+            <tendermint::block::Id as Protobuf<celestia_tm_version::types::BlockId>>::decode(
+                self.header.last_block_id.as_ref(),
+            )
             .expect("must not call prev_hash on block with no predecessor")
-            .hash
-            .as_ref()
+            .hash;
+        *cached_hash = Some(TmHash(hash.clone()));
+        TmHash(hash)
     }
 }
 
@@ -218,29 +395,32 @@ fn next_pfb(mut data: &mut BlobRefIterator) -> Result<(MsgPayForBlobs, TxPositio
     ))
 }
 
-// #[cfg(test)]
-// mod tests {
-//     use crate::CelestiaHeaderResponse;
-
-//     const HEADER_RESPONSE_JSON: &[u8] = include_bytes!("./header_response.json");
-//     #[test]
-//     fn test_tm_header_serde() {
-//         // TODO: Find breakage here. Then, re-enable this test.
-//         let original_header: CelestiaHeaderResponse =
-//             serde_json::from_slice(HEADER_RESPONSE_JSON).unwrap();
-
-//         let header: tendermint::block::Header = original_header.header;
-
-//         let serialized_header: Vec<u8> = serde_json::to_vec_pretty(&header).unwrap();
-//         println!(
-//             "header: {}",
-//             String::from_utf8(serialized_header.clone()).unwrap()
-//         );
-//         // let deserialized_header: tendermint::block::Header =
-//         //     serde_json::from_slice(&serialized_header).unwrap();
-
-//         let serialized_header = postcard::to_stdvec(&header).unwrap();
-//         let deserialized_header: tendermint::block::Header =
-//             postcard::from_bytes(&serialized_header).unwrap();
-//     }
-// }
+#[cfg(test)]
+mod tests {
+    use crate::{CelestiaHeaderResponse, CompactHeader};
+
+    const HEADER_RESPONSE_JSON: &[u8] = include_bytes!("./header_response.json");
+
+    #[test]
+    fn test_compact_header_serde() {
+        let original_header: CelestiaHeaderResponse =
+            serde_json::from_slice(HEADER_RESPONSE_JSON).unwrap();
+
+        let header: CompactHeader = original_header.header.into();
+
+        let serialized_header = postcard::to_stdvec(&header).unwrap();
+        let deserialized_header: CompactHeader = postcard::from_bytes(&serialized_header).unwrap();
+        assert_eq!(deserialized_header, header)
+    }
+
+    #[test]
+    fn test_compact_header_hash() {
+        let original_header: CelestiaHeaderResponse =
+            serde_json::from_slice(HEADER_RESPONSE_JSON).unwrap();
+
+        let tm_header = original_header.header.clone();
+        let compact_header: CompactHeader = original_header.header.into();
+
+        assert_eq!(tm_header.hash(), compact_header.hash());
+    }
+}

src/da_service.rs

@@ -140,10 +140,7 @@ impl DaService for CelestiaService {
             }
 
             let filtered_block = FilteredCelestiaBlock {
-                header: CelestiaHeader {
-                    header: unmarshalled_header.header,
-                    dah,
-                },
+                header: CelestiaHeader::new(dah, unmarshalled_header.header.into()),
                 rollup_data: rollup_shares,
                 relevant_pfbs: pfd_map,
                 rollup_rows,

src/share_commit.rs

@@ -1,6 +1,6 @@
 use crate::shares::{self, Share};
 
-use tendermint::merkle::simple_hash_from_byte_vectors;
+use tendermint::{crypto::default::Sha256, merkle::simple_hash_from_byte_vectors};
 
 // /// Calculates the size of the smallest square that could be used to commit
 // /// to this message, following Celestia's "non-interactive default rules"
@@ -54,12 +54,7 @@ pub fn recreate_commitment(
         }
         subtree_roots.push(tree.root());
     }
-    let h = simple_hash_from_byte_vectors(
-        subtree_roots
-            .into_iter()
-            .map(|x| x.as_ref().to_vec())
-            .collect(),
-    );
+    let h = simple_hash_from_byte_vectors::<Sha256>(&subtree_roots);
     Ok(h)
 }
 

src/types.rs

@@ -8,7 +8,7 @@ use sovereign_sdk::{
     services::da::SlotData,
     Bytes,
 };
-use tendermint::merkle;
+use tendermint::{crypto::default::Sha256, merkle};
 
 use crate::{
     pfb::MsgPayForBlobs,
@@ -144,16 +144,14 @@ impl CelestiaHeader {
     pub fn validate_dah(&self) -> Result<(), ValidationError> {
         let rows_iter = self.dah.row_roots.iter();
         let cols_iter = self.dah.column_roots.iter();
-        let byte_vecs = rows_iter
-            .chain(cols_iter)
-            .map(|hash| hash.0.to_vec())
-            .collect();
-        let root = merkle::simple_hash_from_byte_vectors(byte_vecs);
+        let byte_vecs: Vec<&NamespacedHash> = rows_iter.chain(cols_iter).collect();
+        let root = merkle::simple_hash_from_byte_vectors::<Sha256>(&byte_vecs);
         let data_hash = self
             .header
             .data_hash
+            .as_ref()
             .ok_or(ValidationError::MissingDataHash)?;
-        if &root != <tendermint::Hash as AsRef<[u8]>>::as_ref(&data_hash) {
+        if &root != &data_hash.0 {
             return Err(ValidationError::InvalidDataRoot);
         }
         Ok(())