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

polydat_core/dsl/
traversal.rs
