1use molframe_core::ExecutionContext;
4use std::any::{Any, TypeId, type_name};
5use std::fmt;
6use std::marker::PhantomData;
7use std::sync::Arc;
8use std::sync::atomic::{AtomicU64, Ordering};
9use thiserror::Error;
10
11pub(crate) type Value = Arc<dyn Any + Send + Sync>;
12pub(crate) type Kernel =
13 Arc<dyn Fn(&[Value], &ExecutionContext) -> WorkflowResult<Value> + Send + Sync>;
14
15static NEXT_GRAPH_ID: AtomicU64 = AtomicU64::new(1);
16
17#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
19pub enum Cost {
20 Borrow,
22 Adopt,
24 Decode,
26 Copy,
28 Materialize,
30}
31
32impl fmt::Display for Cost {
33 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
34 formatter.write_str(match self {
35 Self::Borrow => "borrow",
36 Self::Adopt => "adopt",
37 Self::Decode => "decode",
38 Self::Copy => "copy",
39 Self::Materialize => "materialize",
40 })
41 }
42}
43
44#[derive(Clone, Debug, PartialEq, Eq)]
46pub struct OperationMetadata {
47 pub name: &'static str,
49 pub cse_key: Option<Box<str>>,
51 pub cost: Cost,
53 pub estimated_output_bytes: usize,
55 pub uses_spatial_index: bool,
57}
58
59impl OperationMetadata {
60 #[must_use]
62 pub const fn new(name: &'static str, cost: Cost) -> Self {
63 Self {
64 name,
65 cse_key: None,
66 cost,
67 estimated_output_bytes: 0,
68 uses_spatial_index: false,
69 }
70 }
71
72 #[must_use]
74 pub fn with_cse_key(mut self, key: impl Into<Box<str>>) -> Self {
75 self.cse_key = Some(key.into());
76 self
77 }
78
79 #[must_use]
81 pub const fn with_estimated_output_bytes(mut self, bytes: usize) -> Self {
82 self.estimated_output_bytes = bytes;
83 self
84 }
85
86 #[must_use]
88 pub const fn with_spatial_index(mut self) -> Self {
89 self.uses_spatial_index = true;
90 self
91 }
92}
93
94pub struct Node<T> {
96 pub(crate) graph: u64,
97 pub(crate) index: usize,
98 pub(crate) marker: PhantomData<fn() -> T>,
99}
100
101impl<T> Copy for Node<T> {}
102
103impl<T> Clone for Node<T> {
104 fn clone(&self) -> Self {
105 *self
106 }
107}
108
109impl<T> fmt::Debug for Node<T> {
110 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
111 formatter
112 .debug_struct("Node")
113 .field("graph", &self.graph)
114 .field("index", &self.index)
115 .field("type", &type_name::<T>())
116 .finish()
117 }
118}
119
120pub type Input<T> = Node<T>;
122
123#[derive(Clone, Debug)]
125pub struct Output<T> {
126 pub(crate) node: Node<T>,
127 pub(crate) name: Box<str>,
128}
129
130impl<T> Output<T> {
131 #[must_use]
133 pub fn name(&self) -> &str {
134 &self.name
135 }
136}
137
138#[derive(Clone)]
139pub(crate) struct GraphNode {
140 pub(crate) type_id: TypeId,
141 pub(crate) type_name: &'static str,
142 pub(crate) dependencies: Vec<usize>,
143 pub(crate) expected_inputs: Vec<TypeId>,
144 pub(crate) input_name: Option<Box<str>>,
145 pub(crate) metadata: OperationMetadata,
146 pub(crate) kernel: Option<Kernel>,
147}
148
149impl fmt::Debug for GraphNode {
150 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
151 formatter
152 .debug_struct("GraphNode")
153 .field("type_name", &self.type_name)
154 .field("dependencies", &self.dependencies)
155 .field("input_name", &self.input_name)
156 .field("metadata", &self.metadata)
157 .finish_non_exhaustive()
158 }
159}
160
161#[derive(Clone, Debug)]
163pub struct Workflow {
164 pub(crate) id: u64,
165 pub(crate) nodes: Vec<GraphNode>,
166 pub(crate) outputs: Vec<(Box<str>, usize, TypeId)>,
167}
168
169impl Default for Workflow {
170 fn default() -> Self {
171 Self::new()
172 }
173}
174
175impl Workflow {
176 #[must_use]
178 pub fn new() -> Self {
179 Self {
180 id: NEXT_GRAPH_ID.fetch_add(1, Ordering::Relaxed),
181 nodes: Vec::new(),
182 outputs: Vec::new(),
183 }
184 }
185
186 pub fn input<T: Send + Sync + 'static>(
192 &mut self,
193 name: impl Into<Box<str>>,
194 cost: Cost,
195 ) -> Result<Input<T>, WorkflowBuildError> {
196 let name = name.into();
197 if name.is_empty() {
198 return Err(WorkflowBuildError::EmptyName);
199 }
200 if self
201 .nodes
202 .iter()
203 .any(|node| node.input_name.as_deref() == Some(&name))
204 {
205 return Err(WorkflowBuildError::DuplicateInput(name));
206 }
207 let index = self.nodes.len();
208 self.nodes.push(GraphNode {
209 type_id: TypeId::of::<T>(),
210 type_name: type_name::<T>(),
211 dependencies: Vec::new(),
212 expected_inputs: Vec::new(),
213 input_name: Some(name),
214 metadata: OperationMetadata::new("input", cost),
215 kernel: None,
216 });
217 Ok(self.node(index))
218 }
219
220 pub fn map<A, T, F>(
226 &mut self,
227 input: Node<A>,
228 metadata: OperationMetadata,
229 kernel: F,
230 ) -> Result<Node<T>, WorkflowBuildError>
231 where
232 A: Send + Sync + 'static,
233 T: Send + Sync + 'static,
234 F: Fn(&A, &ExecutionContext) -> WorkflowResult<T> + Send + Sync + 'static,
235 {
236 self.ensure_graph(input.graph)?;
237 let erased: Kernel = Arc::new(move |values, context| {
238 let value = downcast::<A>(&values[0])?;
239 kernel(value, context).map(|result| Arc::new(result) as Value)
240 });
241 Ok(self.push_operation::<T>(vec![input.index], vec![TypeId::of::<A>()], metadata, erased))
242 }
243
244 pub fn map2<A, B, T, F>(
250 &mut self,
251 first: Node<A>,
252 second: Node<B>,
253 metadata: OperationMetadata,
254 kernel: F,
255 ) -> Result<Node<T>, WorkflowBuildError>
256 where
257 A: Send + Sync + 'static,
258 B: Send + Sync + 'static,
259 T: Send + Sync + 'static,
260 F: Fn(&A, &B, &ExecutionContext) -> WorkflowResult<T> + Send + Sync + 'static,
261 {
262 self.ensure_graph(first.graph)?;
263 self.ensure_graph(second.graph)?;
264 let erased: Kernel = Arc::new(move |values, context| {
265 let first = downcast::<A>(&values[0])?;
266 let second = downcast::<B>(&values[1])?;
267 kernel(first, second, context).map(|result| Arc::new(result) as Value)
268 });
269 Ok(self.push_operation::<T>(
270 vec![first.index, second.index],
271 vec![TypeId::of::<A>(), TypeId::of::<B>()],
272 metadata,
273 erased,
274 ))
275 }
276
277 pub fn output<T: Send + Sync + 'static>(
283 &mut self,
284 name: impl Into<Box<str>>,
285 node: Node<T>,
286 ) -> Result<Output<T>, WorkflowBuildError> {
287 self.ensure_graph(node.graph)?;
288 let name = name.into();
289 if name.is_empty() {
290 return Err(WorkflowBuildError::EmptyName);
291 }
292 if self
293 .outputs
294 .iter()
295 .any(|(existing, _, _)| existing == &name)
296 {
297 return Err(WorkflowBuildError::DuplicateOutput(name));
298 }
299 self.outputs
300 .push((name.clone(), node.index, TypeId::of::<T>()));
301 Ok(Output { node, name })
302 }
303
304 fn node<T>(&self, index: usize) -> Node<T> {
305 Node {
306 graph: self.id,
307 index,
308 marker: PhantomData,
309 }
310 }
311
312 fn push_operation<T: Send + Sync + 'static>(
313 &mut self,
314 dependencies: Vec<usize>,
315 expected_inputs: Vec<TypeId>,
316 metadata: OperationMetadata,
317 kernel: Kernel,
318 ) -> Node<T> {
319 let index = self.nodes.len();
320 self.nodes.push(GraphNode {
321 type_id: TypeId::of::<T>(),
322 type_name: type_name::<T>(),
323 dependencies,
324 expected_inputs,
325 input_name: None,
326 metadata,
327 kernel: Some(kernel),
328 });
329 self.node(index)
330 }
331
332 fn ensure_graph(&self, graph: u64) -> Result<(), WorkflowBuildError> {
333 if graph == self.id {
334 Ok(())
335 } else {
336 Err(WorkflowBuildError::ForeignNode)
337 }
338 }
339}
340
341fn downcast<T: Send + Sync + 'static>(value: &Value) -> WorkflowResult<&T> {
342 value
343 .downcast_ref::<T>()
344 .ok_or(WorkflowError::InternalTypeMismatch {
345 expected: type_name::<T>(),
346 })
347}
348
349#[derive(Clone, Debug, Error, PartialEq, Eq)]
351pub enum WorkflowBuildError {
352 #[error("workflow names must not be empty")]
354 EmptyName,
355 #[error("duplicate workflow input: {0}")]
357 DuplicateInput(Box<str>),
358 #[error("duplicate workflow output: {0}")]
360 DuplicateOutput(Box<str>),
361 #[error("node belongs to another workflow")]
363 ForeignNode,
364}
365
366#[derive(Debug, Error)]
368pub enum WorkflowError {
369 #[error("workflow has no outputs")]
371 NoOutputs,
372 #[error("node {node} references missing dependency {dependency}")]
374 MissingDependency {
375 node: usize,
377 dependency: usize,
379 },
380 #[error("node {node} has an incompatible dependency type")]
382 TypeMismatch {
383 node: usize,
385 },
386 #[error("workflow graph contains a cycle")]
388 Cycle,
389 #[error("missing workflow input: {0}")]
391 MissingInput(Box<str>),
392 #[error("workflow input {name:?} has the wrong type; expected {expected}")]
394 InputType {
395 name: Box<str>,
397 expected: &'static str,
399 },
400 #[error("workflow execution cancelled")]
402 Cancelled,
403 #[error(
405 "workflow needs an estimated {estimated} bytes but the context permits {budget}; raise the memory budget: Python ExecutionContext(memory_budget=...), Rust ExecutionContext::builder().memory_budget(MemoryBudget::new(..)?), CLI --memory-budget, or env MOLFRAME_MEMORY_BUDGET"
406 )]
407 MemoryRefused {
408 estimated: usize,
410 budget: usize,
412 },
413 #[error("operation {operation:?} failed: {message}")]
415 Operation {
416 operation: &'static str,
418 message: Box<str>,
420 },
421 #[error("internal workflow type mismatch; expected {expected}")]
423 InternalTypeMismatch {
424 expected: &'static str,
426 },
427 #[error("workflow result {name:?} has the wrong requested type")]
429 ResultType {
430 name: Box<str>,
432 },
433 #[error("workflow result is unavailable: {0}")]
435 MissingResult(Box<str>),
436}
437
438pub type WorkflowResult<T> = Result<T, WorkflowError>;
440
441molframe_core::diagnostic_from!(WorkflowBuildError, |_error| molframe_core::Code::E6101);
442
443molframe_core::diagnostic_from!(WorkflowError, |error| match error {
444 WorkflowError::NoOutputs
445 | WorkflowError::MissingDependency { .. }
446 | WorkflowError::TypeMismatch { .. }
447 | WorkflowError::Cycle
448 | WorkflowError::MissingInput(_)
449 | WorkflowError::InputType { .. }
450 | WorkflowError::ResultType { .. }
451 | WorkflowError::MissingResult(_) => molframe_core::Code::E6101,
452 WorkflowError::Cancelled => molframe_core::Code::E1904,
453 WorkflowError::MemoryRefused { .. } => molframe_core::Code::E7001,
454 WorkflowError::Operation { .. } => molframe_core::Code::E5101,
455 WorkflowError::InternalTypeMismatch { .. } => molframe_core::Code::E9001,
456});