Skip to main content

otf_pixels_core/
graph.rs

1//! The immutable lazy operation graph.
2//!
3//! Chaining does no pixel work: each call wraps the current node in a new one,
4//! producing a DAG of [`Arc<Node>`] (ARCHITECTURE §Layer 3). Sharing is free —
5//! cloning an [`Image`] clones an `Arc`, and two pipelines branching from a
6//! common prefix share those nodes, so the evaluator computes them once.
7//!
8//! Descriptors are resolved **as the graph is built**. By the time a node
9//! exists, its output shape is already known, so [`Image::metadata`] is a field
10//! read rather than a traversal.
11
12use crate::{Format, ImageDescriptor, Metadata, Op, PixelsError, Producer, Region, Result};
13use std::sync::Arc;
14use std::sync::atomic::{AtomicU64, Ordering};
15
16/// A unique identifier for a graph node.
17///
18/// Identity is per-node, not per-op: two `crop` nodes with identical parameters
19/// are distinct. The M1 evaluator memoizes on this so a shared subgraph
20/// evaluates once, and M2's tile cache keys on `(NodeId, Region)`.
21#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
22pub struct NodeId(u64);
23
24impl NodeId {
25    /// The identifier's raw value, for diagnostics and cache keys.
26    #[must_use]
27    pub const fn get(self) -> u64 {
28        self.0
29    }
30
31    /// Allocate the next process-unique identifier.
32    fn next() -> Self {
33        static COUNTER: AtomicU64 = AtomicU64::new(0);
34        Self(COUNTER.fetch_add(1, Ordering::Relaxed))
35    }
36}
37
38/// What a node does: originate pixels, or transform its inputs.
39#[derive(Debug, Clone)]
40enum NodeKind {
41    /// A root node that produces pixels from a decoder or buffer.
42    Source(Arc<dyn Producer>),
43    /// An interior node that transforms its inputs.
44    Op(Arc<dyn Op>),
45}
46
47/// A node in the immutable operation graph.
48///
49/// Nodes are never mutated after construction. Build them through [`Image`]
50/// rather than directly.
51#[derive(Debug)]
52pub struct Node {
53    id: NodeId,
54    kind: NodeKind,
55    inputs: Vec<Arc<Node>>,
56    descriptor: ImageDescriptor,
57}
58
59impl Node {
60    /// This node's unique identifier.
61    #[must_use]
62    pub const fn id(&self) -> NodeId {
63        self.id
64    }
65
66    /// The shape of this node's output, resolved at build time.
67    #[must_use]
68    pub const fn descriptor(&self) -> ImageDescriptor {
69        self.descriptor
70    }
71
72    /// This node's inputs, in the order the op consumes them.
73    #[must_use]
74    pub fn inputs(&self) -> &[Arc<Node>] {
75        &self.inputs
76    }
77
78    /// The op this node applies, or [`None`] if it is a source.
79    #[must_use]
80    pub fn op(&self) -> Option<&Arc<dyn Op>> {
81        match &self.kind {
82            NodeKind::Op(op) => Some(op),
83            NodeKind::Source(_) => None,
84        }
85    }
86
87    /// The producer at this node, or [`None`] if it is an interior op.
88    #[must_use]
89    pub fn producer(&self) -> Option<&Arc<dyn Producer>> {
90        match &self.kind {
91            NodeKind::Source(producer) => Some(producer),
92            NodeKind::Op(_) => None,
93        }
94    }
95
96    /// A short name for this node, for diagnostics.
97    #[must_use]
98    pub fn name(&self) -> &'static str {
99        match &self.kind {
100            NodeKind::Source(producer) => producer.name(),
101            NodeKind::Op(op) => op.name(),
102        }
103    }
104
105    /// The number of nodes reachable from here, counting shared nodes once.
106    ///
107    /// Useful for asserting in tests that chaining built the graph it should
108    /// have, and that a shared prefix really is shared.
109    #[must_use]
110    pub fn node_count(self: &Arc<Self>) -> usize {
111        let mut seen = std::collections::HashSet::new();
112        let mut stack = vec![Arc::clone(self)];
113        while let Some(node) = stack.pop() {
114            if !seen.insert(node.id) {
115                continue;
116            }
117            stack.extend(node.inputs.iter().map(Arc::clone));
118        }
119        seen.len()
120    }
121}
122
123impl Drop for Node {
124    /// Tear the graph down iteratively.
125    ///
126    /// The derived drop glue would recurse once per node: a node drops its
127    /// `inputs`, each of which drops *its* inputs, and so on. A long chain —
128    /// which callers build simply by chaining many ops — would then overflow
129    /// the stack at teardown, aborting the process. Since an abort is exactly
130    /// what the failure model forbids, teardown uses an explicit worklist.
131    ///
132    /// Only uniquely-owned inputs are unwrapped: a node still shared by
133    /// another branch of the DAG is left for its last owner to drop.
134    fn drop(&mut self) {
135        let mut stack: Vec<Arc<Self>> = std::mem::take(&mut self.inputs);
136        while let Some(node) = stack.pop() {
137            if let Some(mut node) = Arc::into_inner(node) {
138                // Move the grandchildren onto our worklist before `node` drops,
139                // so its own `drop` finds nothing left to recurse into.
140                stack.append(&mut node.inputs);
141            }
142        }
143    }
144}
145
146/// A lazily evaluated image: a handle onto one node of an op graph.
147///
148/// Constructing an `Image` and chaining ops onto it performs **no** pixel work
149/// and reads no source bytes beyond the header (SPEC §Guarantees 3). Pixels
150/// move only when a terminal pulls them.
151///
152/// `Image` is cheap to clone — it shares graph nodes rather than copying them —
153/// and is `Send + Sync`.
154#[derive(Debug, Clone)]
155pub struct Image {
156    node: Arc<Node>,
157    format: Format,
158}
159
160impl Image {
161    /// Build an image rooted at `producer`.
162    ///
163    /// `format` is the container the pixels came from, reported by
164    /// [`Image::metadata`]. Use [`Format::Raw`] for pixels the caller supplied
165    /// directly.
166    #[must_use]
167    pub fn from_producer(producer: Arc<dyn Producer>, format: Format) -> Self {
168        let descriptor = producer.descriptor();
169        Self {
170            node: Arc::new(Node {
171                id: NodeId::next(),
172                kind: NodeKind::Source(producer),
173                inputs: Vec::new(),
174                descriptor,
175            }),
176            format,
177        }
178    }
179
180    /// The graph node this handle points at.
181    #[must_use]
182    pub fn node(&self) -> &Arc<Node> {
183        &self.node
184    }
185
186    /// The shape of this image, resolved when the node was built.
187    #[must_use]
188    pub fn descriptor(&self) -> ImageDescriptor {
189        self.node.descriptor
190    }
191
192    /// Header-only facts about this image.
193    ///
194    /// Free: descriptors flowed forward at graph-build time, so this reads a
195    /// field and decodes nothing.
196    ///
197    /// # Errors
198    ///
199    /// Infallible for graphs built through this API. It returns [`Result`] so
200    /// that formats whose headers are parsed lazily can report a malformed
201    /// header here without a breaking signature change.
202    pub fn metadata(&self) -> Result<Metadata> {
203        Ok(Metadata::new(&self.node.descriptor, self.format))
204    }
205
206    /// Chain a single-input op onto this image.
207    ///
208    /// The op's output descriptor is computed now, so an op that cannot apply
209    /// to this input fails here rather than at evaluation time.
210    ///
211    /// # Errors
212    ///
213    /// Propagates [`Op::output_descriptor`], and returns
214    /// [`PixelsError::Graph`] if `op` does not take exactly one input.
215    pub fn apply(&self, op: Arc<dyn Op>) -> Result<Self> {
216        Self::combine(std::slice::from_ref(self), op)
217    }
218
219    /// Chain a multi-input op over `inputs`.
220    ///
221    /// This is how ops like `composite` join two branches of a graph. The
222    /// container format reported by [`Image::metadata`] is taken from the first
223    /// input.
224    ///
225    /// # Errors
226    ///
227    /// Returns [`PixelsError::Graph`] if `inputs` is empty or its length does
228    /// not match [`Op::arity`], and propagates [`Op::output_descriptor`].
229    pub fn combine(inputs: &[Self], op: Arc<dyn Op>) -> Result<Self> {
230        let Some(first) = inputs.first() else {
231            return Err(PixelsError::graph(format!(
232                "op `{}` needs at least one input",
233                op.name()
234            )));
235        };
236        if inputs.len() != op.arity() {
237            return Err(PixelsError::graph(format!(
238                "op `{}` takes {} input(s), got {}",
239                op.name(),
240                op.arity(),
241                inputs.len()
242            )));
243        }
244        let descriptors: Vec<ImageDescriptor> = inputs.iter().map(Self::descriptor).collect();
245        let descriptor = op.output_descriptor(&descriptors)?;
246        let format = first.format;
247        Ok(Self {
248            node: Arc::new(Node {
249                id: NodeId::next(),
250                kind: NodeKind::Op(op),
251                inputs: inputs.iter().map(|image| Arc::clone(&image.node)).collect(),
252                descriptor,
253            }),
254            format,
255        })
256    }
257
258    /// The region covering this whole image.
259    #[must_use]
260    pub fn region(&self) -> Region {
261        self.node.descriptor.region()
262    }
263}
264
265#[cfg(test)]
266#[allow(
267    clippy::unwrap_used,
268    clippy::indexing_slicing,
269    reason = "tests operate on known-good values and assert shapes directly"
270)]
271mod tests {
272    use super::*;
273    use crate::testing::{ConstantOp, CountingProducer};
274    use crate::{AccessPattern, Op, PixelFormat, Tile, TileMut};
275
276    fn source(width: u32, height: u32) -> Image {
277        let producer =
278            CountingProducer::new(ImageDescriptor::new(width, height, PixelFormat::Gray8).unwrap());
279        Image::from_producer(Arc::new(producer), Format::Raw)
280    }
281
282    #[test]
283    fn chaining_builds_a_dag_without_touching_pixels() {
284        let producer = Arc::new(CountingProducer::new(
285            ImageDescriptor::new(4, 4, PixelFormat::Gray8).unwrap(),
286        ));
287        let image = Image::from_producer(Arc::clone(&producer) as Arc<dyn Producer>, Format::Raw);
288        let chained = image.apply(Arc::new(ConstantOp::new(7))).unwrap();
289        let _ = chained.apply(Arc::new(ConstantOp::new(9))).unwrap();
290        assert_eq!(
291            producer.produce_calls(),
292            0,
293            "graph construction must not pull pixels"
294        );
295    }
296
297    #[test]
298    fn metadata_is_available_without_evaluation() {
299        let producer = Arc::new(CountingProducer::new(
300            ImageDescriptor::new(6, 3, PixelFormat::Gray8).unwrap(),
301        ));
302        let image = Image::from_producer(Arc::clone(&producer) as Arc<dyn Producer>, Format::Raw);
303        let meta = image.metadata().unwrap();
304        assert_eq!((meta.width, meta.height), (6, 3));
305        assert_eq!(meta.format, Format::Raw);
306        assert_eq!(meta.pixel, PixelFormat::Gray8);
307        assert_eq!(producer.produce_calls(), 0);
308    }
309
310    #[test]
311    fn descriptors_flow_forward_through_the_chain() {
312        /// An op that halves its input's width, to prove shapes propagate.
313        #[derive(Debug)]
314        struct Halve;
315        impl Op for Halve {
316            fn name(&self) -> &'static str {
317                "halve"
318            }
319            fn output_descriptor(&self, inputs: &[ImageDescriptor]) -> Result<ImageDescriptor> {
320                let input = inputs
321                    .first()
322                    .ok_or_else(|| PixelsError::graph("no input"))?;
323                input.resized(input.width / 2, input.height)
324            }
325            fn input_regions(&self, out: Region, _: &[ImageDescriptor]) -> Result<Vec<Region>> {
326                Ok(vec![out])
327            }
328            fn access_pattern(&self) -> AccessPattern {
329                AccessPattern::Sequential
330            }
331            fn compute(&self, _: &[Tile<'_>], _: &mut TileMut<'_>) -> Result<()> {
332                Ok(())
333            }
334        }
335        let image = source(16, 4)
336            .apply(Arc::new(Halve))
337            .unwrap()
338            .apply(Arc::new(Halve))
339            .unwrap();
340        assert_eq!(image.descriptor().width, 4);
341        assert_eq!(image.descriptor().height, 4);
342    }
343
344    #[test]
345    fn an_op_rejecting_its_input_fails_at_build_time() {
346        /// An op that refuses every input, to prove build-time validation.
347        #[derive(Debug)]
348        struct Refuses;
349        impl Op for Refuses {
350            fn name(&self) -> &'static str {
351                "refuses"
352            }
353            fn output_descriptor(&self, _: &[ImageDescriptor]) -> Result<ImageDescriptor> {
354                Err(PixelsError::unsupported("never applicable"))
355            }
356            fn input_regions(&self, out: Region, _: &[ImageDescriptor]) -> Result<Vec<Region>> {
357                Ok(vec![out])
358            }
359            fn compute(&self, _: &[Tile<'_>], _: &mut TileMut<'_>) -> Result<()> {
360                Ok(())
361            }
362        }
363        let err = source(4, 4).apply(Arc::new(Refuses)).unwrap_err();
364        assert_eq!(err.code(), crate::ErrorCode::Unsupported);
365    }
366
367    #[test]
368    fn branches_share_their_common_prefix() {
369        let base = source(4, 4).apply(Arc::new(ConstantOp::new(1))).unwrap();
370        let left = base.apply(Arc::new(ConstantOp::new(2))).unwrap();
371        let right = base.apply(Arc::new(ConstantOp::new(3))).unwrap();
372        // Both branches point at the same prefix node, not a copy of it.
373        assert_eq!(left.node().inputs()[0].id(), right.node().inputs()[0].id());
374        // source + base + left = 3 distinct nodes on the left branch.
375        assert_eq!(left.node().node_count(), 3);
376    }
377
378    #[test]
379    fn cloning_an_image_shares_the_node() {
380        let image = source(4, 4);
381        let clone = image.clone();
382        assert_eq!(image.node().id(), clone.node().id());
383        assert!(Arc::ptr_eq(image.node(), clone.node()));
384    }
385
386    #[test]
387    fn node_ids_are_unique() {
388        let a = source(2, 2);
389        let b = source(2, 2);
390        assert_ne!(a.node().id(), b.node().id());
391        assert_ne!(a.node().id().get(), b.node().id().get());
392    }
393
394    #[test]
395    fn arity_mismatch_is_a_graph_error() {
396        let image = source(4, 4);
397        let err =
398            Image::combine(&[image.clone(), image], Arc::new(ConstantOp::new(1))).unwrap_err();
399        assert_eq!(err.code(), crate::ErrorCode::Graph);
400        let err = Image::combine(&[], Arc::new(ConstantOp::new(1))).unwrap_err();
401        assert_eq!(err.code(), crate::ErrorCode::Graph);
402    }
403
404    #[test]
405    fn images_are_send_and_sync() {
406        const fn assert_send_sync<T: Send + Sync>() {}
407        assert_send_sync::<Image>();
408        assert_send_sync::<Arc<Node>>();
409    }
410}