polydat_core/dsl/
traversal.rs

1// Copyright 2024-2026 Jonathan Shook
2// SPDX-License-Identifier: Apache-2.0
3
4//! Compile-time support for the `for` construct (SRD 113 steps 2 and
5//! 3): element typing, body-to-child-program lowering, and the
6//! metadata a parent program carries for each traversal and producer.
7//!
8//! A `for` body compiles exactly once, at parent compile time, into a
9//! child [`PolydatProgram`] keyed by the statement's lexical position.
10//! Its element names become `IterationExtern` inputs typed from the
11//! comprehension's sources; outer wires it references become cascade
12//! externs typed from the parent's manifest. Activation (step 4) only
13//! ever allocates state over that program.
14
15use std::collections::{BTreeSet, HashMap};
16use std::sync::{Arc, Mutex};
17
18use crate::ast::PortType;
19use crate::iteration::comprehension::source::{LiteralValue, Source};
20use crate::iteration::comprehension::{Comprehension, StreamerValue};
21use crate::kernel::PolydatProgram;
22
23use super::ast::{
24    Arg, Binding, BindingModifier, CallExpr, Expr, ExternPort, ForSource, ForSourceKind, ForStmt,
25    InputDecl, PolydatFile, Statement,
26};
27use super::lexer::Span;
28
29/// A compiled traversal: one `for` statement and the child program its
30/// body lowered to.
31#[derive(Debug, Clone)]
32pub struct Traversal {
33    /// Lexical position of the `for` statement in its parent.
34    pub span: Span,
35    /// The text after `for`, as written.
36    pub source_text: String,
37    /// The comprehension traversed, with a producer reference already
38    /// resolved to the producer's comprehension.
39    pub comprehension: Comprehension,
40    /// Element names and their compile-time types, in tuple order.
41    pub elements: Vec<(String, PortType)>,
42    /// Outer wires the body references, with the parent's types. Bound
43    /// from the parent at activation.
44    pub cascade: Vec<(String, PortType)>,
45    /// The body's program on the interpreter. Compiled once; every
46    /// interpreter activation shares it.
47    pub program: Arc<PolydatProgram>,
48    /// The body as the parent compiled it, for activations on the other
49    /// engines (engine parity, step 8): compiled once per engine, on the
50    /// first activation that asks.
51    pub body: Arc<BodySource>,
52}
53
54/// A traversal body as its parent compiled it: the child file and the
55/// compiler settings the parent used, so the same body compiles on any
56/// engine, once, keyed by the engine as the interpreter's program is
57/// keyed by the body's position (SRD 113 §5.1).
58pub struct BodySource {
59    pub(crate) file: PolydatFile,
60    pub(crate) source_text: String,
61    pub(crate) source_dir: Option<std::path::PathBuf>,
62    pub(crate) lib_paths: Vec<std::path::PathBuf>,
63    pub(crate) strict: bool,
64    pub(crate) context_label: String,
65    pub(crate) cursor_limit: Option<u64>,
66    pub(crate) pragmas: super::pragmas::PragmaSet,
67    /// The modules the parent program had resolved when the body was
68    /// lowered, its own definitions included, so the body sees them
69    /// wherever it compiles, as the parent did.
70    pub(super) modules: HashMap<String, super::modules::ResolvedModule>,
71    /// The body's program per engine, built on first use.
72    pub(crate) programs: Mutex<HashMap<crate::Engine, Arc<dyn crate::kernel::KernelProgram>>>,
73}
74
75impl BodySource {
76    /// The body's source as the parent lowered it: the implicit `cycle`
77    /// input, one extern per element and per cascaded wire, then the
78    /// body's statements.
79    pub fn source_text(&self) -> &str {
80        &self.source_text
81    }
82}
83
84impl std::fmt::Debug for BodySource {
85    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
86        f.debug_struct("BodySource")
87            .field("context", &self.context_label)
88            .field("statements", &self.file.statements.len())
89            .finish()
90    }
91}
92
93impl Traversal {
94    /// The body's program on `engine`, compiled on the first call for
95    /// that engine and shared by every activation after it, as the
96    /// interpreter's program is (SRD 113 §5.1, §5.2), its own `for`
97    /// statements included: an activation on any engine opens them.
98    pub fn program_on(
99        &self,
100        engine: crate::Engine,
101    ) -> Result<Arc<dyn crate::kernel::KernelProgram>, crate::KernelError> {
102        if matches!(engine, crate::Engine::Interpreter(_)) {
103            return Ok(self.program.clone());
104        }
105        let mut programs = self
106            .body
107            .programs
108            .lock()
109            .unwrap_or_else(|poisoned| poisoned.into_inner());
110        if let Some(program) = programs.get(&engine) {
111            return Ok(program.clone());
112        }
113        let program = super::compile::Compiler::compile_body_on(&self.body, engine)?.into_program();
114        programs.insert(engine, program.clone());
115        Ok(program)
116    }
117}
118
119/// A producer binding, `name := for ...`, recorded on the program that
120/// declares it so traversals over `name` resolve at compile time.
121#[derive(Debug, Clone)]
122pub struct Producer {
123    /// The wire the producer binds.
124    pub name: String,
125    /// Where the binding appears.
126    pub span: Span,
127    /// The text after `for`, as written.
128    pub source_text: String,
129    /// The comprehension, with derivations resolved.
130    pub comprehension: Comprehension,
131}
132
133/// Split a parsed file into the statements the parent compiles directly,
134/// the `for` statements to lower into child programs, and the producer
135/// bindings to record. Order within each group is preserved.
136///
137/// Producer bindings stay in the parent as `const name := streamer(...)`
138/// calls carrying the resolved comprehension, so the wire exists with a
139/// `Streamer` value on it (SRD 113 §3.1). Derivations resolve against
140/// producers bound earlier in the file, in document order.
141pub fn strip_for_forms(
142    file: &PolydatFile,
143) -> Result<(PolydatFile, Vec<ForStmt>, Vec<Producer>), String> {
144    let mut parent = Vec::with_capacity(file.statements.len());
145    let mut fors = Vec::new();
146    let mut producers: Vec<Producer> = Vec::new();
147    for stmt in &file.statements {
148        match stmt {
149            Statement::For(f) => fors.push(f.clone()),
150            Statement::Binding(b) if matches!(b.value, Expr::For(_)) => {
151                let Expr::For(source) = &b.value else {
152                    unreachable!()
153                };
154                let comprehension = resolve_source(source, &producers)?;
155                let name = b.targets.join(",");
156                let value = StreamerValue::new(source.text.clone(), comprehension.clone());
157                parent.push(Statement::Binding(Binding {
158                    targets: b.targets.clone(),
159                    value: Expr::Call(CallExpr {
160                        func: "streamer".into(),
161                        args: vec![Arg::Positional(Expr::StringLit(value.to_json(), b.span))],
162                        span: b.span,
163                    }),
164                    modifier: BindingModifier::CONST,
165                    type_annotation: None,
166                    span: b.span,
167                }));
168                producers.push(Producer {
169                    name,
170                    span: b.span,
171                    source_text: source.text.clone(),
172                    comprehension,
173                });
174            }
175            other => parent.push(other.clone()),
176        }
177    }
178    Ok((PolydatFile { statements: parent }, fors, producers))
179}
180
181/// Resolve a traversal's source to a comprehension: inline text as is, a
182/// producer reference to the producer bound in the same scope, and a
183/// derivation to the base producer with its filter and order applied.
184pub fn resolve_source(source: &ForSource, producers: &[Producer]) -> Result<Comprehension, String> {
185    let find = |name: &str| -> Result<Comprehension, String> {
186        producers
187            .iter()
188            .rev()
189            .find(|p| p.name == name)
190            .map(|p| p.comprehension.clone())
191            .ok_or_else(|| {
192                let known: Vec<&str> = producers.iter().map(|p| p.name.as_str()).collect();
193                format!(
194                    "`for {}` at line {}, col {}: no producer named '{name}' is bound in this scope{}",
195                    source.text,
196                    source.span.line,
197                    source.span.col,
198                    if known.is_empty() { String::new() } else { format!("; producers here: {}", known.join(", ")) }
199                )
200            })
201    };
202    match &source.kind {
203        ForSourceKind::Comprehension(c) => Ok(c.clone()),
204        ForSourceKind::Producer(name) => find(name),
205        ForSourceKind::Derived {
206            base,
207            filter,
208            order,
209        } => {
210            let mut c = find(base)?;
211            if let Some(pred) = filter {
212                c = Comprehension::filter(c, pred.clone());
213            }
214            if let Some(spec) = order {
215                let (strategy, truncation) = parse_order(spec).map_err(|e| {
216                    format!(
217                        "`for {}` at line {}, col {}: {e}",
218                        source.text, source.span.line, source.span.col
219                    )
220                })?;
221                c = Comprehension::order(c, strategy, truncation);
222            }
223            Ok(c)
224        }
225    }
226}
227
228/// Parse an `order` spec such as `halton/5` into the algebra's strategy
229/// and truncation by running it through the comprehension parser on a
230/// one-clause carrier.
231fn parse_order(
232    spec: &str,
233) -> Result<(crate::iteration::comprehension::StrategyName, Option<u64>), String> {
234    let carrier = format!("__o in 0..1 order {spec}");
235    let legacy = crate::iteration::comprehension::parse::parse_comprehension_text(&carrier)?;
236    let algebra = crate::iteration::comprehension::spec::legacy_to_algebra(&legacy)
237        .map_err(|e| e.to_string())?;
238    match algebra {
239        Comprehension::Order {
240            strategy,
241            truncation,
242            ..
243        } => Ok((strategy, truncation)),
244        other => Err(format!(
245            "order spec `{spec}` did not produce an ordering (got {other:?})"
246        )),
247    }
248}
249
250/// Type each element name of a comprehension from its source, per SRD
251/// 113 §3.3. `probe` types a generator call expression the way the
252/// enclosing compiler would.
253pub fn element_types(
254    comprehension: &Comprehension,
255    probe: &mut dyn FnMut(&str) -> Result<PortType, String>,
256) -> Result<Vec<(String, PortType)>, String> {
257    let mut out = Vec::new();
258    collect_element_types(comprehension, probe, &mut out)?;
259    Ok(out)
260}
261
262fn collect_element_types(
263    c: &Comprehension,
264    probe: &mut dyn FnMut(&str) -> Result<PortType, String>,
265    out: &mut Vec<(String, PortType)>,
266) -> Result<(), String> {
267    match c {
268        Comprehension::Clause { name, source } => {
269            if out.iter().any(|(n, _)| n == name) {
270                return Ok(());
271            }
272            let ty = source_type(name, source, probe)?;
273            out.push((name.clone(), ty));
274            Ok(())
275        }
276        Comprehension::Cartesian { children } | Comprehension::Zip { children, .. } => {
277            for child in children {
278                collect_element_types(child, probe, out)?;
279            }
280            Ok(())
281        }
282        Comprehension::Union { children } => {
283            // Union children share one tuple shape; the first child's
284            // types stand for all of them.
285            if let Some(first) = children.first() {
286                collect_element_types(first, probe, out)?;
287            }
288            Ok(())
289        }
290        Comprehension::Filter { child, .. } | Comprehension::Order { child, .. } => {
291            collect_element_types(child, probe, out)
292        }
293    }
294}
295
296fn source_type(
297    name: &str,
298    source: &Source,
299    probe: &mut dyn FnMut(&str) -> Result<PortType, String>,
300) -> Result<PortType, String> {
301    match source {
302        Source::Literal { values } => {
303            let mut ty: Option<PortType> = None;
304            for v in values {
305                let t = match v {
306                    LiteralValue::Int(_) => PortType::U64,
307                    LiteralValue::Float(_) => PortType::F64,
308                    LiteralValue::String(_) => PortType::Str,
309                    LiteralValue::Bool(_) => PortType::Bool,
310                    LiteralValue::Json(_) => PortType::Json,
311                };
312                match ty {
313                    None => ty = Some(t),
314                    // An int among floats widens; anything else is mixed.
315                    Some(PortType::F64) if t == PortType::U64 => {}
316                    Some(PortType::U64) if t == PortType::F64 => ty = Some(PortType::F64),
317                    Some(prev) if prev != t => {
318                        return Err(format!(
319                            "element '{name}': literal list mixes {prev:?} and {t:?} values; a comprehension element has one type"
320                        ));
321                    }
322                    Some(_) => {}
323                }
324            }
325            ty.ok_or_else(|| format!("element '{name}': literal list is empty"))
326        }
327        Source::IntRange { .. } => Ok(PortType::U64),
328        Source::ContinuousInterval { .. } | Source::Distribution { .. } => Ok(PortType::F64),
329        // Workload parameter lists are text until a host types them.
330        Source::WorkloadParamList { .. } => Ok(PortType::Str),
331        Source::Generator { expr, .. } => {
332            let head = expr.trim();
333            if head.starts_with("partitions(")
334                || head.starts_with("subdivide(")
335                || head.ends_with(".partitions")
336            {
337                return Ok(PortType::Ext);
338            }
339            probe(head)
340                .map_err(|e| format!("element '{name}': cannot type generator `{head}`: {e}"))
341        }
342    }
343}
344
345/// Names a body declares itself, so they are not cascaded from the
346/// parent: binding targets, element names, declared inputs and externs,
347/// module names, and cursors with their projection sources.
348fn body_declared(body: &[Statement], elements: &[(String, PortType)]) -> BTreeSet<String> {
349    let mut names: BTreeSet<String> = elements.iter().map(|(n, _)| n.clone()).collect();
350    names.insert("cycle".to_string());
351    for stmt in body {
352        match stmt {
353            Statement::Binding(b) => names.extend(b.targets.iter().cloned()),
354            Statement::InputDecl(d) => {
355                names.insert(d.name.clone());
356            }
357            Statement::ExternPort(p) => {
358                names.insert(p.name.clone());
359            }
360            Statement::ModuleDef(m) => {
361                names.insert(m.name.clone());
362            }
363            Statement::Cursor(c) => {
364                names.insert(c.name.clone());
365            }
366            Statement::Pragma { .. } => {}
367            Statement::For(f) => {
368                // A nested traversal's own elements are its business,
369                // but a producer it binds is visible after it.
370                let _ = f;
371            }
372            Statement::Tile(t) => {
373                names.insert(t.name.clone());
374            }
375        }
376    }
377    names
378}
379
380/// Wires a tile's holes and branch conditions reference.
381fn tile_references(pieces: &[super::ast::TilePiece], out: &mut BTreeSet<String>) {
382    use super::ast::TilePiece;
383    use super::refs::collect_expr_refs;
384    for piece in pieces {
385        match piece {
386            TilePiece::Static(_) => {}
387            TilePiece::Hole(h) => collect_expr_refs(&h.expr, out),
388            TilePiece::Projection { body, .. } => tile_references(body, out),
389            TilePiece::Branch {
390                cond,
391                then,
392                otherwise,
393                ..
394            } => {
395                collect_expr_refs(cond, out);
396                tile_references(then, out);
397                if let Some(o) = otherwise {
398                    tile_references(o, out);
399                }
400            }
401        }
402    }
403}
404
405/// Names a body references, gathered from binding expressions, cursor
406/// constructors and `over` clauses, extern defaults, and nested `for`
407/// bodies. Nested bodies contribute their outer references so the
408/// cascade reaches through every level.
409fn body_references(body: &[Statement], out: &mut BTreeSet<String>) {
410    use super::refs::collect_expr_refs;
411    for stmt in body {
412        match stmt {
413            Statement::Binding(b) => collect_expr_refs(&b.value, out),
414            Statement::ExternPort(p) => {
415                if let Some(d) = &p.default {
416                    collect_expr_refs(d, out);
417                }
418            }
419            Statement::Cursor(c) => {
420                collect_expr_refs(&c.constructor, out);
421                if let Some(over) = &c.over {
422                    collect_expr_refs(over, out);
423                }
424            }
425            Statement::For(f) => {
426                let mut inner = BTreeSet::new();
427                body_references(&f.body, &mut inner);
428                let own = body_declared(&f.body, &[]);
429                let elems: BTreeSet<String> = f.source.element_names().into_iter().collect();
430                for n in inner {
431                    if !own.contains(&n) && !elems.contains(&n) {
432                        out.insert(n);
433                    }
434                }
435            }
436            Statement::Tile(t) => tile_references(&t.pieces, out),
437            Statement::InputDecl(_) | Statement::ModuleDef(_) | Statement::Pragma { .. } => {}
438        }
439    }
440}
441
442/// Build the child file for a traversal body: an implicit `cycle`
443/// input, one extern per element, one cascade extern per outer wire
444/// the body references and the parent exposes, then the body itself.
445/// Returns the file and the cascade list.
446pub fn child_file(
447    f: &ForStmt,
448    comprehension: &Comprehension,
449    elements: &[(String, PortType)],
450    type_of: &dyn Fn(&str) -> Option<PortType>,
451) -> Result<(PolydatFile, Vec<(String, PortType)>), String> {
452    // §7: a body may not declare a coordinate other than `cycle`.
453    for stmt in &f.body {
454        if let Statement::InputDecl(d) = stmt
455            && d.name != "cycle"
456        {
457            return Err(format!(
458                "`for {}` at line {}, col {}: a traversal body cannot declare input '{}'; only `cycle` is a coordinate inside a body, and the comprehension supplies the rest",
459                f.source.text, f.span.line, f.span.col, d.name
460            ));
461        }
462    }
463    let declared = body_declared(&f.body, elements);
464    let mut referenced = BTreeSet::new();
465    body_references(&f.body, &mut referenced);
466    // Outer wires the comprehension's own sources reference through
467    // `{name}` also cascade, so the tuple evaluation sees them.
468    referenced.extend(comprehension.referenced_source_names());
469
470    let mut cascade = Vec::new();
471    for name in referenced {
472        if declared.contains(&name) {
473            continue;
474        }
475        let ty = type_of(&name);
476        if let Some(ty) = ty {
477            cascade.push((name, ty));
478        }
479        // Names the parent does not expose stay unresolved; the child
480        // compile reports them as unknown wires with the body's spans.
481    }
482
483    let span = f.span;
484    let mut statements = Vec::with_capacity(f.body.len() + elements.len() + cascade.len() + 1);
485    if !f.body.iter().any(|s| matches!(s, Statement::InputDecl(_))) {
486        statements.push(Statement::InputDecl(InputDecl {
487            name: "cycle".into(),
488            ty: Some("u64".into()),
489            span,
490        }));
491    }
492    for (name, ty) in elements {
493        statements.push(Statement::ExternPort(ExternPort {
494            name: name.clone(),
495            typ: ty.to_keyword().to_string(),
496            default: None,
497            span,
498        }));
499    }
500    for (name, ty) in &cascade {
501        statements.push(Statement::ExternPort(ExternPort {
502            name: name.clone(),
503            typ: ty.to_keyword().to_string(),
504            default: None,
505            span,
506        }));
507    }
508    statements.extend(f.body.iter().cloned());
509    Ok((PolydatFile { statements }, cascade))
510}
polydat_core/dsl/traversal.rs

polydat_core/dsl/
traversal.rs
