zcash_voting 0.5.2

//! Vote commitment tree sync and VAN witness generation.
//!
//! Manages per-round in-memory `TreeClient` instances that sync incrementally
//! from a chain node via HTTP, then generates Merkle authentication paths
//! (witnesses) for Vote Authority Notes (VANs) needed by ZKP #2.

use std::collections::HashMap;
use std::sync::{Arc, Mutex};

use vote_commitment_tree::{MerklePath, TreeClient};
use vote_commitment_tree_client::http_sync_api::HttpTreeSyncApi;

use crate::storage::VotingDb;
use crate::types::VotingError;
use crate::HyperTransport;

/// VAN Merkle witness for ZKP #2.
///
/// Contains the authentication path, leaf position, and anchor height needed
/// by `build_vote_commitment`. Generated by `generate_van_witness` after syncing
/// the vote commitment tree.
#[derive(Clone, Debug)]
pub struct VanWitness {
    /// 24 sibling hashes (32 bytes each) from leaf to root.
    pub auth_path: Vec<[u8; 32]>,
    /// Leaf position of the VAN in the tree.
    pub position: u32,
    /// Block height at which the tree was snapshotted.
    pub anchor_height: u32,
}

impl From<(MerklePath, u32)> for VanWitness {
    fn from((path, anchor_height): (MerklePath, u32)) -> Self {
        let auth_path = path
            .auth_path()
            .iter()
            .map(|h| h.to_bytes())
            .collect();
        Self {
            auth_path,
            position: path.position(),
            anchor_height,
        }
    }
}

/// Manages per-round in-memory vote commitment trees for VAN witness generation.
///
/// Wraps a `HashMap<String, TreeClient>` behind a `Mutex` for thread-safe
/// per-round incremental sync. Each round gets its own `TreeClient`, created
/// lazily on first `sync` call for that round.
pub struct VoteTreeSync {
    clients: Mutex<HashMap<String, TreeClient>>,
    transport: Arc<HyperTransport>,
}

impl VoteTreeSync {
    pub fn new() -> Self {
        Self {
            clients: Mutex::new(HashMap::new()),
            transport: Arc::new(HyperTransport::new()),
        }
    }

    /// Sync the vote commitment tree for a specific round from a chain node.
    ///
    /// Creates a per-round `TreeClient` on first call, then syncs incrementally
    /// on subsequent calls. VAN positions from ALL bundles are automatically
    /// marked for witness generation before syncing.
    ///
    /// Returns the latest synced block height.
    pub fn sync(&self, db: &VotingDb, round_id: &str, node_url: &str) -> Result<u32, VotingError> {
        let bundle_count = db.get_bundle_count(round_id)?;

        let mut guard = self.clients.lock().map_err(|e| VotingError::Internal {
            message: format!("tree client lock poisoned: {}", e),
        })?;

        let client = guard
            .entry(round_id.to_string())
            .or_insert_with(TreeClient::empty);

        for bi in 0..bundle_count {
            if let Ok(pos) = db.load_van_position(round_id, bi) {
                client.mark_position(pos as u64);
            }
        }

        let api = HttpTreeSyncApi::new(node_url, round_id, self.transport.clone());
        client.sync(&api).map_err(|e| VotingError::Internal {
            message: format!("vote tree sync failed: {}", e),
        })?;

        // Empty tree is valid before the first delegation commitment is appended.
        // Report height 0 so callers can proceed instead of failing sync.
        Ok(client.last_synced_height().unwrap_or(0))
    }

    /// Generate a VAN Merkle witness for ZKP #2.
    ///
    /// Requires `sync` to have been called first for this round. Loads the VAN
    /// position for the specified bundle and generates a witness at the given
    /// anchor height.
    pub fn generate_van_witness(
        &self,
        db: &VotingDb,
        round_id: &str,
        bundle_index: u32,
        anchor_height: u32,
    ) -> Result<VanWitness, VotingError> {
        let van_position = db.load_van_position(round_id, bundle_index)?;

        let guard = self.clients.lock().map_err(|e| VotingError::Internal {
            message: format!("tree client lock poisoned: {}", e),
        })?;

        let client = guard.get(round_id).ok_or_else(|| VotingError::InvalidInput {
            message: "must call sync before generate_van_witness".to_string(),
        })?;

        let path = client
            .witness(van_position as u64, anchor_height)
            .ok_or_else(|| VotingError::Internal {
                message: format!(
                    "failed to generate witness for position {} at height {}",
                    van_position, anchor_height
                ),
            })?;

        Ok(VanWitness::from((path, anchor_height)))
    }

    /// Drop the in-memory TreeClient for a round so the next `sync` call
    /// creates a fresh one and does a full resync. This recovers from stale
    /// state that would otherwise cause `StartIndexMismatch` or `RootMismatch`.
    /// If `round_id` is empty, all clients are dropped.
    pub fn reset(&self, round_id: &str) -> Result<(), VotingError> {
        let mut guard = self.clients.lock().map_err(|e| VotingError::Internal {
            message: format!("tree client lock poisoned: {}", e),
        })?;
        if round_id.is_empty() {
            guard.clear();
        } else {
            guard.remove(round_id);
        }
        Ok(())
    }
}