rostrum 14.0.1

An efficient implementation of Electrum Server with token support
Documentation
use std::{
    collections::HashMap,
    sync::{Mutex, MutexGuard},
};

use crate::chaindef::ScriptHash;
use bitcoincash::Txid;
use rayon::prelude::*;

use crate::{
    chaindef::OutPointHash,
    indexes::{
        headerindex::HeaderRow, heightindex::HeightIndexRow, inputindex::InputIndexRow,
        outputindex::OutputIndexRow, outputtokenindex::OutputTokenIndexRow,
        outtoscriptindex::OutToScripthashIndex, scripthashindex::ScriptHashIndexRow,
        tokenoutputindex::TokenOutputIndexRow, unspentindex::UnspentIndexRow, DBRow,
    },
    store::{Row, COLUMN_FAMILIES, META_CF},
};

#[derive(Clone)]
pub struct PrefilledSpentData {
    pub scripthash: ScriptHash,
    pub has_token: bool,
}

pub struct SpentUpdateBatch {
    // (outpoint, txid, vin, height, (scripthash, has_token)) - optionally has token information we otherwise look up.
    #[allow(clippy::type_complexity)]
    spending_metadata: Mutex<Vec<(OutPointHash, Txid, u32, u32, Option<PrefilledSpentData>)>>,

    // flag if we're cleaning out of mempool
    erase_only: bool,

    // prefilled spent data (for mempool etc)
    prefilled_spent_data: HashMap<OutPointHash, PrefilledSpentData>,
}

impl Default for SpentUpdateBatch {
    fn default() -> Self {
        Self::new()
    }
}

impl SpentUpdateBatch {
    pub fn new() -> Self {
        Self {
            spending_metadata: Default::default(),
            erase_only: false,
            prefilled_spent_data: Default::default(),
        }
    }

    pub fn set_erase_only(mut self) -> Self {
        self.erase_only = true;
        self
    }

    pub fn is_erase_only(&self) -> bool {
        self.erase_only
    }

    pub fn add_prefilled_spent_data(
        mut self,
        prefilled_spent_data: HashMap<OutPointHash, PrefilledSpentData>,
    ) -> Self {
        self.prefilled_spent_data.extend(prefilled_spent_data);
        self
    }

    pub fn insert_spending_metadata(
        &self,
        oph: OutPointHash,
        txid: Txid,
        vin: u32,
        height: u32,
        scripthash_and_has_token: Option<PrefilledSpentData>,
    ) {
        let scripthash_and_has_token =
            scripthash_and_has_token.or_else(|| self.prefilled_spent_data.get(&oph).cloned());

        self.spending_metadata.lock().unwrap().push((
            oph,
            txid,
            vin,
            height,
            scripthash_and_has_token,
        ));
    }

    pub fn take_spending_metadata(
        &self,
    ) -> Vec<(OutPointHash, Txid, u32, u32, Option<PrefilledSpentData>)> {
        std::mem::take(&mut *self.spending_metadata.lock().unwrap())
    }
}

/**
 * For collecting multiple database entries and batch writing them in one transaction.
 */
pub struct WriteBatch {
    metadata: Mutex<Vec<Row>>,
    header: Mutex<Vec<Row>>,
    height: Mutex<Vec<Row>>,
    input: Mutex<Vec<Row>>,
    output: Mutex<Vec<Row>>,
    output_token: Mutex<Vec<Row>>,
    scripthash: Mutex<Vec<Row>>,
    token_output: Mutex<Vec<Row>>,
    unspent: Mutex<Vec<Row>>,
    outtoscripthash: Mutex<Vec<Row>>,
}

impl WriteBatch {
    pub fn new() -> Self {
        WriteBatch {
            metadata: Default::default(),
            header: Default::default(),
            height: Default::default(),
            input: Default::default(),
            output: Default::default(),
            output_token: Default::default(),
            scripthash: Default::default(),
            token_output: Default::default(),
            unspent: Default::default(),
            outtoscripthash: Default::default(),
        }
    }

    pub fn cf_to_vec_tuples(&self) -> Vec<(&'static str, &Mutex<Vec<Row>>)> {
        let map = vec![
            (META_CF, &self.metadata),
            (HeaderRow::CF, &self.header),
            (HeightIndexRow::CF, &self.height),
            (InputIndexRow::CF, &self.input),
            (OutputIndexRow::CF, &self.output),
            (OutputTokenIndexRow::CF, &self.output_token),
            (ScriptHashIndexRow::CF, &self.scripthash),
            (TokenOutputIndexRow::CF, &self.token_output),
            (UnspentIndexRow::CF, &self.unspent),
            (OutToScripthashIndex::CF, &self.outtoscripthash),
        ];
        assert!(map.len() == COLUMN_FAMILIES.len());
        map
    }

    pub fn insert<I: IntoParallelIterator<Item = Row>>(&self, cf_name: &str, rows: I) {
        let mut rows: Vec<Row> = rows.into_par_iter().collect();

        let mut locked = match cf_name {
            META_CF => self.metadata.lock().unwrap(),
            HeaderRow::CF => self.header.lock().unwrap(),
            HeightIndexRow::CF => self.height.lock().unwrap(),
            InputIndexRow::CF => self.input.lock().unwrap(),
            OutputIndexRow::CF => self.output.lock().unwrap(),
            OutputTokenIndexRow::CF => self.output_token.lock().unwrap(),
            ScriptHashIndexRow::CF => self.scripthash.lock().unwrap(),
            TokenOutputIndexRow::CF => self.token_output.lock().unwrap(),
            UnspentIndexRow::CF => self.unspent.lock().unwrap(),
            OutToScripthashIndex::CF => self.outtoscripthash.lock().unwrap(),
            _ => panic!("unknown cf family {cf_name}"),
        };
        locked.append(&mut rows);
    }

    /**
     * How many entries are in the current batch.
     */
    pub fn len(&self) -> usize {
        let locks: Vec<MutexGuard<Vec<Row>>> = self
            .cf_to_vec_tuples()
            .into_iter()
            .map(|(_, column)| column.lock().unwrap())
            .collect();
        locks.iter().map(|rows| rows.len()).sum()
    }
    /**
     * No entires in writebatch
     */
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }
}

impl Default for WriteBatch {
    fn default() -> Self {
        Self::new()
    }
}