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