newton-core 0.7.5

use alloy::{primitives::Address, providers::Provider, sol};

use crate::newton_prover_task_manager::NewtonProverTaskManager;

sol! {
    #[sol(rpc)]
    interface ITaskResponseHandlerVerifier {
        function certificateVerifier() external view returns (address);
    }
}

/// Resolved destination-chain verifier wiring derived from TaskManager.
#[derive(Debug, Clone, Copy)]
pub struct ResolvedDestinationVerifier {
    /// task response handler currently set on task manager
    pub task_response_handler: Address,
    /// certificate verifier currently set on task response handler
    pub certificate_verifier: Address,
    /// source chain AVS service manager read from task manager
    pub source_chain_avs: Address,
}

/// Resolve active destination verification wiring from task manager state.
pub async fn resolve_destination_verifier<P>(
    task_manager_addr: Address,
    provider: P,
) -> Result<ResolvedDestinationVerifier, eyre::Error>
where
    P: Provider + Clone,
{
    let task_manager = NewtonProverTaskManager::new(task_manager_addr, provider.clone());
    let task_response_handler = task_manager
        .taskResponseHandler()
        .call()
        .await
        .map_err(|e| eyre::eyre!("failed to get task response handler address: {:?}", e))?;

    let task_response_handler_contract = ITaskResponseHandlerVerifier::new(task_response_handler, provider.clone());
    let certificate_verifier = task_response_handler_contract
        .certificateVerifier()
        .call()
        .await
        .map_err(|e| eyre::eyre!("failed to get certificate verifier from task response handler: {:?}", e))?;

    let source_chain_avs = task_manager
        .serviceManager()
        .call()
        .await
        .map_err(|e| eyre::eyre!("failed to get source chain avs address: {:?}", e))?;

    Ok(ResolvedDestinationVerifier {
        task_response_handler,
        certificate_verifier,
        source_chain_avs,
    })
}

/// How long a resolved destination wiring is reused before it is read again, when no
/// [`TaskResponseHandlerSource`] can name the current handler.
///
/// The wiring is three `eth_call`s (`taskResponseHandler`, `certificateVerifier`,
/// `serviceManager`) that every destination-chain task would otherwise repeat. It only
/// changes on an owner-gated `updateTaskResponseHandler`, so a short TTL bounds how long
/// such a change takes to be picked up.
#[cfg(feature = "rpc")]
pub const DESTINATION_VERIFIER_TTL: std::time::Duration = std::time::Duration::from_secs(60);

/// Names a destination task manager's current response handler without a chain read.
///
/// The handler is the only mutable part of the wiring: the certificate verifier is
/// immutable on the handler and the source AVS is set once on the task manager. So with the
/// handler known, the wiring for it is read once per process and never expires.
#[cfg(feature = "rpc")]
#[async_trait::async_trait]
pub trait TaskResponseHandlerSource: Send + Sync {
    /// `task_manager`'s current handler on `chain_id`, or `None` when the source cannot say
    /// (the TTL-cached chain read then answers).
    async fn current_handler(&self, chain_id: u64, task_manager: Address) -> Option<Address>;
}

#[cfg(feature = "rpc")]
static HANDLER_SOURCE: std::sync::OnceLock<std::sync::Arc<dyn TaskResponseHandlerSource>> = std::sync::OnceLock::new();

/// Install the process-wide [`TaskResponseHandlerSource`]. The first call wins; later calls
/// are ignored and return `false`.
#[cfg(feature = "rpc")]
pub fn set_task_response_handler_source(source: std::sync::Arc<dyn TaskResponseHandlerSource>) -> bool {
    HANDLER_SOURCE.set(source).is_ok()
}

#[cfg(feature = "rpc")]
type VerifierCache =
    std::sync::Mutex<std::collections::HashMap<(Address, String), (std::time::Instant, ResolvedDestinationVerifier)>>;

/// `(task_manager, handler, rpc_url)` → wiring. Never expires: nothing in it can change
/// for a given handler.
#[cfg(feature = "rpc")]
type HandlerWiringCache =
    std::sync::Mutex<std::collections::HashMap<(Address, Address, String), ResolvedDestinationVerifier>>;

#[cfg(feature = "rpc")]
static DESTINATION_VERIFIERS: std::sync::LazyLock<VerifierCache> = std::sync::LazyLock::new(Default::default);

#[cfg(feature = "rpc")]
static HANDLER_WIRING: std::sync::LazyLock<HandlerWiringCache> = std::sync::LazyLock::new(Default::default);

/// [`resolve_destination_verifier`] without per-task chain reads.
///
/// When the installed [`TaskResponseHandlerSource`] names the current handler, the wiring
/// for that handler is read once and reused for the life of the process. Otherwise it falls
/// back to a cache keyed by `(task_manager_addr, rpc_url)` refreshed after
/// [`DESTINATION_VERIFIER_TTL`].
///
/// Traced as `chain.destination_verifier`; `cache` is indexed (handler from the source,
/// wiring cached)|indexed_miss (handler from the source, wiring read now)|hit|miss.
#[cfg(feature = "rpc")]
#[tracing::instrument(
    name = "chain.destination_verifier",
    target = "newton::task_evaluation",
    level = "debug",
    skip(rpc_url),
    fields(chain_id, %task_manager_addr, cache = tracing::field::Empty)
)]
pub async fn resolve_destination_verifier_cached(
    chain_id: u64,
    task_manager_addr: Address,
    rpc_url: &str,
) -> Result<ResolvedDestinationVerifier, eyre::Error> {
    let span = tracing::Span::current();
    let handler = match HANDLER_SOURCE.get() {
        Some(source) => source.current_handler(chain_id, task_manager_addr).await,
        None => None,
    };
    if let Some(handler) = handler {
        let key = (task_manager_addr, handler, rpc_url.to_string());
        let cached = HANDLER_WIRING
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
            .get(&key)
            .copied();
        if let Some(resolved) = cached {
            span.record("cache", "indexed");
            return Ok(resolved);
        }
        span.record("cache", "indexed_miss");
        let resolved = resolve_handler_wiring(task_manager_addr, handler, rpc_url).await?;
        HANDLER_WIRING
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
            .insert(key, resolved);
        return Ok(resolved);
    }

    let key = (task_manager_addr, rpc_url.to_string());
    let cached = DESTINATION_VERIFIERS
        .lock()
        .unwrap_or_else(std::sync::PoisonError::into_inner)
        .get(&key)
        .copied();
    if let Some((at, resolved)) = cached {
        if at.elapsed() < DESTINATION_VERIFIER_TTL {
            span.record("cache", "hit");
            return Ok(resolved);
        }
    }
    span.record("cache", "miss");
    let resolved = resolve_destination_verifier(task_manager_addr, newton_rpc_provider::get_provider(rpc_url)).await?;
    DESTINATION_VERIFIERS
        .lock()
        .unwrap_or_else(std::sync::PoisonError::into_inner)
        .insert(key, (std::time::Instant::now(), resolved));
    Ok(resolved)
}

/// The immutable wiring behind a known `handler`: its certificate verifier and the task
/// manager's source AVS.
#[cfg(feature = "rpc")]
async fn resolve_handler_wiring(
    task_manager_addr: Address,
    handler: Address,
    rpc_url: &str,
) -> Result<ResolvedDestinationVerifier, eyre::Error> {
    let provider = newton_rpc_provider::get_provider(rpc_url);
    let task_manager = NewtonProverTaskManager::new(task_manager_addr, provider.clone());
    let handler_contract = ITaskResponseHandlerVerifier::new(handler, provider.clone());
    let (certificate_verifier, source_chain_avs) = tokio::try_join!(
        async {
            handler_contract
                .certificateVerifier()
                .call()
                .await
                .map_err(|e| eyre::eyre!("failed to get certificate verifier from task response handler: {:?}", e))
        },
        async {
            task_manager
                .serviceManager()
                .call()
                .await
                .map_err(|e| eyre::eyre!("failed to get source chain avs address: {:?}", e))
        },
    )?;
    Ok(ResolvedDestinationVerifier {
        task_response_handler: handler,
        certificate_verifier,
        source_chain_avs,
    })
}