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}