Skip to main content

polydat_core/iteration/comprehension/ir/
op.rs

1// Copyright 2024-2026 Jonathan Shook
2// SPDX-License-Identifier: Apache-2.0
3
4//! Operator IR — comprehension_forms.md §9.1.
5//!
6//! Every well-formed comprehension AST compiles to a finite
7//! sequence of these 8 opcodes. Every operator is a stream
8//! transducer; operands flow as tuple streams via
9//! `advance() -> Option<Tuple>`, never as materialized
10//! `Vec<Tuple>`. The two materialization barriers — non-Lex
11//! `ORDER_MATERIALIZE` and `ZIP(Cycle)`'s buffering of operands
12//! that are not index-addressable — are called out explicitly.
13
14use serde::{Deserialize, Serialize};
15
16use crate::iteration::comprehension::ast::Comprehension;
17use crate::iteration::comprehension::metadata::{CycleOperand, cycle_plan_is_empty};
18use crate::iteration::comprehension::source::Source;
19use crate::iteration::comprehension::strategy::{StrategyName, ZipMode};
20
21/// The 8-opcode IR set (comprehension_forms.md §9.1).
22#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
23#[serde(tag = "op", rename_all = "snake_case")]
24pub enum Op {
25    /// Push a single-name tuple stream produced by `source`.
26    /// Streaming; O(1) per pull above the source's own state.
27    PushClause {
28        /// The name the stream binds.
29        name: String,
30        /// Where its values come from.
31        source: Source,
32    },
33
34    /// Replace the top-N stream operands with one stream that
35    /// enumerates their cross product in Lex order. Streaming.
36    Cartesian {
37        /// Operands combined.
38        n: usize,
39    },
40
41    /// Replace the top-N stream operands with their lockstep
42    /// diagonal. Streaming under Strict/Truncate; under `Cycle`
43    /// each operand is held as `operands` plans
44    /// (comprehension_forms.md §3.3, §6.2):
45    /// indexed, streamed, or buffered.
46    Zip {
47        /// Operands combined.
48        n: usize,
49        /// The length policy.
50        mode: ZipMode,
51        /// Under `Cycle`, how each operand is held, from the
52        /// operands' metadata ([`crate::iteration::comprehension::metadata::cycle_operands`]).
53        /// Empty under Strict and Truncate, and for a program built
54        /// without one, whose operands the interpreter plans from the
55        /// streams themselves.
56        #[serde(default, skip_serializing_if = "Vec::is_empty")]
57        operands: Vec<CycleOperand>,
58    },
59
60    /// Replace the top-N stream operands with a stream that
61    /// concatenates them in operand order. Streaming.
62    Union {
63        /// Operands concatenated.
64        n: usize,
65    },
66
67    /// Wrap the top operand with a per-tuple predicate check.
68    /// Streaming.
69    Filter {
70        /// The predicate, a boolean expression over the tuple.
71        predicate: String,
72    },
73
74    /// Wrap the top operand with a counter / pass-through.
75    /// Used for `order(Lex, _)` per comprehension_forms.md §10.2 R1; this is
76    /// the "streaming" order opcode.
77    OrderStreaming {
78        /// The streaming order's kind.
79        kind: OrderStreamingKind,
80        /// The output cap, if any.
81        truncation: Option<u64>,
82    },
83
84    /// MATERIALIZATION BARRIER. Push a stream of the tuples a
85    /// non-`Lex` order selects from `input`. Truncation is the output
86    /// cap.
87    ///
88    /// The order is evaluated as a traversal evaluates it
89    /// ([`crate::iteration::comprehension::runtime::evaluate_indexed`]),
90    /// on the first pull, in the empty scope: the strategy selects
91    /// positions from the input's evaluated shape and length, and
92    /// the stream computes each selected tuple as it emits it. Per
93    /// comprehension_forms.md §10.2 R2, over an index-addressable
94    /// input the working set is the selection, O(output); over a
95    /// filter it is the filter's input and the positions of the
96    /// survivors the strategy keeps (§5 V5); over any other input it
97    /// is the input's tuples. An order over a continuous axis samples
98    /// the input's space and holds the samples. No IR is emitted for
99    /// `input`: the op pops nothing and pushes one stream.
100    ///
101    /// `input_index_fn` is the input's addressing scheme as the
102    /// metadata propagator claims it at compile time (§10.7.6),
103    /// which the bounds checker reads; `None` when it claims none.
104    OrderMaterialize {
105        /// The strategy applied.
106        strategy: StrategyName,
107        /// The output cap, if any.
108        truncation: Option<u64>,
109        /// The authored seed a seeded strategy (`Shuffle`, `Lhs`)
110        /// derives its state from; its fixed default when `None`.
111        seed: Option<u64>,
112        /// The input's IndexFn from its metadata at compile time.
113        input_index_fn: Option<crate::iteration::comprehension::metadata::IndexFn>,
114        /// The comprehension ordered.
115        input: Box<Comprehension>,
116    },
117
118    /// Bind the top stream as the comprehension's result.
119    /// Must be the last opcode in a well-formed Program.
120    Dispense,
121}
122
123/// Variant marker for [`Op::OrderStreaming`]. Lex is the one
124/// streaming strategy (comprehension_forms.md §6.2's footprint
125/// table); the kind is an enum so the opcode set stays fixed
126/// whatever strategies stream.
127#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
128#[serde(rename_all = "snake_case")]
129pub enum OrderStreamingKind {
130    /// Lexicographic order: the natural enumeration, counted.
131    Lex,
132}
133
134impl Op {
135    /// Arity for stack-effect computation: how many stream
136    /// operands this opcode pops, and how many it pushes.
137    /// Always pushes 1 for stream-producing ops; `Dispense`
138    /// pushes 0 (it consumes the final stream).
139    pub fn stack_effect(&self) -> (usize, usize) {
140        match self {
141            Op::PushClause { .. } => (0, 1),
142            Op::Cartesian { n } => (*n, 1),
143            Op::Zip { n, .. } => (*n, 1),
144            Op::Union { n } => (*n, 1),
145            Op::Filter { .. } => (1, 1),
146            Op::OrderStreaming { .. } => (1, 1),
147            Op::OrderMaterialize { .. } => (0, 1),
148            Op::Dispense => (1, 0),
149        }
150    }
151
152    /// `true` if this opcode is a materialization barrier per
153    /// comprehension_forms.md §6.2 and §6.3. Used by the bounds checker. A `Cycle` zip
154    /// is one when it buffers an operand, or when it carries no plan
155    /// and may have to; one with an operand known empty holds nothing.
156    pub fn is_barrier(&self) -> bool {
157        match self {
158            Op::OrderMaterialize { .. } => true,
159            Op::Zip {
160                mode: ZipMode::Cycle,
161                operands,
162                ..
163            } => {
164                operands.is_empty()
165                    || (!cycle_plan_is_empty(operands)
166                        && operands
167                            .iter()
168                            .any(|o| matches!(o, CycleOperand::Buffered { .. })))
169            }
170            _ => false,
171        }
172    }
173}
174
175#[cfg(test)]
176mod tests {
177    use super::*;
178
179    #[test]
180    fn stack_effect_basics() {
181        assert_eq!(
182            Op::PushClause {
183                name: "k".into(),
184                source: Source::Literal { values: vec![] },
185            }
186            .stack_effect(),
187            (0, 1)
188        );
189        assert_eq!(Op::Cartesian { n: 3 }.stack_effect(), (3, 1));
190        assert_eq!(Op::Dispense.stack_effect(), (1, 0));
191    }
192
193    fn input() -> Box<Comprehension> {
194        Box::new(Comprehension::clause(
195            "k",
196            Source::IntRange {
197                lo: 0,
198                hi: 50,
199                step: 1,
200            },
201        ))
202    }
203
204    #[test]
205    fn barrier_classification() {
206        assert!(
207            Op::OrderMaterialize {
208                strategy: StrategyName::Halton,
209                truncation: Some(10),
210                seed: None,
211                input_index_fn: None,
212                input: input(),
213            }
214            .is_barrier()
215        );
216        assert!(
217            Op::Zip {
218                n: 2,
219                mode: ZipMode::Cycle,
220                operands: vec![
221                    CycleOperand::Streamed,
222                    CycleOperand::Buffered { bound: Some(3) }
223                ],
224            }
225            .is_barrier()
226        );
227        assert!(
228            !Op::Zip {
229                n: 2,
230                mode: ZipMode::Cycle,
231                operands: vec![CycleOperand::Indexed, CycleOperand::Streamed],
232            }
233            .is_barrier(),
234            "indexing and streaming buffer nothing"
235        );
236        assert!(
237            !Op::Zip {
238                n: 2,
239                mode: ZipMode::Strict,
240                operands: Vec::new(),
241            }
242            .is_barrier()
243        );
244        assert!(
245            !Op::OrderStreaming {
246                kind: OrderStreamingKind::Lex,
247                truncation: None,
248            }
249            .is_barrier()
250        );
251    }
252
253    #[test]
254    fn serde_round_trip() {
255        let op = Op::OrderMaterialize {
256            strategy: StrategyName::Halton,
257            truncation: Some(50),
258            seed: None,
259            input_index_fn: Some(
260                crate::iteration::comprehension::metadata::IndexFn::Lattice {
261                    axis_sizes: vec![10, 5],
262                },
263            ),
264            input: input(),
265        };
266        let json = serde_json::to_string(&op).unwrap();
267        let back: Op = serde_json::from_str(&json).unwrap();
268        assert_eq!(op, back);
269    }
270}