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}