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}