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::BackendSession;
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::CpuBackend;
23/// use tenferro_tensor::{BackendSessionHost, Tensor};
24/// use tenferro_runtime::{
25///     ExtensionCacheSelector, ExtensionCacheStore, ExtensionExecutionContext,
26/// };
27///
28/// let mut backend = CpuBackend::new();
29/// let mut caches = ExtensionCacheStore::new();
30/// backend.with_backend_session(|session| {
31///     let mut context = ExtensionExecutionContext::new(session, &mut caches);
32///     let lhs = Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap();
33///     let rhs = Tensor::from_vec_col_major(vec![2], vec![3.0_f64, 4.0]).unwrap();
34///     let output = context.backend_mut().add(&lhs, &rhs).unwrap();
35///
36///     assert_eq!(output.as_slice::<f64>().unwrap(), &[4.0, 6.0]);
37///     assert_eq!(context.caches().stats(ExtensionCacheSelector::All).entries, 0);
38/// });
39/// ```
40///
41/// A session borrow cannot escape the call that supplied it.
42///
43/// ```compile_fail
44/// use tenferro_cpu::{with_cpu_exec_session, CpuBackend, CpuExecSession};
45/// use tenferro_runtime::{ExtensionCacheStore, ExtensionExecutionContext};
46/// use tenferro_tensor::BackendSessionHost;
47///
48/// fn leak_context<'a>(
49///     backend: &'a mut CpuBackend,
50///     caches: &'a mut ExtensionCacheStore,
51/// ) -> ExtensionExecutionContext<'a, CpuExecSession<'a>> {
52///     backend.with_backend_session(move |session| {
53///         with_cpu_exec_session(session, |cpu_session| {
54///             ExtensionExecutionContext::new(cpu_session, caches)
55///         })
56///         .unwrap()
57///     })
58/// }
59/// ```
60pub struct ExtensionExecutionContext<'a, B: BackendSession + ?Sized> {
61    backend: &'a mut B,
62    caches: &'a mut ExtensionCacheStore,
63}
64
65impl<B: BackendSession + ?Sized> fmt::Debug for ExtensionExecutionContext<'_, B> {
66    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
67        f.debug_struct("ExtensionExecutionContext")
68            .field("backend_type", &std::any::type_name::<B>())
69            .field("caches", &self.caches)
70            .finish_non_exhaustive()
71    }
72}
73
74impl<'a, B: BackendSession + ?Sized> ExtensionExecutionContext<'a, B> {
75    /// Build a context from externally-owned backend and cache state.
76    pub fn new(backend: &'a mut B, caches: &'a mut ExtensionCacheStore) -> Self {
77        Self { backend, caches }
78    }
79
80    /// Borrow the backend for non-mutating inspection.
81    pub fn backend(&self) -> &B {
82        self.backend
83    }
84
85    /// Borrow the backend mutably for extension execution.
86    pub fn backend_mut(&mut self) -> &mut B {
87        self.backend
88    }
89
90    /// Borrow the extension runtime cache store.
91    pub fn caches(&self) -> &ExtensionCacheStore {
92        self.caches
93    }
94
95    /// Borrow the extension runtime cache store mutably.
96    pub fn caches_mut(&mut self) -> &mut ExtensionCacheStore {
97        self.caches
98    }
99
100    /// Borrow backend and extension cache store as disjoint mutable parts.
101    pub fn parts_mut(&mut self) -> (&mut B, &mut ExtensionCacheStore) {
102        (self.backend, self.caches)
103    }
104}
105
106#[cfg(test)]
107mod tests {
108    use super::*;
109    use tenferro_cpu::CpuBackend;
110    use tenferro_tensor::{BackendSession, BackendSessionHost, Tensor};
111
112    use crate::ExtensionCacheSelector;
113
114    #[test]
115    fn context_accepts_non_owning_backend_session() {
116        let mut backend = CpuBackend::new();
117        let mut caches = ExtensionCacheStore::new();
118
119        backend.with_backend_session(|session| {
120            let mut context = ExtensionExecutionContext::new(session, &mut caches);
121            let _: &dyn BackendSession = context.backend();
122            let lhs = Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap();
123            let rhs = Tensor::from_vec_col_major(vec![2], vec![3.0_f64, 4.0]).unwrap();
124            let output = context.backend_mut().add(&lhs, &rhs).unwrap();
125
126            assert_eq!(output.as_slice::<f64>().unwrap(), &[4.0, 6.0]);
127            assert_eq!(
128                context.caches().stats(ExtensionCacheSelector::All).entries,
129                0
130            );
131
132            let (_, caches) = context.parts_mut();
133            assert_eq!(caches.stats(ExtensionCacheSelector::All).entries, 0);
134        });
135    }
136}