Skip to main content

molframe_engine/
graph.rs

1//! Graph construction and the public typed handles.
2
3use 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/// The data-movement cost attached to an input or operation.
18#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
19pub enum Cost {
20    /// Reuses storage owned elsewhere for the duration of execution.
21    Borrow,
22    /// Retains compatible owner-backed storage without copying it.
23    Adopt,
24    /// Converts an encoded representation into native storage.
25    Decode,
26    /// Duplicates an existing native representation.
27    Copy,
28    /// Gathers or computes a new representation on explicit request.
29    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/// Static information used by compilation, explanation and memory admission.
45#[derive(Clone, Debug, PartialEq, Eq)]
46pub struct OperationMetadata {
47    /// Canonical operation name shown by `explain`.
48    pub name: &'static str,
49    /// Stable semantic key. Equal keys and equal dependencies are CSE candidates.
50    pub cse_key: Option<Box<str>>,
51    /// Dominant data-movement cost.
52    pub cost: Cost,
53    /// Conservative bytes retained by the operation result.
54    pub estimated_output_bytes: usize,
55    /// Whether the operation needs a reusable spatial index.
56    pub uses_spatial_index: bool,
57}
58
59impl OperationMetadata {
60    /// Metadata for an operation without a semantic CSE key.
61    #[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    /// Adds a stable key used for common-subexpression elimination.
73    #[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    /// Adds a conservative retained-size estimate.
80    #[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    /// Marks an operation as a spatial-index consumer.
87    #[must_use]
88    pub const fn with_spatial_index(mut self) -> Self {
89        self.uses_spatial_index = true;
90        self
91    }
92}
93
94/// A typed immutable expression node.
95pub 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
120/// A typed named workflow input.
121pub type Input<T> = Node<T>;
122
123/// A typed named workflow output handle.
124#[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    /// Output name used in named result lookup.
132    #[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/// A symbolic computation graph under construction.
162#[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    /// Starts an empty workflow.
177    #[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    /// Declares a typed named input.
187    ///
188    /// # Errors
189    ///
190    /// Returns an error when the name is empty or already used by an input.
191    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    /// Adds one typed unary operation.
221    ///
222    /// # Errors
223    ///
224    /// Returns an error when `input` belongs to another workflow.
225    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    /// Adds one typed binary operation.
245    ///
246    /// # Errors
247    ///
248    /// Returns an error when either input belongs to another workflow.
249    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    /// Publishes a typed node under a stable output name.
278    ///
279    /// # Errors
280    ///
281    /// Returns an error for an empty/duplicate name or a foreign node.
282    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/// Graph-construction errors detected before compilation.
350#[derive(Clone, Debug, Error, PartialEq, Eq)]
351pub enum WorkflowBuildError {
352    /// Names must be non-empty.
353    #[error("workflow names must not be empty")]
354    EmptyName,
355    /// An input name was declared twice.
356    #[error("duplicate workflow input: {0}")]
357    DuplicateInput(Box<str>),
358    /// An output name was declared twice.
359    #[error("duplicate workflow output: {0}")]
360    DuplicateOutput(Box<str>),
361    /// A node from another workflow was supplied.
362    #[error("node belongs to another workflow")]
363    ForeignNode,
364}
365
366/// Workflow compilation or execution failure.
367#[derive(Debug, Error)]
368pub enum WorkflowError {
369    /// The graph has no named outputs.
370    #[error("workflow has no outputs")]
371    NoOutputs,
372    /// A dependency index is invalid.
373    #[error("node {node} references missing dependency {dependency}")]
374    MissingDependency {
375        /// Node being validated.
376        node: usize,
377        /// Missing dependency index.
378        dependency: usize,
379    },
380    /// A dependency type does not match the operation signature.
381    #[error("node {node} has an incompatible dependency type")]
382    TypeMismatch {
383        /// Node being validated.
384        node: usize,
385    },
386    /// A cycle prevents stable topological execution.
387    #[error("workflow graph contains a cycle")]
388    Cycle,
389    /// A required named input was not supplied.
390    #[error("missing workflow input: {0}")]
391    MissingInput(Box<str>),
392    /// A named input has the wrong concrete type.
393    #[error("workflow input {name:?} has the wrong type; expected {expected}")]
394    InputType {
395        /// Input name.
396        name: Box<str>,
397        /// Expected Rust type.
398        expected: &'static str,
399    },
400    /// Execution was cancelled between nodes.
401    #[error("workflow execution cancelled")]
402    Cancelled,
403    /// Static output-lifetime planning exceeds the execution budget.
404    #[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        /// Conservative workflow estimate.
409        estimated: usize,
410        /// Context memory ceiling.
411        budget: usize,
412    },
413    /// A kernel rejected its input or could not produce its result.
414    #[error("operation {operation:?} failed: {message}")]
415    Operation {
416        /// Operation name.
417        operation: &'static str,
418        /// Kernel error.
419        message: Box<str>,
420    },
421    /// Internal type validation disagreed with the typed builder.
422    #[error("internal workflow type mismatch; expected {expected}")]
423    InternalTypeMismatch {
424        /// Expected concrete Rust type.
425        expected: &'static str,
426    },
427    /// A named result was requested with the wrong type.
428    #[error("workflow result {name:?} has the wrong requested type")]
429    ResultType {
430        /// Result name.
431        name: Box<str>,
432    },
433    /// A result handle does not belong to this execution.
434    #[error("workflow result is unavailable: {0}")]
435    MissingResult(Box<str>),
436}
437
438/// Result type used by workflow kernels.
439pub 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});