spark_rust/wallet/leaf_manager/

mod.rs

use crate::error::{validation::ValidationError, wallet::WalletError, SparkSdkError};
use hashbrown::HashMap;
use parking_lot::RwLock;
use serde::{Deserialize, Serialize};
use spark_protos::spark::TreeNode;
use std::sync::Arc;
use uuid::Uuid;

use super::internal_handlers::traits::leaves::LeafSelectionResponse;

struct SparkLeafEntry {
    leaf: SparkLeaf,
    status: SparkNodeStatus,
    unlocking_id: Option<String>,
}

type LeafMap = Arc<RwLock<HashMap<String, SparkLeafEntry>>>;

pub(crate) struct LeafManager {
    /// The map of leaf nodes
    leaves: LeafMap,
}

#[derive(Debug, Clone)]
pub(crate) enum SparkLeaf {
    Bitcoin(TreeNode),
}

impl SparkLeaf {
    pub(crate) fn get_id(&self) -> &String {
        match self {
            SparkLeaf::Bitcoin(leaf) => &leaf.id,
        }
    }

    pub(crate) fn get_value(&self) -> u64 {
        match self {
            SparkLeaf::Bitcoin(leaf) => leaf.value,
        }
    }

    pub(crate) fn is_bitcoin(&self) -> bool {
        matches!(self, SparkLeaf::Bitcoin(_))
    }

    pub(crate) fn get_tree_node(&self) -> Result<TreeNode, SparkSdkError> {
        match self {
            SparkLeaf::Bitcoin(leaf) => Ok(leaf.clone()),
        }
    }
}

/// Status of a user-owned Spark node. These are mostly leaves, but they can also represent aggregatable branch nodes (parents  ) if the user owns all children of the parent. The status of a transfer (or derivative thereof) may require a swap operation first. In this case, the status won't show the swap operation, since it is a sub-operation of the user-intended flow.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum SparkNodeStatus {
    /// Available for a new operation
    Available,

    /// Aggregatable parent
    AggregatableParent,

    /// Currently used for a transfer operation. It is locked and will no longer exist if the transfer operation succeeds.
    Transfer,

    /// Currently used for a split operation. It is locked and will no longer exist if the split operation succeeds.
    Split,

    /// Currently used for a swap operation. It is locked and will no longer exist if the swap operation succeeds.
    Swap,

    /// Currently used for a fee query that will be followed by a transfer (or derivative thereof) operation. It is locked for the operation that pre-fetches the fee. Since fee requests can also be called externally by the user, these locks should have a timeout duration.
    FeeQuery,

    /// Exiting through the SSP with cooperative exit
    CooperativeExit,

    /// Refreshing timelock
    RefreshTimelock,
}

impl SparkNodeStatus {
    fn generate_unlocking_id(&self) -> Option<String> {
        match self {
            SparkNodeStatus::Transfer | SparkNodeStatus::Split | SparkNodeStatus::Swap => {
                Some(Uuid::now_v7().to_string())
            }
            _ => None,
        }
    }
}

type LeafFilterCallback = Box<dyn Fn(&SparkLeaf) -> bool>;

#[derive(Debug)]
pub(crate) struct LockLeavesResponse {
    pub(crate) unlocking_id: Option<String>,
    pub(crate) leaves: Vec<SparkLeaf>,
}

pub(crate) type LeafFilterFunction = fn(&SparkLeaf) -> bool;

impl LeafManager {
    pub(crate) fn new() -> Self {
        Self {
            leaves: Arc::new(parking_lot::RwLock::new(HashMap::new())),
        }
    }

    #[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))]
    pub(crate) fn filter_nodes(&self, cb: Option<LeafFilterFunction>) -> Vec<SparkLeaf> {
        let mut nodes = Vec::new();
        let guard = self.leaves.read();
        for node in guard.values() {
            if cb.as_ref().is_some_and(|f| f(&node.leaf)) {
                nodes.push(node.leaf.clone());
            }
        }

        drop(guard);
        nodes
    }

    #[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))]
    pub(crate) fn filter_nodes_by_ids(&self, leaf_ids: &Vec<String>) -> Vec<SparkLeaf> {
        let guard = self.leaves.read();
        let mut nodes = Vec::new();
        for leaf_id in leaf_ids {
            if let Some(entry) = guard.get(leaf_id) {
                nodes.push(entry.leaf.clone());
            }
        }

        drop(guard);
        nodes
    }

    #[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))]
    pub(crate) fn filter_and_lock_nodes(
        &self,
        cb: Option<LeafFilterFunction>,
        new_status: Option<SparkNodeStatus>,
    ) -> Vec<SparkLeaf> {
        // Check if all nodes are lockable
        let mut nodes = Vec::new();
        let unlocking_id = match &new_status {
            Some(status) => status.generate_unlocking_id(),
            None => None,
        };

        let mut guard = self.leaves.write();
        for node in guard.values_mut() {
            if cb.as_ref().is_some_and(|f| f(&node.leaf)) {
                if new_status.is_some() {
                    node.unlocking_id = unlocking_id.clone();
                    node.status = new_status.clone().unwrap();
                }
                nodes.push(node.leaf.clone());
            }
        }

        drop(guard);
        nodes
    }

    #[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))]
    pub(crate) fn delete_and_insert_leaves_atomically(
        &self,
        delete_ids: &Vec<String>,
        insert_leaves: &Vec<SparkLeaf>,
    ) -> Result<(), SparkSdkError> {
        let mut guard = self.leaves.write();
        for id in delete_ids {
            let removed = guard.remove(id);
            if removed.is_none() {
                #[cfg(feature = "telemetry")]
                tracing::warn!(
                    "Leaf not found in wallet when deleting and inserting leaves atomically: {}",
                    id
                );
            }
        }

        for leaf in insert_leaves {
            guard.insert(
                leaf.get_id().clone(),
                SparkLeafEntry {
                    leaf: leaf.clone(),
                    status: SparkNodeStatus::Available,
                    unlocking_id: None,
                },
            );
        }

        Ok(())
    }

    #[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))]
    pub(crate) fn get_all_available_leaves(
        &self,
        lock_callback: Option<LeafFilterCallback>,
        new_status: Option<SparkNodeStatus>,
    ) -> (Vec<SparkLeaf>, Option<String>) {
        let mut nodes = Vec::new();
        let mut guard = self.leaves.write();
        let unlocking_id = match &new_status {
            Some(status) => status.generate_unlocking_id(),
            None => None,
        };

        for node in guard.values_mut() {
            if node.status == SparkNodeStatus::Available
                && lock_callback.as_ref().is_some_and(|f| f(&node.leaf))
            {
                if new_status.is_some() {
                    node.status = new_status.clone().unwrap().clone();
                    node.unlocking_id = unlocking_id.clone();
                }
                nodes.push(node.leaf.clone());
            }
        }

        drop(guard);
        (nodes, unlocking_id)
    }

    #[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))]
    pub(crate) fn lock_leaf_ids(
        &self,
        leaf_ids: &Vec<String>,
        new_status: SparkNodeStatus,
    ) -> Result<LeafSelectionResponse, SparkSdkError> {
        // make sure all leaves exist and are available
        let mut leaves = Vec::new();
        let mut guard = self.leaves.write();
        for leaf_id in leaf_ids {
            let get_leaf = guard.get(leaf_id);
            if get_leaf.is_none() {
                drop(guard);
                return Err(SparkSdkError::from(WalletError::LeafNotFoundInWallet {
                    leaf_id: leaf_id.clone(),
                }));
            }
            let leaf = get_leaf.unwrap();
            if leaf.status != SparkNodeStatus::Available {
                drop(guard);
                return Err(SparkSdkError::from(WalletError::LeafNotAvailableForUse {
                    leaf_id: leaf_id.clone(),
                }));
            }
            leaves.push(leaf.leaf.clone());
        }

        let unlocking_id = new_status.generate_unlocking_id();
        for leaf_id in leaf_ids {
            let leaf = guard.get_mut(leaf_id).unwrap();
            leaf.status = new_status.clone();
            leaf.unlocking_id = unlocking_id.clone();
        }

        // get total value of leaves
        let total_value = leaves.iter().map(|l| l.get_value()).sum();

        Ok(LeafSelectionResponse {
            leaves,
            total_value,
            unlocking_id,
            exact_amount: true,
        })
    }

    #[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))]
    pub(crate) fn remove_all_leaves(&self) -> Result<(), SparkSdkError> {
        let mut guard = self.leaves.write();
        guard.clear();
        Ok(())
    }

    #[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))]
    pub(crate) fn insert_leaves(
        &self,
        new_leaves: Vec<SparkLeaf>,
        delete_leaves_first: bool,
    ) -> Result<(), SparkSdkError> {
        let mut guard = self.leaves.write();

        if delete_leaves_first {
            guard.clear();
        }

        // check if any of the leaves already exist before insertion
        for leaf in &new_leaves {
            let id = leaf.get_id();
            if guard.contains_key(id.as_str()) {
                return Err(SparkSdkError::from(
                    WalletError::LeafAlreadyExistsInWallet {
                        leaf_id: id.clone(),
                    },
                ));
            }
        }

        // insert all leaves since we've verified none exist
        for leaf in new_leaves {
            let id = leaf.get_id();
            let leaf_entry = SparkLeafEntry {
                leaf: leaf.clone(),
                status: SparkNodeStatus::Available,
                unlocking_id: None,
            };
            guard.insert(id.clone(), leaf_entry);
        }

        Ok(())
    }

    #[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))]
    pub(crate) fn get_available_bitcoin_value(&self, filter_cb: Option<LeafFilterFunction>) -> u64 {
        // set the default filter for mapping btc leaves
        let default_filter = |ln: &SparkLeaf| ln.is_bitcoin();

        // combine the default filter with the custom filter
        let filter: Box<dyn Fn(&SparkLeaf) -> bool> = if let Some(custom_filter) = filter_cb {
            let combined_filter =
                move |node: &SparkLeaf| default_filter(node) && custom_filter(node);
            Box::new(combined_filter)
        } else {
            Box::new(move |node: &SparkLeaf| default_filter(node))
        };

        let guard = self.leaves.read();
        let mut available_btc_sum = 0;
        for node in guard.values() {
            if node.status != SparkNodeStatus::Available {
                continue;
            }

            if filter(&node.leaf) {
                available_btc_sum += node.leaf.get_value();
            }
        }

        drop(guard);
        available_btc_sum
    }

    #[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))]
    pub(crate) fn get_available_bitcoin_leaves(
        &self,
        filter_cb: Option<LeafFilterFunction>,
        new_status: SparkNodeStatus,
    ) -> Vec<SparkLeaf> {
        // set the default filter for mapping btc leaves
        let default_filter = |ln: &SparkLeaf| ln.is_bitcoin();

        // combine the default filter with the custom filter
        let filter: Box<dyn Fn(&SparkLeaf) -> bool> = if let Some(custom_filter) = filter_cb {
            let combined_filter =
                move |node: &SparkLeaf| default_filter(node) && custom_filter(node);
            Box::new(combined_filter)
        } else {
            Box::new(move |node: &SparkLeaf| default_filter(node))
        };

        let unlocking_id = Uuid::now_v7().to_string();

        let mut guard = self.leaves.write();
        let mut available_btc_leaves = Vec::new();
        for node in guard.values_mut() {
            if node.status != SparkNodeStatus::Available {
                continue;
            }

            node.status = new_status.clone();
            node.unlocking_id = Some(unlocking_id.clone());

            if filter(&node.leaf) {
                available_btc_leaves.push(node.leaf.clone());
            }
        }

        drop(guard);
        available_btc_leaves
    }

    /// Lock all available bitcoin leaves.
    #[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))]
    pub(crate) fn lock_available_bitcoin_leaves(
        &self,
        new_status: SparkNodeStatus,
    ) -> LockLeavesResponse {
        // set the default filter for mapping btc leaves
        let unlocking_id = Uuid::now_v7().to_string();
        let default_filter: Box<dyn Fn(&SparkLeaf) -> bool> = Box::new(|ln| ln.is_bitcoin());

        let mut guard = self.leaves.write();
        let mut available_btc_leaves = Vec::new();
        for node in guard.values_mut() {
            if node.status != SparkNodeStatus::Available {
                continue;
            }

            if default_filter(&node.leaf) {
                node.status = new_status.clone();
                node.unlocking_id = Some(unlocking_id.clone());
                available_btc_leaves.push(node.leaf.clone());
            }
        }

        drop(guard);
        LockLeavesResponse {
            unlocking_id: Some(unlocking_id),
            leaves: available_btc_leaves,
        }
    }

    #[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))]
    pub(crate) fn select_leaves(
        &self,
        target_amount: u64,
        new_status: SparkNodeStatus,
    ) -> Result<LeafSelectionResponse, SparkSdkError> {
        // check if target amount is 0
        if target_amount == 0 {
            return Err(SparkSdkError::from(ValidationError::InvalidInput {
                field: "Target amount cannot be 0".to_string(),
            }));
        }

        // prepare the filter, if token pubkey is None, filter for Bitcoin leaves
        let filter = Box::new(|_node: &SparkLeaf| true);

        // get all leaves that match the filter
        let mut guard = self.leaves.write();
        let mut filtered_leaves = Vec::new();
        for node in guard.values() {
            if filter(&node.leaf) {
                filtered_leaves.push(node.leaf.clone());
            }
        }

        // sort leaves in descending order by value
        filtered_leaves.sort_by_key(|b| std::cmp::Reverse(b.get_value()));

        let unlocking_id = new_status.generate_unlocking_id();
        let mut total_value = 0;
        let mut leaves = Vec::new();
        for leaf in filtered_leaves {
            if target_amount - total_value >= leaf.get_value() {
                let leaf_id = leaf.get_id().clone();

                total_value += leaf.get_value();
                leaves.push(leaf);

                // lock the leaf
                let leaf_entry = guard.get_mut(&leaf_id).unwrap();
                leaf_entry.status = new_status.clone();
                leaf_entry.unlocking_id = unlocking_id.clone();
            }
        }

        drop(guard);

        Ok(LeafSelectionResponse {
            leaves,
            total_value,
            unlocking_id: unlocking_id.clone(),
            exact_amount: total_value == target_amount,
        })
    }

    #[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))]
    pub(crate) fn unlock_leaves(
        &self,
        unlocking_id: String,
        leaf_ids: &Vec<String>,
        delete: bool,
    ) -> Result<(), SparkSdkError> {
        let mut leaves = self.leaves.write();

        // make sure that all leaves are in transfer lock with the same request id
        for leaf_id in leaf_ids {
            let leaf = leaves.get(leaf_id).ok_or_else(|| {
                SparkSdkError::from(WalletError::LeafNotFoundInWallet {
                    leaf_id: leaf_id.clone(),
                })
            })?;

            if leaf.unlocking_id != Some(unlocking_id.clone()) {
                return Err(SparkSdkError::from(WalletError::LeafNotUsingExpectedLock {
                    expected: unlocking_id.clone(),
                    actual: leaf.unlocking_id.clone().unwrap_or_default(),
                }));
            }
        }

        if delete {
            for leaf_id in leaf_ids {
                leaves.remove(leaf_id);
            }

            drop(leaves);
            return Ok(());
        }

        // update status to available
        for leaf_id in leaf_ids {
            let leaf = leaves.get_mut(leaf_id).unwrap();
            leaf.status = SparkNodeStatus::Available;
            leaf.unlocking_id = None;
        }

        drop(leaves);
        Ok(())
    }
}