spark_rust/wallet/internal_handlers/implementations/

leaves.rs

use spark_protos::spark::{query_nodes_request::Source, QueryNodesRequest};
use tonic::{async_trait, Request};

use crate::{
    error::{network::NetworkError, SparkSdkError, WalletError},
    signer::traits::SparkSigner,
    wallet::{
        internal_handlers::traits::leaves::{LeafSelectionResponse, LeavesInternalHandlers},
        leaf_manager::{SparkLeaf, SparkNodeStatus},
    },
    SparkSdk,
};

#[async_trait]
impl<S: SparkSigner + Send + Sync + Clone + 'static> LeavesInternalHandlers<S> for SparkSdk<S> {
    #[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))]
    async fn sync_leaves(&self) -> Result<(), SparkSdkError> {
        let mut spark_client = self.config.spark_config.get_spark_connection(None).await?;

        let identity_pubkey = self.get_spark_address()?;
        let network = self.config.spark_config.network;
        let mut request = Request::new(QueryNodesRequest {
            include_parents: true,
            source: Some(Source::OwnerIdentityPubkey(
                identity_pubkey.serialize().to_vec(),
            )),
            network: network.marshal_proto(),
        });

        self.add_authorization_header_to_request(&mut request, None);

        let response = spark_client
            .query_nodes(request)
            .await
            .map_err(|status| SparkSdkError::from(NetworkError::Status(status)))?;
        let queried_nodes = response.into_inner().nodes;

        // select the available leaves only
        let available_leaves_with_ids = queried_nodes
            .into_iter()
            .filter(|node| node.1.status.to_uppercase() == "AVAILABLE")
            .collect::<Vec<_>>();

        let available_leaves = available_leaves_with_ids
            .into_iter()
            .map(|(_, node)| SparkLeaf::Bitcoin(node))
            .collect();

        // refresh the leaf manager with the new leaves
        self.leaf_manager.insert_leaves(available_leaves, true)?;

        Ok(())
    }

    #[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))]
    async fn prepare_leaves_for_amount(
        &self,
        target_amount: u64,
    ) -> Result<LeafSelectionResponse, SparkSdkError> {
        if self.leaf_manager.get_available_bitcoin_value(None) < target_amount {
            return Err(SparkSdkError::from(
                WalletError::LeafSelectionInsufficientFunds {},
            ));
        }

        // Attempt leaf selection with the target amount. If the exact amount is found, return the response.
        let leaf_selection_response =
            self.leaf_manager
                .select_leaves(target_amount, None, SparkNodeStatus::Transfer)?;
        if leaf_selection_response.exact_amount {
            return Ok(leaf_selection_response);
        }

        // We don't have the exact amount available, so we'll need to swap leaves.
        let leaf_ids = leaf_selection_response
            .leaves
            .iter()
            .map(|l| l.get_id().clone())
            .collect();

        // Unlock the leaves before attempting to swap them.
        self.leaf_manager.unlock_leaves(
            leaf_selection_response.unlocking_id.unwrap(),
            &leaf_ids,
            false,
        )?;

        self.request_leaves_swap(target_amount).await?;

        // Retry leaf selection now since we should have the exact amount available.
        self.leaf_manager
            .select_leaves(target_amount, None, SparkNodeStatus::Transfer)
    }

    #[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))]
    async fn optimize_leaves(&self) -> Result<(), SparkSdkError> {
        // TODO leaves: add status to track if there is any optimization.
        //              we already keep status of the leaves, but this seems to be for all the leaves.
        //              this requires further thinking as a part of the next PR.

        // TODO leaves: or if optimization is in progress (if we are optimizing for all leaves)
        if !self.are_leaves_inefficient()? {
            return Ok(());
        }

        // TODO leaves: an appropriate status must be set here, and ideally passed into `are_leaves_inefficient` with the unlocking_id
        let leaves = self
            .leaf_manager
            .get_available_bitcoin_leaves(None, SparkNodeStatus::Available);

        let target_amount = leaves
            .iter()
            .map(|leaf| leaf.get_tree_node().unwrap().value)
            .sum();

        // TODO leaves: again, we have many unlocked gaps between operations in `are_leaves_efficient` and until below. We must think about this.
        if leaves.len() > 1 {
            self.request_leaves_swap(target_amount).await?;
        }

        self.sync_leaves().await?;

        Ok(())
    }

    #[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))]
    fn are_leaves_inefficient(&self) -> Result<bool, SparkSdkError> {
        // Get all available leaves
        let leaves = self
            .leaf_manager
            .get_available_bitcoin_leaves(None, SparkNodeStatus::Available); // TODO: we must use a function that does *not* request a status

        if leaves.len() <= 1 {
            return Ok(false);
        }

        // Calculate total amount across all Bitcoin leaves
        let total_btc_amount = leaves
            .iter()
            .fold(0, |acc, leaf| acc + leaf.get_tree_node().unwrap().value); // safe unwrap

        let next_lower_power_of_two = total_btc_amount.ilog2();

        let mut remaining_amount = total_btc_amount;
        let mut optimal_leaves_length = 0;

        // Iterate from highest power to lowest (e.g., from 2^63 down to 2^0)
        for i in (0..=next_lower_power_of_two).rev() {
            let denomination = 1u64.checked_shl(i).unwrap_or(0);
            if denomination == 0 {
                continue; // Skip if shift would overflow
            }

            while remaining_amount >= denomination {
                remaining_amount -= denomination;
                optimal_leaves_length += 1;
            }
        }

        let is_inefficient = leaves.len() > optimal_leaves_length * 5;

        #[cfg(feature = "telemetry")]
        tracing::debug!("Leaves are inefficient: {}", is_inefficient);

        Ok(is_inefficient)
    }
}

#[allow(dead_code)]
fn next_lower_power_of_two(total_amount: u64) -> u64 {
    if total_amount == 0 {
        0
    } else {
        // This safely gets the highest power of 2 that's less than or equal to total_amount
        1u64.checked_shl(total_amount.ilog2()).unwrap_or(1u64 << 63)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_next_lower_power_of_two() {
        assert_eq!(next_lower_power_of_two(0), 0);

        assert_eq!(next_lower_power_of_two(1), 1);
        assert_eq!(next_lower_power_of_two(2), 2);
        assert_eq!(next_lower_power_of_two(4), 4);
        assert_eq!(next_lower_power_of_two(8), 8);
        assert_eq!(next_lower_power_of_two(16), 16);
        assert_eq!(next_lower_power_of_two(256), 256);
        assert_eq!(next_lower_power_of_two(65536), 65536);

        assert_eq!(next_lower_power_of_two(3), 2);
        assert_eq!(next_lower_power_of_two(5), 4);
        assert_eq!(next_lower_power_of_two(7), 4);
        assert_eq!(next_lower_power_of_two(9), 8);
        assert_eq!(next_lower_power_of_two(15), 8);
        assert_eq!(next_lower_power_of_two(17), 16);
        assert_eq!(next_lower_power_of_two(31), 16);
        assert_eq!(next_lower_power_of_two(100), 64);
        assert_eq!(next_lower_power_of_two(1000), 512);

        assert_eq!(next_lower_power_of_two(u64::MAX), 1u64 << 63);
        assert_eq!(next_lower_power_of_two((1u64 << 63) + 1), 1u64 << 63);
        assert_eq!(next_lower_power_of_two((1u64 << 62) + 1), 1u64 << 62);
    }
}