pe_graph/

matrix_hook.rs

//! Matrix integration hook — trait for the optional matrix layer.
//!
//! pe-graph defines [`MatrixHook`] as the extension point for convergence
//! tracking and learned routing. pe-matrix provides the implementation;
//! pe-runtime wires them together. pe-graph never depends on pe-matrix.
//!
//! When no hook is provided, the Pregel engine handles `NodeResult::Converge`
//! by applying the partial_update (degrading to Update) and uses the user's
//! conditional edge function for routing — unchanged from pre-matrix behavior.

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

/// Hook for the optional matrix layer.
///
/// Injected into the Pregel engine by pe-runtime when the matrix layer is
/// active. The engine calls these methods at the right moments; the hook
/// implementation does the actual convergence tracking and routing.
///
/// `Send + Sync` because the engine may run nodes in parallel.
///
/// # Thread safety
///
/// The hook is shared across the BSP loop (single-threaded per superstep)
/// but must be `Send + Sync` to satisfy trait object requirements. Interior
/// mutability (via Mutex) is expected since the BSP loop is sequential.
pub trait MatrixHook: Send + Sync {
    /// Called when a node returns `NodeResult::Converge`.
    ///
    /// Records the convergence signal metadata. The engine has already
    /// extracted and applied the partial_update — this receives only
    /// the signal values (contribution, surprise, quality) and the
    /// originating node name.
    fn on_converge(&self, node_name: &str, actual_contribution: f64, surprise: f64, quality: f64);

    /// Resolve routing for a conditional edge using learned probabilities.
    ///
    /// Called instead of the user's router function when the matrix layer
    /// is active. `from` is the source node, `candidates` are the possible
    /// target nodes (from the conditional edge's router output).
    ///
    /// Returns the selected candidate(s). If this returns `None`, the engine
    /// falls back to the user's router function (graceful degradation).
    fn route(&self, from: &str, candidates: &[String]) -> Option<Vec<String>>;

    /// Current C value (aggregate confidence).
    fn c_value(&self) -> f64;

    /// Current completion estimate.
    fn completion(&self) -> f64;

    /// Whether convergence threshold has been reached.
    fn is_converged(&self) -> bool;

    /// Record that a transition occurred (for learning).
    fn record_transition(&self, from: &str, to: &str, quality: f64);
}

/// Wrapper around `Arc<dyn MatrixHook>` for ergonomic use in the engine.
///
/// Cloneable, shareable, optional.
#[derive(Clone)]
pub struct MatrixHookHandle(pub(crate) Arc<dyn MatrixHook>);

impl MatrixHookHandle {
    /// Create a new handle wrapping a hook implementation.
    pub fn new(hook: impl MatrixHook + 'static) -> Self {
        Self(Arc::new(hook))
    }

    /// Delegate to the inner hook.
    pub fn on_converge(
        &self,
        node_name: &str,
        actual_contribution: f64,
        surprise: f64,
        quality: f64,
    ) {
        self.0
            .on_converge(node_name, actual_contribution, surprise, quality);
    }

    /// Delegate routing to the inner hook.
    pub fn route(&self, from: &str, candidates: &[String]) -> Option<Vec<String>> {
        self.0.route(from, candidates)
    }

    /// Current C value.
    pub fn c_value(&self) -> f64 {
        self.0.c_value()
    }

    /// Current completion estimate.
    pub fn completion(&self) -> f64 {
        self.0.completion()
    }

    /// Whether convergence has been reached.
    pub fn is_converged(&self) -> bool {
        self.0.is_converged()
    }

    /// Record a transition observation.
    pub fn record_transition(&self, from: &str, to: &str, quality: f64) {
        self.0.record_transition(from, to, quality);
    }
}

impl std::fmt::Debug for MatrixHookHandle {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("MatrixHookHandle")
            .field("c_value", &self.c_value())
            .field("completion", &self.completion())
            .field("is_converged", &self.is_converged())
            .finish()
    }
}

/// Concrete implementation that bridges pe-matrix types into the hook trait.
///
/// Created by pe-runtime when wiring up the matrix layer. Uses `Mutex`
/// for interior mutability since the BSP loop is sequential.
///
/// # Thread safety assumption
///
/// This type assumes **single-threaded sequential access** within each
/// BSP superstep. The `Mutex` satisfies the `Sync` bound but is never
/// contended in practice. If a `Mutex::lock()` fails (poisoned), the
/// hook silently returns defaults — this can only happen if the engine
/// panics while holding the lock, which the BSP loop prevents.
pub struct DefaultMatrixHook<T: ConvergenceRecorder, R: RoutingResolver> {
    tracker: Mutex<T>,
    router: Mutex<R>,
}

impl<T: ConvergenceRecorder, R: RoutingResolver> DefaultMatrixHook<T, R> {
    /// Create a new hook from a tracker and router.
    pub fn new(tracker: T, router: R) -> Self {
        Self {
            tracker: Mutex::new(tracker),
            router: Mutex::new(router),
        }
    }
}

// Mutex<T>: Sync when T: Send — no extra +Sync bounds needed on T/R.
impl<T: ConvergenceRecorder, R: RoutingResolver> MatrixHook for DefaultMatrixHook<T, R> {
    fn on_converge(&self, _node_name: &str, actual_contribution: f64, surprise: f64, quality: f64) {
        if let Ok(mut t) = self.tracker.lock() {
            t.record(actual_contribution, surprise, quality);
        }
    }

    fn route(&self, from: &str, candidates: &[String]) -> Option<Vec<String>> {
        self.router.lock().ok()?.resolve(from, candidates)
    }

    fn c_value(&self) -> f64 {
        self.tracker.lock().map(|t| t.c_value()).unwrap_or(0.0)
    }

    fn completion(&self) -> f64 {
        self.tracker.lock().map(|t| t.completion()).unwrap_or(0.0)
    }

    fn is_converged(&self) -> bool {
        self.tracker
            .lock()
            .map(|t| t.is_converged())
            .unwrap_or(false)
    }

    fn record_transition(&self, from: &str, to: &str, quality: f64) {
        if let Ok(mut r) = self.router.lock() {
            r.learn(from, to, quality);
        }
    }
}

/// Trait for convergence tracking — implemented by pe-matrix ConvergenceTracker.
pub trait ConvergenceRecorder: Send {
    /// Record a convergence observation.
    fn record(&mut self, actual_contribution: f64, surprise: f64, quality: f64);
    /// Current C value.
    fn c_value(&self) -> f64;
    /// Current completion estimate.
    fn completion(&self) -> f64;
    /// Whether convergence threshold is met.
    fn is_converged(&self) -> bool;
}

/// Trait for routing resolution — implemented by pe-matrix MatrixRouter.
pub trait RoutingResolver: Send {
    /// Select candidates using learned routing. Returns None to fall back.
    fn resolve(&self, from: &str, candidates: &[String]) -> Option<Vec<String>>;
    /// Record a transition for learning.
    fn learn(&mut self, from: &str, to: &str, quality: f64);
}