fugue/program/ast.rs
1//! The serializable AST, its JSON front end, structural validation, and the
2//! canonical text printer.
3
4use std::fmt;
5
6use serde::{Deserialize, Serialize};
7
8use super::ProgramError;
9
10/// The program format version this build reads and writes (the
11/// `fugue_program` field of every serialized [`Program`]).
12///
13/// It changes only when a program that one version accepts would mean
14/// something else, or fail, under another; new node kinds that an old reader
15/// would reject outright also bump it, so an old reader fails with
16/// [`ProgramError::UnsupportedVersion`] instead of an obscure decode error.
17pub const FORMAT_VERSION: u32 = 1;
18
19/// How deeply a program may nest, counted as the nesting depth of its JSON
20/// form (every expression node is one level, a block is two).
21///
22/// serde_json refuses documents nested deeper than 127 levels; capping
23/// programs below that, identically in *both* front ends (and in
24/// [`Program::compile`](Program::compile)), guarantees that every program one
25/// front end accepts the other accepts too, and bounds the parser's and the
26/// interpreter's recursion on hostile input.
27pub const MAX_NESTING: usize = 100;
28
29/// A program: statements and a return expression, plus the format version.
30///
31/// A `Program` is plain data. Build one with [`Program::parse`] (the text
32/// syntax), [`Program::from_json`], or directly from the AST types; then
33/// [`compile`](Program::compile) it against a [`Registry`](super::Registry)
34/// and [`Data`](super::Data) to get a [`CompiledProgram`](super::CompiledProgram)
35/// that builds real `Model`s.
36///
37/// `Display` prints the canonical text form, which [`Program::parse`] reads
38/// back to an equal program.
39///
40/// ```rust
41/// use fugue::program::Program;
42///
43/// let text = Program::parse(
44/// r#"let p <- sample(addr!("p"), Beta(2.0, 2.0));
45/// observe(addr!("y"), Bernoulli(p), true);
46/// pure(p)"#,
47/// )
48/// .unwrap();
49/// let json = Program::from_json(&text.to_json().unwrap()).unwrap();
50/// assert_eq!(text, json);
51/// assert_eq!(Program::parse(&text.to_string()).unwrap(), text);
52/// ```
53#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
54#[serde(deny_unknown_fields)]
55pub struct Program {
56 /// Format version; [`FORMAT_VERSION`] for programs this build writes.
57 pub fugue_program: u32,
58 /// The statements, in order.
59 pub body: Vec<Stmt>,
60 /// The expression `pure(..)` returns.
61 pub ret: Expr,
62}
63
64/// A statement. The JSON form is an object tagged by `"stmt"`.
65#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
66#[serde(tag = "stmt", rename_all = "snake_case", deny_unknown_fields)]
67pub enum Stmt {
68 /// `let var <- sample(addr, dist);` — draw from `dist` at `addr` and bind
69 /// the draw (with the site's natural type) to `var`.
70 Sample {
71 /// The variable the draw is bound to.
72 var: String,
73 /// The site's address.
74 addr: Addr,
75 /// The distribution to draw from.
76 dist: DistCall,
77 },
78 /// `let var = value;` — bind a new variable (shadowing any outer one
79 /// until the end of the enclosing block).
80 Let {
81 /// The new variable.
82 var: String,
83 /// Its value.
84 value: Expr,
85 },
86 /// `var = value;` — reassign a variable bound by an enclosing `let`
87 /// (or by the data). The new value is visible for the rest of that
88 /// variable's scope, including later iterations of an enclosing loop.
89 Assign {
90 /// The variable to reassign.
91 var: String,
92 /// Its new value.
93 value: Expr,
94 },
95 /// `observe(addr, dist, value);` — condition on `value` under `dist`.
96 Observe {
97 /// The observation's address.
98 addr: Addr,
99 /// The distribution the value is scored under.
100 dist: DistCall,
101 /// The observed value.
102 value: Expr,
103 },
104 /// `factor(logw);` — add `logw` to the execution's log-weight.
105 Factor {
106 /// The log-weight.
107 logw: Expr,
108 },
109 /// `for var in start..end { body }` — run `body` with `var` bound to each
110 /// integer of the half-open range, evaluated once on entry.
111 For {
112 /// The loop variable (scoped to the body).
113 var: String,
114 /// Inclusive start.
115 start: Expr,
116 /// Exclusive end.
117 end: Expr,
118 /// The loop body.
119 body: Vec<Stmt>,
120 },
121 /// `if cond { then } else { else }` (`else` optional; `else if` chains
122 /// nest an `If` as the only statement of the else branch).
123 If {
124 /// The condition.
125 cond: Expr,
126 /// Statements run when the condition is true.
127 #[serde(rename = "then")]
128 then_branch: Vec<Stmt>,
129 /// Statements run when it is false (omitted from JSON when empty).
130 #[serde(rename = "else", default, skip_serializing_if = "Vec::is_empty")]
131 else_branch: Vec<Stmt>,
132 },
133 /// `break;` — leave the innermost `for` loop.
134 Break {},
135}
136
137/// A site address: `addr!("name")` or `addr!("name", index)`, built with the
138/// same encoding as fugue's `addr!` macro.
139#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
140#[serde(deny_unknown_fields)]
141pub struct Addr {
142 /// The address name (any string; escaped exactly as `addr!` escapes it).
143 pub name: String,
144 /// The optional index expression (omitted from JSON when absent).
145 #[serde(default, skip_serializing_if = "Option::is_none")]
146 pub index: Option<Expr>,
147}
148
149/// A distribution constructor call: a registered distribution's name and its
150/// argument expressions (`Normal(mu, 1.0)`).
151#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
152#[serde(deny_unknown_fields)]
153pub struct DistCall {
154 /// The registered distribution name.
155 pub name: String,
156 /// The constructor's arguments.
157 #[serde(default)]
158 pub args: Vec<Expr>,
159}
160
161/// An expression. The JSON form is an object tagged by `"op"`.
162#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
163#[serde(tag = "op", rename_all = "snake_case", deny_unknown_fields)]
164pub enum Expr {
165 /// A float literal (`2.0`, `1e-3`); must be finite.
166 F64 {
167 /// The literal's value.
168 value: f64,
169 },
170 /// An integer literal (`2`).
171 Int {
172 /// The literal's value.
173 value: i64,
174 },
175 /// `true` or `false`.
176 Bool {
177 /// The literal's value.
178 value: bool,
179 },
180 /// A variable reference.
181 Var {
182 /// The variable's name.
183 name: String,
184 },
185 /// An array literal `[a, b, c]` (its elements must evaluate to numbers).
186 Array {
187 /// The element expressions.
188 items: Vec<Expr>,
189 },
190 /// `array[index]`.
191 Index {
192 /// The indexed array.
193 array: Box<Expr>,
194 /// The index.
195 index: Box<Expr>,
196 },
197 /// `array.len()`.
198 Len {
199 /// The array.
200 array: Box<Expr>,
201 },
202 /// `-arg`.
203 Neg {
204 /// The operand.
205 arg: Box<Expr>,
206 },
207 /// `!arg`.
208 Not {
209 /// The operand.
210 arg: Box<Expr>,
211 },
212 /// `lhs + rhs`.
213 Add {
214 /// Left operand.
215 lhs: Box<Expr>,
216 /// Right operand.
217 rhs: Box<Expr>,
218 },
219 /// `lhs - rhs`.
220 Sub {
221 /// Left operand.
222 lhs: Box<Expr>,
223 /// Right operand.
224 rhs: Box<Expr>,
225 },
226 /// `lhs * rhs`.
227 Mul {
228 /// Left operand.
229 lhs: Box<Expr>,
230 /// Right operand.
231 rhs: Box<Expr>,
232 },
233 /// `lhs / rhs` (always float division).
234 Div {
235 /// Left operand.
236 lhs: Box<Expr>,
237 /// Right operand.
238 rhs: Box<Expr>,
239 },
240 /// `lhs == rhs`.
241 Eq {
242 /// Left operand.
243 lhs: Box<Expr>,
244 /// Right operand.
245 rhs: Box<Expr>,
246 },
247 /// `lhs != rhs`.
248 Ne {
249 /// Left operand.
250 lhs: Box<Expr>,
251 /// Right operand.
252 rhs: Box<Expr>,
253 },
254 /// `lhs < rhs`.
255 Lt {
256 /// Left operand.
257 lhs: Box<Expr>,
258 /// Right operand.
259 rhs: Box<Expr>,
260 },
261 /// `lhs <= rhs`.
262 Le {
263 /// Left operand.
264 lhs: Box<Expr>,
265 /// Right operand.
266 rhs: Box<Expr>,
267 },
268 /// `lhs > rhs`.
269 Gt {
270 /// Left operand.
271 lhs: Box<Expr>,
272 /// Right operand.
273 rhs: Box<Expr>,
274 },
275 /// `lhs >= rhs`.
276 Ge {
277 /// Left operand.
278 lhs: Box<Expr>,
279 /// Right operand.
280 rhs: Box<Expr>,
281 },
282 /// `lhs && rhs` (short-circuit).
283 And {
284 /// Left operand.
285 lhs: Box<Expr>,
286 /// Right operand.
287 rhs: Box<Expr>,
288 },
289 /// `lhs || rhs` (short-circuit).
290 Or {
291 /// Left operand.
292 lhs: Box<Expr>,
293 /// Right operand.
294 rhs: Box<Expr>,
295 },
296 /// A call of a registered pure function: `exp(x)`, `max(a, b)`.
297 Call {
298 /// The registered function name.
299 func: String,
300 /// The arguments.
301 #[serde(default)]
302 args: Vec<Expr>,
303 },
304}
305
306/// The binary operators, for building and inspecting [`Expr`] generically.
307///
308/// ```rust
309/// use fugue::program::{BinOp, Expr};
310///
311/// let e = Expr::binary(BinOp::Mul, Expr::F64 { value: 2.0 }, Expr::var("p"));
312/// assert_eq!(e.to_string(), "2.0 * p");
313/// assert!(matches!(e.as_binary(), Some((BinOp::Mul, _, _))));
314/// ```
315#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
316pub enum BinOp {
317 /// `+`
318 Add,
319 /// `-`
320 Sub,
321 /// `*`
322 Mul,
323 /// `/`
324 Div,
325 /// `==`
326 Eq,
327 /// `!=`
328 Ne,
329 /// `<`
330 Lt,
331 /// `<=`
332 Le,
333 /// `>`
334 Gt,
335 /// `>=`
336 Ge,
337 /// `&&`
338 And,
339 /// `||`
340 Or,
341}
342
343impl BinOp {
344 /// The operator's text spelling (`"+"`, `"<="`, `"&&"`, ...).
345 pub fn symbol(self) -> &'static str {
346 match self {
347 BinOp::Add => "+",
348 BinOp::Sub => "-",
349 BinOp::Mul => "*",
350 BinOp::Div => "/",
351 BinOp::Eq => "==",
352 BinOp::Ne => "!=",
353 BinOp::Lt => "<",
354 BinOp::Le => "<=",
355 BinOp::Gt => ">",
356 BinOp::Ge => ">=",
357 BinOp::And => "&&",
358 BinOp::Or => "||",
359 }
360 }
361
362 /// Binding strength, Rust's: `||` < `&&` < comparisons < `+ -` < `* /`.
363 pub(crate) fn precedence(self) -> u8 {
364 match self {
365 BinOp::Or => 1,
366 BinOp::And => 2,
367 BinOp::Eq | BinOp::Ne | BinOp::Lt | BinOp::Le | BinOp::Gt | BinOp::Ge => 3,
368 BinOp::Add | BinOp::Sub => 4,
369 BinOp::Mul | BinOp::Div => 5,
370 }
371 }
372}
373
374impl Expr {
375 /// A variable reference.
376 pub fn var(name: impl Into<String>) -> Expr {
377 Expr::Var { name: name.into() }
378 }
379
380 /// The binary expression `lhs op rhs`.
381 pub fn binary(op: BinOp, lhs: Expr, rhs: Expr) -> Expr {
382 let (lhs, rhs) = (Box::new(lhs), Box::new(rhs));
383 match op {
384 BinOp::Add => Expr::Add { lhs, rhs },
385 BinOp::Sub => Expr::Sub { lhs, rhs },
386 BinOp::Mul => Expr::Mul { lhs, rhs },
387 BinOp::Div => Expr::Div { lhs, rhs },
388 BinOp::Eq => Expr::Eq { lhs, rhs },
389 BinOp::Ne => Expr::Ne { lhs, rhs },
390 BinOp::Lt => Expr::Lt { lhs, rhs },
391 BinOp::Le => Expr::Le { lhs, rhs },
392 BinOp::Gt => Expr::Gt { lhs, rhs },
393 BinOp::Ge => Expr::Ge { lhs, rhs },
394 BinOp::And => Expr::And { lhs, rhs },
395 BinOp::Or => Expr::Or { lhs, rhs },
396 }
397 }
398
399 /// `Some((op, lhs, rhs))` for a binary expression, `None` otherwise.
400 pub fn as_binary(&self) -> Option<(BinOp, &Expr, &Expr)> {
401 let (op, lhs, rhs) = match self {
402 Expr::Add { lhs, rhs } => (BinOp::Add, lhs, rhs),
403 Expr::Sub { lhs, rhs } => (BinOp::Sub, lhs, rhs),
404 Expr::Mul { lhs, rhs } => (BinOp::Mul, lhs, rhs),
405 Expr::Div { lhs, rhs } => (BinOp::Div, lhs, rhs),
406 Expr::Eq { lhs, rhs } => (BinOp::Eq, lhs, rhs),
407 Expr::Ne { lhs, rhs } => (BinOp::Ne, lhs, rhs),
408 Expr::Lt { lhs, rhs } => (BinOp::Lt, lhs, rhs),
409 Expr::Le { lhs, rhs } => (BinOp::Le, lhs, rhs),
410 Expr::Gt { lhs, rhs } => (BinOp::Gt, lhs, rhs),
411 Expr::Ge { lhs, rhs } => (BinOp::Ge, lhs, rhs),
412 Expr::And { lhs, rhs } => (BinOp::And, lhs, rhs),
413 Expr::Or { lhs, rhs } => (BinOp::Or, lhs, rhs),
414 _ => return None,
415 };
416 Some((op, lhs, rhs))
417 }
418}
419
420// ---------------------------------------------------------------------------
421// Construction, JSON front end
422// ---------------------------------------------------------------------------
423
424impl Program {
425 /// A program at the current [`FORMAT_VERSION`].
426 pub fn new(body: Vec<Stmt>, ret: Expr) -> Program {
427 Program {
428 fugue_program: FORMAT_VERSION,
429 body,
430 ret,
431 }
432 }
433
434 /// Read a program from its JSON form.
435 ///
436 /// The `fugue_program` version is checked **before** the rest of the
437 /// document is decoded, so a program written by a newer format fails with
438 /// [`ProgramError::UnsupportedVersion`] rather than with whatever node its
439 /// newer syntax trips over first. Unknown fields are rejected, so a
440 /// misspelt key is an error, not a silently different program.
441 ///
442 /// ```rust
443 /// use fugue::program::{Program, ProgramError};
444 ///
445 /// let p = Program::from_json(
446 /// r#"{"fugue_program": 1, "body": [], "ret": {"op": "f64", "value": 1.5}}"#,
447 /// )
448 /// .unwrap();
449 /// assert_eq!(p, Program::parse("pure(1.5)").unwrap());
450 ///
451 /// let newer = r#"{"fugue_program": 2, "body": [], "ret": {"op": "f64", "value": 1.5}}"#;
452 /// assert!(matches!(
453 /// Program::from_json(newer),
454 /// Err(ProgramError::UnsupportedVersion { found: 2, supported: 1 })
455 /// ));
456 /// ```
457 pub fn from_json(json: &str) -> Result<Program, ProgramError> {
458 #[derive(Deserialize)]
459 struct VersionProbe {
460 fugue_program: Option<serde_json::Value>,
461 }
462 let probe: VersionProbe =
463 serde_json::from_str(json).map_err(|e| ProgramError::Json(e.to_string()))?;
464 let version = match probe.fugue_program {
465 None => {
466 return Err(ProgramError::Json(
467 "missing field `fugue_program` (the program format version)".to_string(),
468 ))
469 }
470 Some(v) => v.as_u64().ok_or_else(|| {
471 ProgramError::Json(format!(
472 "`fugue_program` must be a non-negative integer format version, found {v}"
473 ))
474 })?,
475 };
476 if version != u64::from(FORMAT_VERSION) {
477 return Err(ProgramError::UnsupportedVersion {
478 found: version,
479 supported: FORMAT_VERSION,
480 });
481 }
482 let program: Program =
483 serde_json::from_str(json).map_err(|e| ProgramError::Json(e.to_string()))?;
484 program.validate()?;
485 Ok(program)
486 }
487
488 /// The program's JSON form (compact). Fails for a program that would not
489 /// read back: a non-current version, a non-finite literal, a name that is
490 /// not an identifier, or nesting beyond [`MAX_NESTING`].
491 pub fn to_json(&self) -> Result<String, ProgramError> {
492 self.validate()?;
493 serde_json::to_string(self).map_err(|e| ProgramError::Json(e.to_string()))
494 }
495
496 /// The program's JSON form, pretty-printed. Same checks as
497 /// [`to_json`](Program::to_json).
498 pub fn to_json_pretty(&self) -> Result<String, ProgramError> {
499 self.validate()?;
500 serde_json::to_string_pretty(self).map_err(|e| ProgramError::Json(e.to_string()))
501 }
502
503 /// Structural checks shared by both front ends, `to_json` and `compile`:
504 /// the version, finite literals, identifier names, and the nesting cap.
505 /// Name *resolution* (variables, registry entries, `break` placement) is
506 /// [`compile`](Program::compile)'s job.
507 pub(crate) fn validate(&self) -> Result<(), ProgramError> {
508 if self.fugue_program != FORMAT_VERSION {
509 return Err(ProgramError::UnsupportedVersion {
510 found: u64::from(self.fugue_program),
511 supported: FORMAT_VERSION,
512 });
513 }
514 let depth = self.nesting_depth();
515 if depth > MAX_NESTING {
516 return Err(ProgramError::Check(format!(
517 "program nests {depth} levels deep; the format allows at most {MAX_NESTING}"
518 )));
519 }
520 // Iterative walk (the depth is bounded now, but a walk that cannot
521 // overflow is no dearer than one that can).
522 let mut stmts: Vec<&Stmt> = self.body.iter().collect();
523 let mut exprs: Vec<&Expr> = vec![&self.ret];
524 let ident = |kind: &str, name: &str| -> Result<(), ProgramError> {
525 if is_ident(name) {
526 Ok(())
527 } else {
528 Err(ProgramError::Check(format!(
529 "{kind} name `{name}` is not an identifier (or is a reserved word)"
530 )))
531 }
532 };
533 while let Some(s) = stmts.pop() {
534 match s {
535 Stmt::Sample { var, addr, dist } => {
536 ident("variable", var)?;
537 ident("distribution", &dist.name)?;
538 exprs.extend(addr.index.iter());
539 exprs.extend(dist.args.iter());
540 }
541 Stmt::Let { var, value } | Stmt::Assign { var, value } => {
542 ident("variable", var)?;
543 exprs.push(value);
544 }
545 Stmt::Observe { addr, dist, value } => {
546 ident("distribution", &dist.name)?;
547 exprs.extend(addr.index.iter());
548 exprs.extend(dist.args.iter());
549 exprs.push(value);
550 }
551 Stmt::Factor { logw } => exprs.push(logw),
552 Stmt::For {
553 var,
554 start,
555 end,
556 body,
557 } => {
558 ident("variable", var)?;
559 exprs.push(start);
560 exprs.push(end);
561 stmts.extend(body.iter());
562 }
563 Stmt::If {
564 cond,
565 then_branch,
566 else_branch,
567 } => {
568 exprs.push(cond);
569 stmts.extend(then_branch.iter());
570 stmts.extend(else_branch.iter());
571 }
572 Stmt::Break {} => {}
573 }
574 }
575 while let Some(e) = exprs.pop() {
576 match e {
577 Expr::F64 { value } if !value.is_finite() => {
578 return Err(ProgramError::Check(format!(
579 "number literal {value} is not finite (JSON cannot represent it)"
580 )))
581 }
582 Expr::F64 { .. } | Expr::Int { .. } | Expr::Bool { .. } => {}
583 Expr::Var { name } => ident("variable", name)?,
584 Expr::Array { items } => exprs.extend(items.iter()),
585 Expr::Index { array, index } => {
586 exprs.push(array);
587 exprs.push(index);
588 }
589 Expr::Len { array: arg } | Expr::Neg { arg } | Expr::Not { arg } => exprs.push(arg),
590 Expr::Call { func, args } => {
591 ident("function", func)?;
592 exprs.extend(args.iter());
593 }
594 other => {
595 let (_, lhs, rhs) = other.as_binary().expect("remaining variants are binary");
596 exprs.push(lhs);
597 exprs.push(rhs);
598 }
599 }
600 }
601 Ok(())
602 }
603
604 /// The nesting depth of the program's JSON form, computed without
605 /// recursion (so it is safe on any AST, however deep).
606 pub(crate) fn nesting_depth(&self) -> usize {
607 // Depths follow the serialized shape: the program object is level 1,
608 // its `body` array level 2, each statement object level 3; a
609 // statement's direct expressions sit one level below it, a block two
610 // (array + statement object), a distribution's arguments three
611 // (dist object + args array + expression).
612 let mut max = 2; // program object + body array
613 let mut stmts: Vec<(&Stmt, usize)> = self.body.iter().map(|s| (s, 3)).collect();
614 let mut exprs: Vec<(&Expr, usize)> = vec![(&self.ret, 2)];
615 while let Some((s, d)) = stmts.pop() {
616 max = max.max(d);
617 match s {
618 Stmt::Sample { addr, dist, .. } | Stmt::Observe { addr, dist, .. } => {
619 // addr/dist objects at d + 1, the (always written) args
620 // array at d + 2, its expressions at d + 3.
621 max = max.max(d + 2);
622 if let Some(ix) = &addr.index {
623 exprs.push((ix, d + 2));
624 }
625 exprs.extend(dist.args.iter().map(|a| (a, d + 3)));
626 if let Stmt::Observe { value, .. } = s {
627 exprs.push((value, d + 1));
628 }
629 }
630 Stmt::Let { value, .. } | Stmt::Assign { value, .. } => exprs.push((value, d + 1)),
631 Stmt::Factor { logw } => exprs.push((logw, d + 1)),
632 Stmt::For {
633 start, end, body, ..
634 } => {
635 exprs.push((start, d + 1));
636 exprs.push((end, d + 1));
637 max = max.max(d + 1); // body array (even when empty)
638 stmts.extend(body.iter().map(|s| (s, d + 2)));
639 }
640 Stmt::If {
641 cond,
642 then_branch,
643 else_branch,
644 } => {
645 exprs.push((cond, d + 1));
646 max = max.max(d + 1); // then array (always written)
647 stmts.extend(then_branch.iter().map(|s| (s, d + 2)));
648 stmts.extend(else_branch.iter().map(|s| (s, d + 2)));
649 }
650 Stmt::Break {} => {}
651 }
652 }
653 while let Some((e, d)) = exprs.pop() {
654 max = max.max(d);
655 match e {
656 Expr::F64 { .. } | Expr::Int { .. } | Expr::Bool { .. } | Expr::Var { .. } => {}
657 Expr::Array { items: args } | Expr::Call { args, .. } => {
658 max = max.max(d + 1); // the array itself
659 exprs.extend(args.iter().map(|a| (a, d + 2)));
660 }
661 Expr::Index { array, index } => {
662 exprs.push((array, d + 1));
663 exprs.push((index, d + 1));
664 }
665 Expr::Len { array: arg } | Expr::Neg { arg } | Expr::Not { arg } => {
666 exprs.push((arg, d + 1))
667 }
668 other => {
669 let (_, lhs, rhs) = other.as_binary().expect("remaining variants are binary");
670 exprs.push((lhs, d + 1));
671 exprs.push((rhs, d + 1));
672 }
673 }
674 }
675 max
676 }
677}
678
679/// Words the text syntax reserves; they cannot name variables, distributions
680/// or functions.
681pub(crate) const RESERVED: &[&str] = &[
682 "let", "mut", "observe", "factor", "for", "in", "if", "else", "break", "pure", "true", "false",
683];
684
685/// An ASCII identifier (`[A-Za-z_][A-Za-z0-9_]*`) that is not reserved.
686pub(crate) fn is_ident(s: &str) -> bool {
687 let mut chars = s.chars();
688 matches!(chars.next(), Some(c) if c.is_ascii_alphabetic() || c == '_')
689 && chars.all(|c| c.is_ascii_alphanumeric() || c == '_')
690 && !RESERVED.contains(&s)
691}
692
693// ---------------------------------------------------------------------------
694// Canonical text printer
695// ---------------------------------------------------------------------------
696
697/// The canonical text form; [`Program::parse`] reads it back to an equal
698/// program (for any program [`to_json`](Program::to_json) accepts).
699impl fmt::Display for Program {
700 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
701 for s in &self.body {
702 write_stmt(f, s, 0)?;
703 }
704 write!(f, "pure({})", self.ret)
705 }
706}
707
708/// Prints the expression in text syntax with the minimal parentheses that
709/// read back to the same tree.
710impl fmt::Display for Expr {
711 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
712 write_expr(f, self)
713 }
714}
715
716fn indent(f: &mut fmt::Formatter<'_>, depth: usize) -> fmt::Result {
717 for _ in 0..depth {
718 f.write_str(" ")?;
719 }
720 Ok(())
721}
722
723fn write_block(f: &mut fmt::Formatter<'_>, stmts: &[Stmt], depth: usize) -> fmt::Result {
724 f.write_str("{\n")?;
725 for s in stmts {
726 write_stmt(f, s, depth + 1)?;
727 }
728 indent(f, depth)?;
729 f.write_str("}")
730}
731
732fn write_stmt(f: &mut fmt::Formatter<'_>, s: &Stmt, depth: usize) -> fmt::Result {
733 indent(f, depth)?;
734 match s {
735 Stmt::Sample { var, addr, dist } => {
736 writeln!(f, "let {var} <- sample({addr}, {dist});")
737 }
738 Stmt::Let { var, value } => writeln!(f, "let {var} = {value};"),
739 Stmt::Assign { var, value } => writeln!(f, "{var} = {value};"),
740 Stmt::Observe { addr, dist, value } => {
741 writeln!(f, "observe({addr}, {dist}, {value});")
742 }
743 Stmt::Factor { logw } => writeln!(f, "factor({logw});"),
744 Stmt::For {
745 var,
746 start,
747 end,
748 body,
749 } => {
750 write!(f, "for {var} in {start}..{end} ")?;
751 write_block(f, body, depth)?;
752 f.write_str("\n")
753 }
754 Stmt::If { .. } => {
755 write_if(f, s, depth)?;
756 f.write_str("\n")
757 }
758 Stmt::Break {} => f.write_str("break;\n"),
759 }
760}
761
762fn write_if(f: &mut fmt::Formatter<'_>, s: &Stmt, depth: usize) -> fmt::Result {
763 let Stmt::If {
764 cond,
765 then_branch,
766 else_branch,
767 } = s
768 else {
769 unreachable!("write_if is only called on `If`")
770 };
771 write!(f, "if {cond} ")?;
772 write_block(f, then_branch, depth)?;
773 match else_branch.as_slice() {
774 [] => Ok(()),
775 [nested @ Stmt::If { .. }] => {
776 f.write_str(" else ")?;
777 write_if(f, nested, depth)
778 }
779 stmts => {
780 f.write_str(" else ")?;
781 write_block(f, stmts, depth)
782 }
783 }
784}
785
786impl fmt::Display for Addr {
787 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
788 f.write_str("addr!(")?;
789 write_str_literal(f, &self.name)?;
790 if let Some(ix) = &self.index {
791 write!(f, ", {ix}")?;
792 }
793 f.write_str(")")
794 }
795}
796
797impl fmt::Display for DistCall {
798 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
799 write!(f, "{}(", self.name)?;
800 write_list(f, &self.args)?;
801 f.write_str(")")
802 }
803}
804
805fn write_list(f: &mut fmt::Formatter<'_>, items: &[Expr]) -> fmt::Result {
806 for (i, a) in items.iter().enumerate() {
807 if i > 0 {
808 f.write_str(", ")?;
809 }
810 write_expr(f, a)?;
811 }
812 Ok(())
813}
814
815/// A Rust string literal the lexer reads back to exactly `s`.
816fn write_str_literal(f: &mut fmt::Formatter<'_>, s: &str) -> fmt::Result {
817 f.write_str("\"")?;
818 for c in s.chars() {
819 match c {
820 '\\' => f.write_str("\\\\")?,
821 '"' => f.write_str("\\\"")?,
822 '\n' => f.write_str("\\n")?,
823 '\r' => f.write_str("\\r")?,
824 '\t' => f.write_str("\\t")?,
825 '\0' => f.write_str("\\0")?,
826 c => write!(f, "{c}")?,
827 }
828 }
829 f.write_str("\"")
830}
831
832/// Precedence of an expression as printed: binary operators their own,
833/// prefix operators and negative literals 6, postfix/primary forms 7.
834fn expr_prec(e: &Expr) -> u8 {
835 match e {
836 Expr::F64 { value } if value.is_sign_negative() => 6,
837 Expr::Int { value } if *value < 0 => 6,
838 Expr::Neg { .. } | Expr::Not { .. } => 6,
839 other => other.as_binary().map_or(7, |(op, _, _)| op.precedence()),
840 }
841}
842
843fn write_expr(f: &mut fmt::Formatter<'_>, e: &Expr) -> fmt::Result {
844 match e {
845 // `{:?}` keeps a `.0` or an exponent, so the literal reads back as a
846 // float; it prints the shortest digits that round-trip.
847 Expr::F64 { value } => write!(f, "{value:?}"),
848 Expr::Int { value } => write!(f, "{value}"),
849 Expr::Bool { value } => write!(f, "{value}"),
850 Expr::Var { name } => f.write_str(name),
851 Expr::Array { items } => {
852 f.write_str("[")?;
853 write_list(f, items)?;
854 f.write_str("]")
855 }
856 Expr::Index { array, index } => {
857 write_operand(f, array, 7)?;
858 write!(f, "[{index}]")
859 }
860 Expr::Len { array } => {
861 write_operand(f, array, 7)?;
862 f.write_str(".len()")
863 }
864 Expr::Neg { arg } => {
865 // The parser folds `-2.0` into one negative literal, so a `Neg`
866 // of a literal prints as `-(2.0)` to read back as a `Neg`.
867 if matches!(**arg, Expr::F64 { .. } | Expr::Int { .. }) {
868 write!(f, "-({arg})")
869 } else {
870 f.write_str("-")?;
871 write_operand(f, arg, 6)
872 }
873 }
874 Expr::Not { arg } => {
875 f.write_str("!")?;
876 write_operand(f, arg, 6)
877 }
878 Expr::Call { func, args } => {
879 write!(f, "{func}(")?;
880 write_list(f, args)?;
881 f.write_str(")")
882 }
883 other => {
884 let (op, lhs, rhs) = other.as_binary().expect("remaining variants are binary");
885 let p = op.precedence();
886 // Left-associative, and comparisons do not chain: the left operand
887 // needs parentheses below this precedence (at it, for a
888 // comparison), the right operand at or below it.
889 let lhs_min = if p == 3 { p + 1 } else { p };
890 write_operand(f, lhs, lhs_min)?;
891 write!(f, " {} ", op.symbol())?;
892 write_operand(f, rhs, p + 1)
893 }
894 }
895}
896
897/// Print `e`, parenthesized unless it binds at least as tightly as `min`.
898fn write_operand(f: &mut fmt::Formatter<'_>, e: &Expr, min: u8) -> fmt::Result {
899 if expr_prec(e) >= min {
900 write_expr(f, e)
901 } else {
902 write!(f, "({e})")
903 }
904}