tenferro-runtime 0.9.0

//! Runtime-owned context passed to prepared extension operations.
//!
//! Extension dispatch is owned by [`crate::Runtime`] through installed
//! [`crate::ExtensionModule`] values. This module intentionally exposes only the
//! backend/cache context that prepared operations receive at execution time.

use std::fmt;

use tenferro_tensor::Engine;

use crate::extension_cache::ExtensionCacheStore;

/// Backend and cache state passed to one prepared extension execution.
///
/// Extension crates should obtain this value from their hidden
/// [`crate::PreparedOperationExecutor`] bridge and use it only for the duration
/// of that call.
///
/// # Examples
///
/// ```rust
/// use tenferro_cpu::CpuEngine;
/// use tenferro_tensor::{Tensor, TensorElementwise, TensorRead};
/// use tenferro_runtime::{
///     ExtensionCacheSelector, ExtensionCacheStore, ExtensionExecutionContext,
/// };
///
/// let mut engine = CpuEngine::serial();
/// let mut caches = ExtensionCacheStore::new();
/// let mut context = ExtensionExecutionContext::new(&mut engine, &mut caches);
/// let lhs = Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap();
/// let rhs = Tensor::from_vec_col_major(vec![2], vec![3.0_f64, 4.0]).unwrap();
/// let output = context
///     .backend_mut()
///     .add_read(TensorRead::from_tensor(&lhs), TensorRead::from_tensor(&rhs))
///     .unwrap();
///
/// assert_eq!(output.as_slice::<f64>().unwrap(), &[4.0, 6.0]);
/// assert_eq!(context.caches().stats(ExtensionCacheSelector::All).entries, 0);
/// # Ok::<(), Box<dyn std::error::Error>>(())
/// ```
///
/// The context cannot outlive the engine and cache state it borrows.
///
/// ```compile_fail
/// use tenferro_cpu::CpuEngine;
/// use tenferro_runtime::{ExtensionCacheStore, ExtensionExecutionContext};
///
/// fn leak_context() -> ExtensionExecutionContext<'static, CpuEngine> {
///     let mut engine = CpuEngine::serial();
///     let mut caches = ExtensionCacheStore::new();
///     ExtensionExecutionContext::new(&mut engine, &mut caches)
/// }
/// ```
pub struct ExtensionExecutionContext<'a, B: Engine + ?Sized> {
    backend: &'a mut B,
    caches: &'a mut ExtensionCacheStore,
}

impl<B: Engine + ?Sized> fmt::Debug for ExtensionExecutionContext<'_, B> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("ExtensionExecutionContext")
            .field("backend_type", &std::any::type_name::<B>())
            .field("caches", &self.caches)
            .finish_non_exhaustive()
    }
}

impl<'a, B: Engine + ?Sized> ExtensionExecutionContext<'a, B> {
    /// Build a context from externally-owned backend and cache state.
    pub fn new(backend: &'a mut B, caches: &'a mut ExtensionCacheStore) -> Self {
        Self { backend, caches }
    }

    /// Borrow the backend for non-mutating inspection.
    pub fn backend(&self) -> &B {
        self.backend
    }

    /// Borrow the backend mutably for extension execution.
    pub fn backend_mut(&mut self) -> &mut B {
        self.backend
    }

    /// Borrow the extension runtime cache store.
    pub fn caches(&self) -> &ExtensionCacheStore {
        self.caches
    }

    /// Borrow the extension runtime cache store mutably.
    pub fn caches_mut(&mut self) -> &mut ExtensionCacheStore {
        self.caches
    }

    /// Borrow backend and extension cache store as disjoint mutable parts.
    pub fn parts_mut(&mut self) -> (&mut B, &mut ExtensionCacheStore) {
        (self.backend, self.caches)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use tenferro_cpu::CpuEngine;
    use tenferro_tensor::{Engine, Tensor, TensorElementwise, TensorRead};

    use crate::ExtensionCacheSelector;

    #[test]
    fn context_accepts_a_caller_held_engine() {
        let mut engine = CpuEngine::serial();
        let mut caches = ExtensionCacheStore::new();

        let mut context = ExtensionExecutionContext::new(&mut engine, &mut caches);
        let _: &dyn Engine = context.backend();
        let lhs = Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap();
        let rhs = Tensor::from_vec_col_major(vec![2], vec![3.0_f64, 4.0]).unwrap();
        let output = context
            .backend_mut()
            .add_read(TensorRead::from_tensor(&lhs), TensorRead::from_tensor(&rhs))
            .unwrap();

        assert_eq!(output.as_slice::<f64>().unwrap(), &[4.0, 6.0]);
        assert_eq!(
            context.caches().stats(ExtensionCacheSelector::All).entries,
            0
        );

        let (_, caches) = context.parts_mut();
        assert_eq!(caches.stats(ExtensionCacheSelector::All).entries, 0);
    }
}