brk/_src/parser/databases/address_to_address_index.rs

use std::{
    collections::BTreeMap,
    fs, mem,
    path::{Path, PathBuf},
};

use allocative::Allocative;
use itertools::Itertools;
use snkrj::{AnyDatabase, Database};

use crate::structs::{Address, Config, U8x19, U8x31};

use super::{AnyDatabaseGroup, Metadata};

type Value = u32;
type U8x19Database = Database<U8x19, Value>;
type U8x31Database = Database<U8x31, Value>;
type U32Database = Database<u32, Value>;

type P2PKDatabase = U8x19Database;
type P2PKHDatabase = U8x19Database;
type P2SHDatabase = U8x19Database;
type P2WPKHDatabase = U8x19Database;
type P2WSHDatabase = U8x31Database;
type P2TRDatabase = U8x31Database;
type UnknownDatabase = U32Database;
type OpReturnDatabase = U32Database;
type PushOnlyDatabase = U32Database;
type EmptyDatabase = U32Database;
type MultisigDatabase = U32Database;

#[derive(Allocative)]
pub struct AddressToAddressIndex {
    path: PathBuf,
    pub metadata: Metadata,

    #[allocative(skip)]
    p2pk: BTreeMap<u16, P2PKDatabase>,
    #[allocative(skip)]
    p2pkh: BTreeMap<u16, P2PKHDatabase>,
    #[allocative(skip)]
    p2sh: BTreeMap<u16, P2SHDatabase>,
    #[allocative(skip)]
    p2wpkh: BTreeMap<u16, P2WPKHDatabase>,
    #[allocative(skip)]
    p2wsh: BTreeMap<u16, P2WSHDatabase>,
    #[allocative(skip)]
    p2tr: BTreeMap<u16, P2TRDatabase>,
    #[allocative(skip)]
    op_return: Option<OpReturnDatabase>,
    #[allocative(skip)]
    push_only: Option<PushOnlyDatabase>,
    #[allocative(skip)]
    unknown: Option<UnknownDatabase>,
    #[allocative(skip)]
    empty: Option<EmptyDatabase>,
    #[allocative(skip)]
    multisig: Option<MultisigDatabase>,
}

impl AddressToAddressIndex {
    pub fn open_db(&mut self, address: &Address) {
        match address {
            Address::Empty(_) => {
                self.open_empty();
            }
            Address::Unknown(_) => {
                self.open_unknown();
            }
            Address::OpReturn(_) => {
                self.open_op_return();
            }
            Address::PushOnly(_) => {
                self.open_push_only();
            }
            Address::MultiSig(_) => {
                self.open_multisig();
            }
            Address::P2PK((prefix, _)) => {
                self.open_p2pk(*prefix);
            }
            Address::P2PKH((prefix, _)) => {
                self.open_p2pkh(*prefix);
            }
            Address::P2SH((prefix, _)) => {
                self.open_p2sh(*prefix);
            }
            Address::P2WPKH((prefix, _)) => {
                self.open_p2wpkh(*prefix);
            }
            Address::P2WSH((prefix, _)) => {
                self.open_p2wsh(*prefix);
            }
            Address::P2TR((prefix, _)) => {
                self.open_p2tr(*prefix);
            }
        }
    }

    /// Doesn't check if the database is open contrary to `safe_get` which does and opens if needed.
    /// Though it makes it easy to use with rayon
    pub fn unsafe_get(&self, address: &Address) -> Option<&Value> {
        match address {
            Address::Empty(key) => self.empty.as_ref().unwrap().get(key),
            Address::Unknown(key) => self.unknown.as_ref().unwrap().get(key),
            Address::OpReturn(key) => self.op_return.as_ref().unwrap().get(key),
            Address::PushOnly(key) => self.push_only.as_ref().unwrap().get(key),
            Address::MultiSig(key) => self.multisig.as_ref().unwrap().get(key),
            Address::P2PK((prefix, key)) => self.p2pk.get(prefix).unwrap().get(key),
            Address::P2PKH((prefix, key)) => self.p2pkh.get(prefix).unwrap().get(key),
            Address::P2SH((prefix, key)) => self.p2sh.get(prefix).unwrap().get(key),
            Address::P2WPKH((prefix, key)) => self.p2wpkh.get(prefix).unwrap().get(key),
            Address::P2WSH((prefix, key)) => self.p2wsh.get(prefix).unwrap().get(key),
            Address::P2TR((prefix, key)) => self.p2tr.get(prefix).unwrap().get(key),
        }
    }

    pub fn get_from_ram(&self, address: &Address) -> Option<&Value> {
        match address {
            Address::Empty(key) => self.empty.as_ref().unwrap().get_from_ram(key),
            Address::Unknown(key) => self.unknown.as_ref().unwrap().get_from_ram(key),
            Address::OpReturn(key) => self.op_return.as_ref().unwrap().get_from_ram(key),
            Address::PushOnly(key) => self.push_only.as_ref().unwrap().get_from_ram(key),
            Address::MultiSig(key) => self.multisig.as_ref().unwrap().get_from_ram(key),
            Address::P2PK((prefix, key)) => self.p2pk.get(prefix).unwrap().get_from_ram(key),
            Address::P2PKH((prefix, key)) => self.p2pkh.get(prefix).unwrap().get_from_ram(key),
            Address::P2SH((prefix, key)) => self.p2sh.get(prefix).unwrap().get_from_ram(key),
            Address::P2WPKH((prefix, key)) => self.p2wpkh.get(prefix).unwrap().get_from_ram(key),
            Address::P2WSH((prefix, key)) => self.p2wsh.get(prefix).unwrap().get_from_ram(key),
            Address::P2TR((prefix, key)) => self.p2tr.get(prefix).unwrap().get_from_ram(key),
        }
    }

    pub fn insert(&mut self, address: Address, value: Value) -> Option<Value> {
        self.metadata.called_insert();

        match address {
            Address::Empty(key) => self.open_empty().insert(key, value),
            Address::Unknown(key) => self.open_unknown().insert(key, value),
            Address::OpReturn(key) => self.open_op_return().insert(key, value),
            Address::PushOnly(key) => self.open_push_only().insert(key, value),
            Address::MultiSig(key) => self.open_multisig().insert(key, value),
            Address::P2PK((prefix, rest)) => self.open_p2pk(prefix).insert(rest, value),
            Address::P2PKH((prefix, rest)) => self.open_p2pkh(prefix).insert(rest, value),
            Address::P2SH((prefix, rest)) => self.open_p2sh(prefix).insert(rest, value),
            Address::P2WPKH((prefix, rest)) => self.open_p2wpkh(prefix).insert(rest, value),
            Address::P2WSH((prefix, rest)) => self.open_p2wsh(prefix).insert(rest, value),
            Address::P2TR((prefix, rest)) => self.open_p2tr(prefix).insert(rest, value),
        }
    }

    fn path_to_group_prefixes(path: &Path) -> Vec<u16> {
        let folder = fs::read_dir(path);

        if folder.is_err() {
            return vec![];
        }

        folder
            .unwrap()
            .map(|entry| {
                entry
                    .unwrap()
                    .path()
                    .file_name()
                    .unwrap()
                    .to_str()
                    .unwrap()
                    .to_owned()
                    .parse::<u16>()
                    .unwrap()
            })
            .collect_vec()
    }

    fn path_p2pk(&self) -> PathBuf {
        self.path().join("p2pk")
    }

    pub fn open_p2pk(&mut self, prefix: u16) -> &mut P2PKDatabase {
        let path = self.path_p2pk();
        self.p2pk.entry(prefix).or_insert_with(|| {
            let path = path.join(prefix.to_string());
            Database::open(path).unwrap()
        })
    }

    fn open_all_p2pk(&mut self) {
        let path = self.path_p2pk();
        Self::path_to_group_prefixes(&path)
            .into_iter()
            .for_each(|prefix| {
                self.p2pk.insert(prefix, {
                    let path = path.join(prefix.to_string());
                    Database::open(path).unwrap()
                });
            });
    }

    fn path_p2pkh(&self) -> PathBuf {
        self.path().join("p2pkh")
    }

    pub fn open_p2pkh(&mut self, prefix: u16) -> &mut P2PKHDatabase {
        let path = self.path_p2pkh();

        self.p2pkh.entry(prefix).or_insert_with(|| {
            let path = path.join(prefix.to_string());
            Database::open(path).unwrap()
        })
    }

    fn open_all_p2pkh(&mut self) {
        let path = self.path_p2pkh();
        Self::path_to_group_prefixes(&path)
            .into_iter()
            .for_each(|prefix| {
                self.p2pkh.insert(prefix, {
                    let path = path.join(prefix.to_string());
                    Database::open(path).unwrap()
                });
            });
    }

    fn path_p2sh(&self) -> PathBuf {
        self.path().join("p2sh")
    }

    pub fn open_p2sh(&mut self, prefix: u16) -> &mut P2SHDatabase {
        let path = self.path_p2sh();

        self.p2sh.entry(prefix).or_insert_with(|| {
            let path = path.join(prefix.to_string());
            Database::open(path).unwrap()
        })
    }

    fn open_all_p2sh(&mut self) {
        let path = self.path_p2sh();
        Self::path_to_group_prefixes(&path)
            .into_iter()
            .for_each(|prefix| {
                self.p2sh.insert(prefix, {
                    let path = path.join(prefix.to_string());
                    Database::open(path).unwrap()
                });
            });
    }

    fn path_p2wpkh(&self) -> PathBuf {
        self.path().join("p2wpkh")
    }

    pub fn open_p2wpkh(&mut self, prefix: u16) -> &mut P2WPKHDatabase {
        let path = self.path_p2wpkh();

        self.p2wpkh.entry(prefix).or_insert_with(|| {
            let path = path.join(prefix.to_string());
            Database::open(path).unwrap()
        })
    }

    fn open_all_p2wpkh(&mut self) {
        let path = self.path_p2wpkh();
        Self::path_to_group_prefixes(&path)
            .into_iter()
            .for_each(|prefix| {
                self.p2wpkh.insert(prefix, {
                    let path = path.join(prefix.to_string());
                    Database::open(path).unwrap()
                });
            });
    }

    fn path_p2wsh(&self) -> PathBuf {
        self.path().join("p2wsh")
    }

    pub fn open_p2wsh(&mut self, prefix: u16) -> &mut P2WSHDatabase {
        let path = self.path_p2wsh();

        self.p2wsh.entry(prefix).or_insert_with(|| {
            let path = path.join(prefix.to_string());
            Database::open(path).unwrap()
        })
    }

    fn open_all_p2wsh(&mut self) {
        let path = self.path_p2wsh();
        Self::path_to_group_prefixes(&path)
            .into_iter()
            .for_each(|prefix| {
                self.p2wsh.insert(prefix, {
                    let path = path.join(prefix.to_string());
                    Database::open(path).unwrap()
                });
            });
    }

    fn path_p2tr(&self) -> PathBuf {
        self.path().join("p2tr")
    }

    pub fn open_p2tr(&mut self, prefix: u16) -> &mut P2TRDatabase {
        let path = self.path_p2tr();

        self.p2tr.entry(prefix).or_insert_with(|| {
            let path = path.join(prefix.to_string());
            Database::open(path).unwrap()
        })
    }

    fn open_all_p2tr(&mut self) {
        let path = self.path_p2tr();
        Self::path_to_group_prefixes(&path)
            .into_iter()
            .for_each(|prefix| {
                self.p2tr.insert(prefix, {
                    let path = path.join(prefix.to_string());
                    Database::open(path).unwrap()
                });
            });
    }

    pub fn open_unknown(&mut self) -> &mut UnknownDatabase {
        self.unknown
            .get_or_insert_with(|| Database::open(self.path.join("unknown")).unwrap())
    }

    pub fn open_op_return(&mut self) -> &mut UnknownDatabase {
        self.op_return
            .get_or_insert_with(|| Database::open(self.path.join("op_return")).unwrap())
    }

    pub fn open_push_only(&mut self) -> &mut UnknownDatabase {
        self.push_only
            .get_or_insert_with(|| Database::open(self.path.join("push_only")).unwrap())
    }

    pub fn open_empty(&mut self) -> &mut UnknownDatabase {
        self.empty
            .get_or_insert_with(|| Database::open(self.path.join("empty")).unwrap())
    }

    pub fn open_multisig(&mut self) -> &mut MultisigDatabase {
        self.multisig
            .get_or_insert_with(|| Database::open(self.path.join("multisig")).unwrap())
    }
}

impl AnyDatabaseGroup for AddressToAddressIndex {
    fn import(config: &Config) -> Self {
        let path = config.path_databases().join("address_to_address_index");
        Self {
            metadata: Metadata::import(&path, 1),
            path,

            p2pk: BTreeMap::default(),
            p2pkh: BTreeMap::default(),
            p2sh: BTreeMap::default(),
            p2wpkh: BTreeMap::default(),
            p2wsh: BTreeMap::default(),
            p2tr: BTreeMap::default(),
            op_return: None,
            push_only: None,
            unknown: None,
            empty: None,
            multisig: None,
        }
    }

    fn create_dir_all(&self) -> color_eyre::Result<(), std::io::Error> {
        fs::create_dir_all(self.path_p2pk()).unwrap();
        fs::create_dir_all(self.path_p2pkh()).unwrap();
        fs::create_dir_all(self.path_p2sh()).unwrap();
        fs::create_dir_all(self.path_p2wpkh()).unwrap();
        fs::create_dir_all(self.path_p2wsh()).unwrap();
        fs::create_dir_all(self.path_p2tr())
    }

    fn reset_metadata(&mut self) {
        self.metadata.reset()
    }

    fn drain_to_vec(&mut self) -> Vec<Box<dyn AnyDatabase + Send>> {
        mem::take(&mut self.p2pk)
            .into_values()
            .map(|db| Box::new(db) as Box<dyn AnyDatabase + Send>)
            .chain(
                mem::take(&mut self.p2pkh)
                    .into_values()
                    .map(|db| Box::new(db) as Box<dyn AnyDatabase + Send>),
            )
            .chain(
                mem::take(&mut self.p2sh)
                    .into_values()
                    .map(|db| Box::new(db) as Box<dyn AnyDatabase + Send>),
            )
            .chain(
                mem::take(&mut self.p2wpkh)
                    .into_values()
                    .map(|db| Box::new(db) as Box<dyn AnyDatabase + Send>),
            )
            .chain(
                mem::take(&mut self.p2wsh)
                    .into_values()
                    .map(|db| Box::new(db) as Box<dyn AnyDatabase + Send>),
            )
            .chain(
                mem::take(&mut self.p2tr)
                    .into_values()
                    .map(|db| Box::new(db) as Box<dyn AnyDatabase + Send>),
            )
            .chain(
                [
                    self.unknown.take(),
                    self.op_return.take(),
                    self.push_only.take(),
                    self.empty.take(),
                    self.multisig.take(),
                ]
                .into_iter()
                .flatten()
                .map(|db| Box::new(db) as Box<dyn AnyDatabase + Send>),
            )
            .collect_vec()
    }

    fn open_all(&mut self) {
        self.open_all_p2pk();
        self.open_all_p2pkh();
        self.open_all_p2wpkh();
        self.open_all_p2wsh();
        self.open_all_p2sh();
        self.open_all_p2tr();
    }

    fn metadata(&mut self) -> &mut Metadata {
        &mut self.metadata
    }

    fn path(&self) -> &Path {
        &self.path
    }
}