forest-filecoin 0.34.1

// Copyright 2019-2026 ChainSafe Systems
// SPDX-License-Identifier: Apache-2.0, MIT

use super::{
    BLOCK_BLOOM_LEN, EthBlockBloomStore, EthMappingsStore, SettingsStore, SettingsStoreExt,
    decode_block_bloom, encode_block_bloom,
};
use crate::blocks::TipsetKey;
use crate::db::PersistentStore;
use crate::libp2p_bitswap::{BitswapStoreRead, BitswapStoreReadWrite};
use crate::prelude::*;
use crate::rpc::eth::types::EthHash;
use crate::shim::clock::ChainEpoch;
use crate::utils::db::car_stream::CarBlock;
use crate::utils::multihash::prelude::*;
use ahash::HashMap;
use indexmap::IndexMap;
use nunny::Vec as NonEmpty;
use parking_lot::RwLock;

#[derive(Debug, Default)]
pub struct MemoryDB {
    blockchain_db: RwLock<HashMap<Cid, Vec<u8>>>,
    blockchain_persistent_db: RwLock<HashMap<Cid, Vec<u8>>>,
    settings_db: RwLock<HashMap<String, Vec<u8>>>,
    pub eth_mappings_db: RwLock<HashMap<EthHash, Vec<u8>>>,
    pub ts_lookup_db: RwLock<HashMap<ChainEpoch, TipsetKey>>,
    pub eth_block_bloom_db: RwLock<HashMap<Cid, Vec<u8>>>,
}

impl MemoryDB {
    pub fn blockstore_len(&self) -> usize {
        self.blockchain_db.read().len() + self.blockchain_persistent_db.read().len()
    }

    pub fn blockstore_size_bytes(&self) -> usize {
        self.blockchain_db
            .read()
            .iter()
            .chain(self.blockchain_persistent_db.read().iter())
            .map(|(k, v)| k.to_bytes().len() + v.len())
            .sum()
    }

    pub async fn export_forest_car<W: tokio::io::AsyncWrite + Unpin>(
        &self,
        writer: &mut W,
    ) -> anyhow::Result<()> {
        let roots =
            SettingsStoreExt::read_obj::<TipsetKey>(self, crate::db::setting_keys::HEAD_KEY)?
                .context("chain head is not tracked and cannot be exported")?
                .into_cids();
        self.export_forest_car_with_roots(roots, writer).await
    }

    pub async fn export_forest_car_with_roots<W: tokio::io::AsyncWrite + Unpin>(
        &self,
        roots: NonEmpty<Cid>,
        writer: &mut W,
    ) -> anyhow::Result<()> {
        let blocks = {
            let blockchain_db = self.blockchain_db.read();
            let blockchain_persistent_db = self.blockchain_persistent_db.read();
            blockchain_db
                .iter()
                .chain(blockchain_persistent_db.iter())
                // Sort to make the result CAR deterministic
                .sorted_by_key(|&(&cid, _)| cid)
                .map(|(&cid, data)| {
                    anyhow::Ok(CarBlock {
                        cid,
                        data: data.clone().into(),
                    })
                })
                .collect_vec()
        };
        let frames =
            crate::db::car::forest::Encoder::compress_stream_default(futures::stream::iter(blocks));
        crate::db::car::forest::Encoder::write(writer, roots, frames).await
    }
}

impl SettingsStore for MemoryDB {
    fn read_bin(&self, key: &str) -> anyhow::Result<Option<Vec<u8>>> {
        Ok(self.settings_db.read().get(key).cloned())
    }

    fn write_bin(&self, key: &str, value: &[u8]) -> anyhow::Result<()> {
        self.settings_db
            .write()
            .insert(key.to_owned(), value.to_vec());
        Ok(())
    }

    fn exists(&self, key: &str) -> anyhow::Result<bool> {
        Ok(self.settings_db.read().contains_key(key))
    }

    fn setting_keys(&self) -> anyhow::Result<Vec<String>> {
        Ok(self.settings_db.read().keys().cloned().collect_vec())
    }
}

impl EthMappingsStore for MemoryDB {
    fn read_bin(&self, key: &EthHash) -> anyhow::Result<Option<Vec<u8>>> {
        Ok(self.eth_mappings_db.read().get(key).cloned())
    }

    fn write_bin(&self, key: &EthHash, value: &[u8]) -> anyhow::Result<()> {
        self.eth_mappings_db
            .write()
            .insert(key.to_owned(), value.to_vec());
        Ok(())
    }

    fn exists(&self, key: &EthHash) -> anyhow::Result<bool> {
        Ok(self.eth_mappings_db.read().contains_key(key))
    }

    fn get_message_cids(&self) -> anyhow::Result<Vec<(Cid, u64)>> {
        let cids = self
            .eth_mappings_db
            .read()
            .values()
            .filter_map(|value| fvm_ipld_encoding::from_slice::<(Cid, u64)>(value).ok())
            .collect();

        Ok(cids)
    }

    fn delete(&self, keys: Vec<EthHash>) -> anyhow::Result<()> {
        let mut lock = self.eth_mappings_db.write();
        for hash in keys.iter() {
            lock.remove(hash);
        }
        Ok(())
    }

    fn tipset_key_by_epoch(&self, epoch: ChainEpoch) -> anyhow::Result<Option<TipsetKey>> {
        Ok(self.ts_lookup_db.read().get(&epoch).cloned())
    }

    fn delete_tipset_key_at_epoch(&self, epoch: ChainEpoch) -> anyhow::Result<()> {
        self.ts_lookup_db.write().remove(&epoch);
        Ok(())
    }

    fn set_tipset_key_at_epoch_raw(
        &self,
        epoch: ChainEpoch,
        tsk: &TipsetKey,
    ) -> anyhow::Result<()> {
        self.ts_lookup_db.write().insert(epoch, tsk.clone());
        Ok(())
    }
}

impl EthBlockBloomStore for MemoryDB {
    fn read_bloom(&self, key: &Cid) -> anyhow::Result<Option<[u8; BLOCK_BLOOM_LEN]>> {
        Ok(self
            .eth_block_bloom_db
            .read()
            .get(key)
            .and_then(|entry| decode_block_bloom(entry).map(|(_, bloom)| *bloom)))
    }

    fn write_bloom(
        &self,
        key: &Cid,
        height: ChainEpoch,
        bloom: &[u8; BLOCK_BLOOM_LEN],
    ) -> anyhow::Result<()> {
        self.eth_block_bloom_db
            .write()
            .insert(*key, encode_block_bloom(height, bloom));
        Ok(())
    }

    fn delete_blooms_before_height(&self, height: ChainEpoch) -> anyhow::Result<()> {
        self.eth_block_bloom_db
            .write()
            .retain(|_, entry| decode_block_bloom(entry).is_some_and(|(h, _)| h >= height));
        Ok(())
    }
}

impl Blockstore for MemoryDB {
    fn get(&self, k: &Cid) -> anyhow::Result<Option<Vec<u8>>> {
        Ok(self.blockchain_db.read().get(k).cloned().or(self
            .blockchain_persistent_db
            .read()
            .get(k)
            .cloned()))
    }

    fn put_keyed(&self, k: &Cid, block: &[u8]) -> anyhow::Result<()> {
        self.blockchain_db.write().insert(*k, block.to_vec());
        Ok(())
    }
}

impl PersistentStore for MemoryDB {
    fn put_keyed_persistent(&self, k: &Cid, block: &[u8]) -> anyhow::Result<()> {
        self.blockchain_persistent_db
            .write()
            .insert(*k, block.to_vec());
        Ok(())
    }
}

impl BitswapStoreRead for MemoryDB {
    fn contains(&self, cid: &Cid) -> anyhow::Result<bool> {
        Ok(self.blockchain_db.read().contains_key(cid))
    }

    fn get(&self, cid: &Cid) -> anyhow::Result<Option<Vec<u8>>> {
        Blockstore::get(self, cid)
    }
}

impl BitswapStoreReadWrite for MemoryDB {
    type Hashes = MultihashCode;

    fn insert(&self, block: &crate::libp2p_bitswap::Block64<Self::Hashes>) -> anyhow::Result<()> {
        self.put_keyed(block.cid(), block.data())
    }
}

impl super::HeaviestTipsetKeyProvider for MemoryDB {
    fn heaviest_tipset_key(&self) -> anyhow::Result<Option<TipsetKey>> {
        SettingsStoreExt::read_obj::<TipsetKey>(self, crate::db::setting_keys::HEAD_KEY)
    }

    fn set_heaviest_tipset_key(&self, tsk: &TipsetKey) -> anyhow::Result<()> {
        SettingsStoreExt::write_obj(self, crate::db::setting_keys::HEAD_KEY, tsk)
    }
}

#[derive(Debug, Default, derive_more::Deref)]
/// A memory blockstore that preserves the insertion order
pub struct IndexMapBlockstore {
    inner: RwLock<IndexMap<Cid, Vec<u8>>>,
}

impl IndexMapBlockstore {
    #[allow(dead_code)]
    pub async fn export_forest_car<W: tokio::io::AsyncWrite + Unpin>(
        &self,
        roots: NonEmpty<Cid>,
        writer: &mut W,
    ) -> anyhow::Result<()> {
        let blocks = {
            let inner = self.inner.read();
            let invalid_roots = roots
                .iter()
                .filter(|&c| !inner.contains_key(c))
                .collect_vec();
            anyhow::ensure!(
                invalid_roots.is_empty(),
                "All roots should present in the blockstore, invalid roots: {invalid_roots:?}"
            );
            inner
                .iter()
                .map(|(&cid, data)| {
                    anyhow::Ok(CarBlock {
                        cid,
                        data: data.clone().into(),
                    })
                })
                .collect_vec()
        };
        let frames =
            crate::db::car::forest::Encoder::compress_stream_default(futures::stream::iter(blocks));
        crate::db::car::forest::Encoder::write(writer, roots, frames).await
    }
}

impl Blockstore for IndexMapBlockstore {
    fn get(&self, k: &Cid) -> anyhow::Result<Option<Vec<u8>>> {
        Ok(self.read().get(k).cloned())
    }

    fn put_keyed(&self, k: &Cid, block: &[u8]) -> anyhow::Result<()> {
        self.write().insert(*k, block.to_vec());
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{
        db::{car::ForestCar, setting_keys::HEAD_KEY},
        utils::cid::CidCborExt as _,
    };
    use fil_actors_shared::fvm_ipld_hamt::Hamt;
    use fvm_ipld_encoding::DAG_CBOR;
    use multihash_codetable::Code::Blake2b256;
    use nunny::vec as nonempty;

    #[tokio::test]
    async fn test_export_forest_car() {
        let db = MemoryDB::default();
        let record1 = b"non-persistent";
        let key1 = Cid::new_v1(DAG_CBOR, Blake2b256.digest(record1.as_slice()));
        db.put_keyed(&key1, record1.as_slice()).unwrap();

        let record2 = b"persistent";
        let key2 = Cid::new_v1(DAG_CBOR, Blake2b256.digest(record2.as_slice()));
        db.put_keyed_persistent(&key2, record2.as_slice()).unwrap();

        let mut car_db_bytes = vec![];
        assert!(
            db.export_forest_car(&mut car_db_bytes)
                .await
                .unwrap_err()
                .to_string()
                .contains("chain head is not tracked and cannot be exported")
        );

        db.write_obj(HEAD_KEY, &TipsetKey::from(nonempty![key1]))
            .unwrap();

        car_db_bytes.clear();
        db.export_forest_car(&mut car_db_bytes).await.unwrap();

        let car = ForestCar::new(car_db_bytes).unwrap();
        assert_eq!(car.head_tipset_key(), &nonempty![key1]);
        assert!(car.has(&key1).unwrap());
        assert!(car.has(&key2).unwrap());
    }

    #[test]
    fn block_bloom_encode_decode() {
        let bloom = [0xab; 256];
        let entry = encode_block_bloom(42, &bloom);
        let (height, decoded) = decode_block_bloom(&entry).unwrap();
        assert_eq!(height, 42);
        assert_eq!(decoded, &bloom);
        assert!(decode_block_bloom(&[0, 1, 2]).is_none());
    }

    #[tokio::test]
    async fn test_index_map_blockstore() {
        const BIT_WIDTH: u32 = 5;

        let db = IndexMapBlockstore::default();
        // similate tipset lookup table
        let mut hamt: Hamt<_, TipsetKey, ChainEpoch> = Hamt::new_with_bit_width(&db, BIT_WIDTH);
        let checkpoints = [
            (
                0,
                TipsetKey::from(nunny::vec![Cid::from_cbor_blake2b256(&"1").unwrap()]),
            ),
            (
                5,
                TipsetKey::from(nunny::vec![Cid::from_cbor_blake2b256(&"5").unwrap()]),
            ),
            (
                10,
                TipsetKey::from(nunny::vec![Cid::from_cbor_blake2b256(&"10").unwrap()]),
            ),
        ];
        for (epoch, tsk) in checkpoints.iter().cloned() {
            hamt.set(epoch, tsk).unwrap();
        }
        let hamt_root = hamt.flush().unwrap();
        assert!(db.has(&hamt_root).unwrap(), "hamt root should present");

        // export with invalid roots should fail
        let mut car = vec![];
        db.export_forest_car(
            nunny::vec![Cid::from_cbor_blake2b256(&"invalid").unwrap()],
            &mut car,
        )
        .await
        .unwrap_err();

        // export with hamt root should succeed
        let mut car_bytes = vec![];
        let car_roots = nunny::vec![hamt_root];
        db.export_forest_car(car_roots.clone(), &mut car_bytes)
            .await
            .unwrap();

        let car: ForestCar<Vec<u8>> = ForestCar::new(car_bytes).unwrap();
        let hamt_from_car: Hamt<_, TipsetKey, ChainEpoch> =
            Hamt::load_with_bit_width(&hamt_root, &car, BIT_WIDTH).unwrap();

        let checkpoints_from_memdb_hamt = {
            let mut v = vec![];
            hamt.for_each_cacheless(|epoch, tsk| {
                v.push((*epoch, tsk.clone()));
                anyhow::Ok(())
            })
            .unwrap();
            v
        };
        let checkpoints_from_car_hamt = {
            let mut v = vec![];
            hamt_from_car
                .for_each_cacheless(|epoch, tsk| {
                    v.push((*epoch, tsk.clone()));
                    anyhow::Ok(())
                })
                .unwrap();
            v
        };
        // Cannot compare the 2 hamts as they use different DB types
        assert_eq!(checkpoints_from_memdb_hamt, checkpoints_from_car_hamt);
        // The hamt iteration order is different from the insertion order
        assert_eq!(
            HashMap::from_iter(checkpoints),
            HashMap::from_iter(checkpoints_from_car_hamt)
        );
    }
}