pepper-sync 0.7.0

use std::collections::{BTreeSet, HashMap};
use std::ops::Range;
use std::sync::Arc;

use tokio::sync::{RwLock, mpsc};

use zcash_keys::keys::UnifiedFullViewingKey;
use zcash_protocol::consensus::{self, BlockHeight};
use zcash_transparent::keys::NonHardenedChildIndex;
use zip32::AccountId;

use crate::client::{self, FetchRequest};
use crate::config::TransparentAddressDiscovery;
use crate::error::SyncError;
use crate::keys;
use crate::keys::transparent::{TransparentAddressId, TransparentScope};
use crate::wallet::traits::SyncWallet;
use crate::wallet::{KeyIdInterface, ScanTarget};

use super::MAX_REORG_ALLOWANCE;

/// Discovers all addresses in use by the wallet and returns `scan_targets` for any new relevant transactions to scan transparent
/// bundles.
/// `last_known_chain_height` should be the value before updating to latest chain height.
/// The search starts at the height returned by [`discovery_start_height`].
/// Returns the gap addresses.
// TODO: improve this to not make all the calls at once
pub(crate) async fn address_discovery<W: SyncWallet>(
    consensus_parameters: &impl consensus::Parameters,
    wallet: Arc<RwLock<W>>,
    fetch_request_sender: mpsc::UnboundedSender<FetchRequest>,
    ufvks: &HashMap<AccountId, UnifiedFullViewingKey>,
    last_known_chain_height: BlockHeight,
    chain_height: BlockHeight,
    config: TransparentAddressDiscovery,
) -> Result<HashMap<String, TransparentAddressId>, SyncError<W::Error>> {
    if !config.scopes.external && !config.scopes.internal && !config.scopes.refund {
        return Ok(HashMap::new());
    }

    let wallet_guard = wallet.read().await;
    let wallet_addresses = wallet_guard
        .get_transparent_addresses()
        .map_err(SyncError::WalletError)?
        .clone();
    let transparent_scan_floor = wallet_guard
        .get_sync_state()
        .map_err(SyncError::WalletError)?
        .transparent_scan_floor;
    drop(wallet_guard);
    let mut scan_targets: BTreeSet<ScanTarget> = BTreeSet::new();
    let block_range = Range {
        start: discovery_start_height(
            consensus_parameters,
            last_known_chain_height,
            transparent_scan_floor,
        ),
        end: chain_height + 1,
    };

    // find scan_targets for any new transactions relevant to known addresses
    for address in wallet_addresses.values() {
        let transactions = client::get_transparent_address_transactions(
            fetch_request_sender.clone(),
            consensus_parameters,
            address.clone(),
            block_range.clone(),
        )
        .await?;

        // The transaction is not scanned here, instead the scan target is stored to be later sent to a scan task for these reasons:
        // - We must search for all relevant transactions MAX_REORG_ALLOWANCE blocks below wallet height in case of re-org.
        // These would be scanned again which would be inefficient
        // - In case of re-org, any scanned transactions with heights within the re-org range would be wrongly invalidated
        // - The scan target will cause the surrounding range to be set to high priority which will often also contain shielded notes
        // relevant to the wallet
        // - Scanning a transaction without scanning the surrounding range of compact blocks in the context of a scan task creates
        // complications. Instead of writing all the information into a wallet transaction once, it would result in "incomplete"
        // transactions that only contain transparent outputs and must be updated with shielded notes and other data when scanned.
        // - We would need to add additional processing here to fetch the compact block for transaction metadata such as block time
        // and append this to the wallet.
        // - It allows SyncState to maintain complete knowledge and control of all the tasks that have and will be performed by the
        // sync engine.
        //
        // To summarise, keeping transaction scanning within the scanner is much better co-ordinated and allows us to leverage
        // any new developments to sync state management and scanning. It also separates concerns, with tasks happening in one
        // place and performed once, wherever possible.
        for (height, tx) in &transactions {
            scan_targets.insert(ScanTarget::internal(*height, tx.txid()));
        }
    }

    let mut scopes = Vec::new();
    if config.scopes.external {
        scopes.push(TransparentScope::External);
    }
    if config.scopes.internal {
        scopes.push(TransparentScope::Internal);
    }
    if config.scopes.refund {
        scopes.push(TransparentScope::Refund);
    }

    // discover new addresses and find scan_targets for relevant transactions
    let mut gap_addresses = HashMap::new();
    for (account_id, ufvk) in ufvks {
        if let Some(account_pubkey) = ufvk.p2pkh() {
            for scope in &scopes {
                // start with the first address index previously unused by the wallet
                let mut address_index = if let Some(id) = wallet_addresses
                    .keys()
                    .rfind(|id| id.account_id() == *account_id && id.scope() == *scope)
                {
                    id.address_index().next()
                } else {
                    Some(NonHardenedChildIndex::ZERO)
                }
                .ok_or_else(|| {
                    SyncError::TransparentAddressDerivationError(bip32::Error::ChildNumber)
                })?;
                let mut unused_address_count: usize = 0;
                let mut addresses: Vec<(TransparentAddressId, String)> = Vec::new();

                while unused_address_count < config.gap_limit as usize {
                    let address_id = TransparentAddressId::new(*account_id, *scope, address_index);
                    let address = keys::transparent::derive_address(
                        consensus_parameters,
                        account_pubkey,
                        address_id,
                    )
                    .map_err(SyncError::TransparentAddressDerivationError)?;
                    addresses.push((address_id, address.clone()));

                    let transactions = client::get_transparent_address_transactions(
                        fetch_request_sender.clone(),
                        consensus_parameters,
                        address,
                        block_range.clone(),
                    )
                    .await?;

                    if transactions.is_empty() {
                        unused_address_count += 1;
                    } else {
                        for (height, tx) in &transactions {
                            scan_targets.insert(ScanTarget::internal(*height, tx.txid()));
                        }
                        unused_address_count = 0;
                    }

                    address_index = address_index.next().ok_or_else(|| {
                        SyncError::TransparentAddressDerivationError(bip32::Error::ChildNumber)
                    })?;
                }

                let gap_index = addresses.len().saturating_sub(config.gap_limit as usize);
                let scope_gap_addresses: HashMap<String, TransparentAddressId> = addresses
                    .split_off(gap_index)
                    .into_iter()
                    .map(|(id, address)| (address, id))
                    .collect();
                gap_addresses.extend(scope_gap_addresses);

                let mut wallet_guard = wallet.write().await;
                let wallet_addresses_mut = wallet_guard
                    .get_transparent_addresses_mut()
                    .map_err(SyncError::WalletError)?;
                for (id, address) in addresses {
                    wallet_addresses_mut.insert(id, address);
                }
            }
        }
    }

    let mut wallet_guard = wallet.write().await;
    let sync_state = wallet_guard
        .get_sync_state_mut()
        .map_err(SyncError::WalletError)?;
    for scan_target in scan_targets {
        sync_state.insert_scan_target(scan_target);
    }
    wallet_guard
        .set_save_flag()
        .map_err(SyncError::WalletError)?;

    Ok(gap_addresses)
}

/// Returns the height transparent address discovery searches from.
///
/// The search starts [`MAX_REORG_ALLOWANCE`] blocks below the `last_known_chain_height` in case of re-org.
/// It starts from the block above the `transparent_scan_floor` of the previous sync session if that is lower, as the
/// blocks mined during that session were only covered by compact block transparent data scanning and may still be
/// unscanned.
/// The sapling activation height is the lowest height searched.
fn discovery_start_height(
    consensus_parameters: &impl consensus::Parameters,
    last_known_chain_height: BlockHeight,
    transparent_scan_floor: Option<BlockHeight>,
) -> BlockHeight {
    let sapling_activation_height = consensus_parameters
        .activation_height(consensus::NetworkUpgrade::Sapling)
        .expect("sapling activation height should always return Some");
    let reorg_start_height = last_known_chain_height.saturating_sub(MAX_REORG_ALLOWANCE) + 1;

    transparent_scan_floor
        .map_or(reorg_start_height, |floor| {
            reorg_start_height.min(floor + 1)
        })
        .max(sapling_activation_height)
}

// TODO: process memo encoded address indexes.
// 1. return any memo address ids from scan in ScanResults
// 2. derive the addresses up to that index, add to wallet addresses and send them to GetTaddressTxids
// 3. for each transaction returned:
// a) if the tx is in a range that is not scanned, add scan_targets to sync_state
// b) if the range is scanned and the tx is already in the wallet, rescan the zcash transaction transparent bundles in
// the wallet transaction
// c) if the range is scanned and the tx does not exist in the wallet, fetch the compact block if its not in the wallet
// and scan the transparent bundles

#[cfg(test)]
mod tests {
    use zcash_protocol::consensus::{MAIN_NETWORK, Parameters as _};

    use super::*;

    const LAST_KNOWN_CHAIN_HEIGHT: u32 = 3_000_000;
    const REORG_START_HEIGHT: BlockHeight =
        BlockHeight::from_u32(LAST_KNOWN_CHAIN_HEIGHT - MAX_REORG_ALLOWANCE + 1);
    const BLOCKS_MINED_BELOW_REORG_START: u32 = 500;
    const BLOCKS_MINED_ABOVE_SAPLING_ACTIVATION: u32 = 10;

    fn start_height(transparent_scan_floor: Option<BlockHeight>) -> BlockHeight {
        discovery_start_height(
            &MAIN_NETWORK,
            BlockHeight::from_u32(LAST_KNOWN_CHAIN_HEIGHT),
            transparent_scan_floor,
        )
    }

    /// A wallet saved before the transparent scan floor was stored has no floor. The search covers the blocks that
    /// may have been re-orged.
    #[test]
    fn discovery_starts_at_reorg_allowance_with_no_floor() {
        assert_eq!(start_height(None), REORG_START_HEIGHT);
    }

    /// Fewer blocks were mined during the previous sync session than may have been re-orged.
    #[test]
    fn discovery_starts_at_reorg_allowance_with_floor_above() {
        for floor in [
            REORG_START_HEIGHT,
            BlockHeight::from_u32(LAST_KNOWN_CHAIN_HEIGHT),
        ] {
            assert_eq!(start_height(Some(floor)), REORG_START_HEIGHT);
        }
    }

    /// More blocks were mined during the previous sync session than may have been re-orged. Any of them may still be
    /// unscanned, so the search starts from the first block mined during that session.
    #[test]
    fn discovery_starts_above_floor_below_reorg_allowance() {
        let floor = REORG_START_HEIGHT - BLOCKS_MINED_BELOW_REORG_START;

        assert_eq!(start_height(Some(floor)), floor + 1);
    }

    #[test]
    fn discovery_starts_at_sapling_activation_at_the_lowest() {
        let sapling_activation_height = MAIN_NETWORK
            .activation_height(consensus::NetworkUpgrade::Sapling)
            .unwrap();

        assert_eq!(
            discovery_start_height(
                &MAIN_NETWORK,
                sapling_activation_height + BLOCKS_MINED_ABOVE_SAPLING_ACTIVATION,
                None
            ),
            sapling_activation_height
        );
    }
}