Skip to main content

tenferro_runtime/
extension_execution_context.rs

1//! Runtime-owned context passed to prepared extension operations.
2//!
3//! Extension dispatch is owned by [`crate::Runtime`] through installed
4//! [`crate::ExtensionModule`] values. This module intentionally exposes only the
5//! backend/cache context that prepared operations receive at execution time.
6
7use std::fmt;
8
9use tenferro_tensor::Engine;
10
11use crate::extension_cache::ExtensionCacheStore;
12
13/// Backend and cache state passed to one prepared extension execution.
14///
15/// Extension crates should obtain this value from their hidden
16/// [`crate::PreparedOperationExecutor`] bridge and use it only for the duration
17/// of that call.
18///
19/// # Examples
20///
21/// ```rust
22/// use tenferro_cpu::CpuEngine;
23/// use tenferro_tensor::{Tensor, TensorElementwise, TensorRead};
24/// use tenferro_runtime::{
25///     ExtensionCacheSelector, ExtensionCacheStore, ExtensionExecutionContext,
26/// };
27///
28/// let mut engine = CpuEngine::serial();
29/// let mut caches = ExtensionCacheStore::new();
30/// let mut context = ExtensionExecutionContext::new(&mut engine, &mut caches);
31/// let lhs = Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap();
32/// let rhs = Tensor::from_vec_col_major(vec![2], vec![3.0_f64, 4.0]).unwrap();
33/// let output = context
34///     .backend_mut()
35///     .add_read(TensorRead::from_tensor(&lhs), TensorRead::from_tensor(&rhs))
36///     .unwrap();
37///
38/// assert_eq!(output.as_slice::<f64>().unwrap(), &[4.0, 6.0]);
39/// assert_eq!(context.caches().stats(ExtensionCacheSelector::All).entries, 0);
40/// # Ok::<(), Box<dyn std::error::Error>>(())
41/// ```
42///
43/// The context cannot outlive the engine and cache state it borrows.
44///
45/// ```compile_fail
46/// use tenferro_cpu::CpuEngine;
47/// use tenferro_runtime::{ExtensionCacheStore, ExtensionExecutionContext};
48///
49/// fn leak_context() -> ExtensionExecutionContext<'static, CpuEngine> {
50///     let mut engine = CpuEngine::serial();
51///     let mut caches = ExtensionCacheStore::new();
52///     ExtensionExecutionContext::new(&mut engine, &mut caches)
53/// }
54/// ```
55pub struct ExtensionExecutionContext<'a, B: Engine + ?Sized> {
56    backend: &'a mut B,
57    caches: &'a mut ExtensionCacheStore,
58}
59
60impl<B: Engine + ?Sized> fmt::Debug for ExtensionExecutionContext<'_, B> {
61    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
62        f.debug_struct("ExtensionExecutionContext")
63            .field("backend_type", &std::any::type_name::<B>())
64            .field("caches", &self.caches)
65            .finish_non_exhaustive()
66    }
67}
68
69impl<'a, B: Engine + ?Sized> ExtensionExecutionContext<'a, B> {
70    /// Build a context from externally-owned backend and cache state.
71    pub fn new(backend: &'a mut B, caches: &'a mut ExtensionCacheStore) -> Self {
72        Self { backend, caches }
73    }
74
75    /// Borrow the backend for non-mutating inspection.
76    pub fn backend(&self) -> &B {
77        self.backend
78    }
79
80    /// Borrow the backend mutably for extension execution.
81    pub fn backend_mut(&mut self) -> &mut B {
82        self.backend
83    }
84
85    /// Borrow the extension runtime cache store.
86    pub fn caches(&self) -> &ExtensionCacheStore {
87        self.caches
88    }
89
90    /// Borrow the extension runtime cache store mutably.
91    pub fn caches_mut(&mut self) -> &mut ExtensionCacheStore {
92        self.caches
93    }
94
95    /// Borrow backend and extension cache store as disjoint mutable parts.
96    pub fn parts_mut(&mut self) -> (&mut B, &mut ExtensionCacheStore) {
97        (self.backend, self.caches)
98    }
99}
100
101#[cfg(test)]
102mod tests {
103    use super::*;
104    use tenferro_cpu::CpuEngine;
105    use tenferro_tensor::{Engine, Tensor, TensorElementwise, TensorRead};
106
107    use crate::ExtensionCacheSelector;
108
109    #[test]
110    fn context_accepts_a_caller_held_engine() {
111        let mut engine = CpuEngine::serial();
112        let mut caches = ExtensionCacheStore::new();
113
114        let mut context = ExtensionExecutionContext::new(&mut engine, &mut caches);
115        let _: &dyn Engine = context.backend();
116        let lhs = Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap();
117        let rhs = Tensor::from_vec_col_major(vec![2], vec![3.0_f64, 4.0]).unwrap();
118        let output = context
119            .backend_mut()
120            .add_read(TensorRead::from_tensor(&lhs), TensorRead::from_tensor(&rhs))
121            .unwrap();
122
123        assert_eq!(output.as_slice::<f64>().unwrap(), &[4.0, 6.0]);
124        assert_eq!(
125            context.caches().stats(ExtensionCacheSelector::All).entries,
126            0
127        );
128
129        let (_, caches) = context.parts_mut();
130        assert_eq!(caches.stats(ExtensionCacheSelector::All).entries, 0);
131    }
132}