ckb-test-chain-utils 1.0.0

use ckb_db::RocksDB;
use ckb_db_schema::COLUMNS;
use ckb_store::{ChainDB, ChainStore};
use ckb_systemtime::unix_time_as_millis;
use ckb_types::core::error::OutPointError;
use ckb_types::{
    core::{
        BlockExt, BlockView, EpochExt, HeaderView,
        cell::{CellMetaBuilder, CellProvider, CellStatus, HeaderChecker},
    },
    packed::{Byte32, OutPoint},
};
use std::sync::Arc;
use tempfile::TempDir;

/// A temporary RocksDB for mocking chain storage.
#[doc(hidden)]
#[derive(Clone)]
pub struct MockStore {
    // The fields of a struct are dropped in declaration order.
    // So, put `ChainDB` (`RocksDB`) before `TempDir`.
    //
    // Ref: https://doc.rust-lang.org/reference/destructors.html
    inner: Arc<ChainDB>,
    _tmp_dir: Arc<TempDir>,
}

impl Default for MockStore {
    fn default() -> Self {
        let tmp_dir = TempDir::new().unwrap();
        let db = RocksDB::open_in(&tmp_dir, COLUMNS);
        MockStore {
            inner: Arc::new(ChainDB::new(db, Default::default())),
            _tmp_dir: Arc::new(tmp_dir),
        }
    }
}

impl MockStore {
    /// Create a new `MockStore` with insert parent block into the temporary database for reference.
    #[doc(hidden)]
    pub fn new(parent: &HeaderView, chain_store: &ChainDB) -> Self {
        let block = chain_store.get_block(&parent.hash()).unwrap();
        let epoch_ext = chain_store
            .get_block_epoch_index(&parent.hash())
            .and_then(|index| chain_store.get_epoch_ext(&index))
            .unwrap();
        let parent_block_ext = BlockExt {
            received_at: unix_time_as_millis(),
            total_difficulty: Default::default(),
            total_uncles_count: 0,
            verified: Some(true),
            txs_fees: vec![],
            cycles: None,
            txs_sizes: None,
        };
        let store = Self::default();
        {
            let db_txn = store.store().begin_transaction();
            db_txn
                .insert_block_ext(&block.parent_hash(), &parent_block_ext)
                .unwrap();
            db_txn.commit().unwrap();
        }
        store.insert_block(&block, &epoch_ext);
        store
    }

    /// Return the mock chainDB.
    #[doc(hidden)]
    pub fn store(&self) -> &ChainDB {
        &self.inner
    }

    /// Insert a block into mock chainDB.
    #[doc(hidden)]
    pub fn insert_block(&self, block: &BlockView, epoch_ext: &EpochExt) {
        let db_txn = self.store().begin_transaction();
        let last_block_hash_in_previous_epoch = epoch_ext.last_block_hash_in_previous_epoch();
        db_txn.insert_block(block).unwrap();
        db_txn.attach_block(block).unwrap();
        db_txn
            .insert_block_epoch_index(&block.hash(), &last_block_hash_in_previous_epoch)
            .unwrap();
        db_txn
            .insert_epoch_ext(&last_block_hash_in_previous_epoch, epoch_ext)
            .unwrap();
        {
            let parent_block_ext = self.store().get_block_ext(&block.parent_hash()).unwrap();
            let block_ext = BlockExt {
                received_at: unix_time_as_millis(),
                total_difficulty: parent_block_ext.total_difficulty.to_owned()
                    + block.header().difficulty(),
                total_uncles_count: parent_block_ext.total_uncles_count
                    + block.data().uncles().len() as u64,
                verified: Some(true),
                txs_fees: vec![],
                cycles: None,
                txs_sizes: None,
            };
            db_txn.insert_block_ext(&block.hash(), &block_ext).unwrap();
        }
        db_txn.commit().unwrap();
    }

    /// Remove a block from mock chainDB.
    #[doc(hidden)]
    pub fn remove_block(&self, block: &BlockView) {
        let db_txn = self.store().begin_transaction();
        db_txn.delete_block(block).unwrap();
        db_txn.detach_block(block).unwrap();
        db_txn.commit().unwrap();
    }
}

impl CellProvider for MockStore {
    fn cell(&self, out_point: &OutPoint, _eager_load: bool) -> CellStatus {
        match self.store().get_transaction(&out_point.tx_hash()) {
            Some((tx, _)) => tx
                .outputs()
                .get(out_point.index().into())
                .map(|cell| {
                    let data = tx
                        .outputs_data()
                        .get(out_point.index().into())
                        .expect("output data");

                    let cell_meta = CellMetaBuilder::from_cell_output(cell, data.into())
                        .out_point(out_point.to_owned())
                        .build();

                    CellStatus::live_cell(cell_meta)
                })
                .unwrap_or(CellStatus::Unknown),
            None => CellStatus::Unknown,
        }
    }
}

impl HeaderChecker for MockStore {
    fn check_valid(&self, block_hash: &Byte32) -> Result<(), OutPointError> {
        if self.store().get_block_number(block_hash).is_some() {
            Ok(())
        } else {
            Err(OutPointError::InvalidHeader(block_hash.clone()))
        }
    }
}