fix: update TipsetKey save and load logic to match Lotus

Cargo.toml

@@ -259,6 +259,7 @@ quickcheck = "1"
 quickcheck_macros = "1"
 ra_ap_syntax = "0.0.323"
 regex-automata = "0.4"
+rstest = "0.26"
 serial_test = "3"
 syn = { version = "2", default-features = false, features = ["full", "parsing", "visit", "printing", "extra-traits"] }
 tokio-test = "0.4"

src/blocks/tipset.rs

@@ -3,7 +3,7 @@
 
 use std::{
     fmt,
-    sync::{Arc, OnceLock},
+    sync::{Arc, LazyLock, OnceLock},
 };
 
 use super::{Block, CachingBlockHeader, RawBlockHeader, Ticket};
@@ -12,7 +12,12 @@ use crate::{
     cid_collections::SmallCidNonEmptyVec,
     networks::{calibnet, mainnet},
     shim::clock::ChainEpoch,
-    utils::{cid::CidCborExt, get_size::nunny_vec_heap_size_helper},
+    utils::{
+        cid::CidCborExt,
+        db::{CborStoreExt, car_stream::CarBlock},
+        get_size::nunny_vec_heap_size_helper,
+        multihash::MultihashCode,
+    },
 };
 use ahash::HashMap;
 use anyhow::Context as _;
@@ -21,6 +26,7 @@ use fvm_ipld_blockstore::Blockstore;
 use fvm_ipld_encoding::CborStore;
 use get_size2::GetSize;
 use itertools::Itertools as _;
+use multihash_derive::MultihashDigest as _;
 use num::BigInt;
 use nunny::{Vec as NonEmpty, vec as nonempty};
 use serde::{Deserialize, Serialize};
@@ -43,13 +49,18 @@ use thiserror::Error;
     GetSize,
     derive_more::IntoIterator,
 )]
-#[cfg_attr(test, derive(derive_quickcheck_arbitrary::Arbitrary))]
 pub struct TipsetKey(#[into_iterator(owned, ref)] SmallCidNonEmptyVec);
 
 impl TipsetKey {
     // Special encoding to match Lotus.
     pub fn cid(&self) -> anyhow::Result<Cid> {
-        Ok(Cid::from_cbor_blake2b256(&self.bytes())?)
+        Ok(self.car_block()?.cid)
+    }
+
+    pub fn car_block(&self) -> anyhow::Result<CarBlock> {
+        let data = fvm_ipld_encoding::to_vec(&self.bytes())?;
+        let cid = Cid::from_cbor_encoded_raw_bytes_blake2b256(&data);
+        Ok(CarBlock { cid, data })
     }
 
     /// Returns `true` if the tipset key contains the given CID.
@@ -109,6 +120,37 @@ impl TipsetKey {
     pub fn bytes(&self) -> fvm_ipld_encoding::RawBytes {
         fvm_ipld_encoding::RawBytes::new(self.iter().flat_map(|cid| cid.to_bytes()).collect())
     }
+
+    /// Construct from bytes representation
+    pub fn from_bytes(bytes: fvm_ipld_encoding::RawBytes) -> anyhow::Result<Self> {
+        static BLOCK_HEADER_CID_LEN: LazyLock<usize> = LazyLock::new(|| {
+            let buf = [0_u8; 256];
+            let cid = Cid::new_v1(
+                fvm_ipld_encoding::DAG_CBOR,
+                MultihashCode::Blake2b256.digest(&buf),
+            );
+            cid.encoded_len()
+        });
+
+        let cids: Vec<Cid> = Vec::<u8>::from(bytes)
+            .chunks(*BLOCK_HEADER_CID_LEN)
+            .map(Cid::read_bytes)
+            .try_collect()?;
+
+        Ok(nunny::Vec::new(cids)
+            .map_err(|_| anyhow::anyhow!("tipset key cannot be empty"))?
+            .into())
+    }
+
+    /// Save tipset key to block store
+    pub fn save(&self, bs: &impl Blockstore) -> anyhow::Result<Cid> {
+        bs.put_cbor_default(&self.bytes())
+    }
+
+    /// Load tipset key from block store by its CID
+    pub fn load(bs: &impl Blockstore, cid: &Cid) -> anyhow::Result<Self> {
+        Self::from_bytes(bs.get_cbor_required(cid)?)
+    }
 }
 
 impl From<NonEmpty<Cid>> for TipsetKey {
@@ -654,15 +696,17 @@ mod lotus_json {
 #[cfg(test)]
 mod test {
     use super::*;
-    use crate::blocks::VRFProof;
     use crate::blocks::{
-        CachingBlockHeader, ElectionProof, Ticket, Tipset, TipsetKey, header::RawBlockHeader,
+        CachingBlockHeader, ElectionProof, Ticket, Tipset, TipsetKey, VRFProof,
+        header::RawBlockHeader,
     };
+    use crate::db::MemoryDB;
     use crate::shim::address::Address;
-    use crate::utils::multihash::prelude::*;
     use cid::Cid;
     use fvm_ipld_encoding::DAG_CBOR;
     use num_bigint::BigInt;
+    use quickcheck::Arbitrary;
+    use quickcheck_macros::quickcheck;
     use std::iter;
 
     pub fn mock_block(id: u64, weight: u64, ticket_sequence: u64) -> CachingBlockHeader {
@@ -810,4 +854,35 @@ mod test {
             CreateTipsetError::Empty
         );
     }
+
+    impl Arbitrary for TipsetKey {
+        fn arbitrary(g: &mut quickcheck::Gen) -> Self {
+            let blocks: nunny::Vec<Vec<u8>> = nunny::Vec::arbitrary(g);
+            let cids = nunny::Vec::new(
+                blocks
+                    .into_iter()
+                    .map(|b| {
+                        Cid::new_v1(
+                            fvm_ipld_encoding::DAG_CBOR,
+                            MultihashCode::Blake2b256.digest(&b),
+                        )
+                    })
+                    .collect_vec(),
+            )
+            .expect("infallible");
+            cids.into()
+        }
+    }
+
+    #[quickcheck]
+    fn tipset_key_bytes(tsk: TipsetKey) {
+        let bytes = tsk.bytes();
+        let tsk2 = TipsetKey::from_bytes(bytes).unwrap();
+        assert_eq!(tsk, tsk2);
+
+        let bs = MemoryDB::default();
+        let cid = tsk.save(&bs).unwrap();
+        let tsk3 = TipsetKey::load(&bs, &cid).unwrap();
+        assert_eq!(tsk, tsk3);
+    }
 }

src/chain/mod.rs

@@ -35,6 +35,7 @@ pub struct ExportOptions {
     pub skip_checksum: bool,
     pub include_receipts: bool,
     pub include_events: bool,
+    pub include_tipset_keys: bool,
     pub seen: CidHashSet,
 }
 
@@ -144,6 +145,7 @@ async fn export_to_forest_car<D: Digest>(
         skip_checksum,
         include_receipts,
         include_events,
+        include_tipset_keys,
         seen,
     } = options.unwrap_or_default();
 
@@ -171,6 +173,7 @@ async fn export_to_forest_car<D: Digest>(
         .with_seen(seen)
         .with_message_receipts(include_receipts)
         .with_events(include_events)
+        .with_tipset_keys(include_tipset_keys)
         .track_progress(true),
     );
 

src/chain/store/chain_store.rs

@@ -146,6 +146,7 @@ where
         self.heaviest_tipset_key_provider
             .set_heaviest_tipset_key(ts.key())?;
         *self.heaviest_tipset_cache.write() = Some(ts.clone());
+        ts.key().save(self.blockstore())?;
         if self.publisher.send(HeadChange::Apply(ts)).is_err() {
             debug!("did not publish head change, no active receivers");
         }
@@ -168,23 +169,9 @@ where
         Ok(())
     }
 
-    /// Writes the `TipsetKey` to the blockstore for `EthAPI` queries.
-    pub fn put_tipset_key(&self, tsk: &TipsetKey) -> Result<(), Error> {
-        let tsk_bytes = tsk.bytes();
-        let tsk_cid = self.blockstore().put_cbor_default(&tsk_bytes)?;
-        let hash = tsk_cid.into();
-        self.eth_mappings.write_obj(&hash, tsk)?;
-        Ok(())
-    }
-
     /// Reads the `TipsetKey` from the blockstore for `EthAPI` queries.
     pub fn get_required_tipset_key(&self, hash: &EthHash) -> Result<TipsetKey, Error> {
-        let tsk = self
-            .eth_mappings
-            .read_obj::<TipsetKey>(hash)?
-            .with_context(|| format!("cannot find tipset with hash {hash}"))?;
-
-        Ok(tsk)
+        Ok(TipsetKey::load(self.blockstore(), &hash.to_cid())?)
     }
 
     /// Writes with timestamp the `Hash` to `Cid` mapping to the blockstore for `EthAPI` queries.

src/cli/subcommands/snapshot_cmd.rs

@@ -145,6 +145,7 @@ impl SnapshotCommands {
                     tipset_keys: tipset.key().clone().into(),
                     include_receipts: false,
                     include_events: false,
+                    include_tipset_keys: false,
                     skip_checksum,
                     dry_run,
                 };

src/daemon/db_util.rs

@@ -349,7 +349,7 @@ where
             .headers_delegated_messages(ts.block_headers().iter())?,
     );
     tracing::trace!("Indexing tipset @{}: {}", epoch, &tsk);
-    state_manager.chain_store().put_tipset_key(&tsk)?;
+    tsk.save(state_manager.blockstore())?;
 
     Ok(())
 }

src/daemon/mod.rs

@@ -505,13 +505,10 @@ fn maybe_start_indexer_service(
 
             // Continuously listen for head changes
             loop {
-                let msg = receiver.recv().await?;
+                let HeadChange::Apply(ts) = receiver.recv().await?;
 
-                let HeadChange::Apply(ts) = msg;
                 tracing::debug!("Indexing tipset {}", ts.key());
 
-                chain_store.put_tipset_key(ts.key())?;
-
                 let delegated_messages =
                     chain_store.headers_delegated_messages(ts.block_headers().iter())?;
 

src/db/gc/snapshot.rs

@@ -259,6 +259,7 @@ where
                 skip_checksum: true,
                 include_receipts: true,
                 include_events: true,
+                include_tipset_keys: true,
                 seen: Default::default(),
             }),
         )