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}