probl_syntax/
sexpr.rs

1//! A compact S-expression rendering of the AST, used in tests.
2
3use crate::ast::*;
4use std::fmt::Write;
5
6pub fn program(p: &Program) -> String {
7    let mut out = String::new();
8    for pragma in &p.pragmas {
9        match &pragma.arg {
10            Some(arg) => writeln!(out, "(@{} {})", pragma.name.name, expr(arg)).unwrap(),
11            None => writeln!(out, "(@{})", pragma.name.name).unwrap(),
12        }
13    }
14    for item in &p.items {
15        writeln!(out, "{}", self::item(item)).unwrap();
16    }
17    out
18}
19
20pub fn item(item: &Item) -> String {
21    match item {
22        Item::Fn(f) => {
23            let params: Vec<String> = f
24                .params
25                .iter()
26                .map(|p| match &p.ty {
27                    Some(t) => format!("{}: {}", p.name.name, ty(t)),
28                    None => p.name.name.clone(),
29                })
30                .collect();
31            let ret = f.ret.as_ref().map(|t| format!(" -> {}", ty(t))).unwrap_or_default();
32            format!("(fn {} ({}){} {})", f.name.name, params.join(", "), ret, block(&f.body))
33        }
34        Item::Type(t) => format!("(type {} {})", t.name.name, ty(&t.ty)),
35        Item::Enum(e) => {
36            let names: Vec<&str> = e.variants.iter().map(|v| v.name.as_str()).collect();
37            format!("(enum {} {})", e.name.name, names.join(" "))
38        }
39        Item::Import(i) => format!("(import {:?})", i.path),
40        Item::Stmt(s) => stmt(s),
41    }
42}
43
44pub fn ty(t: &TypeExpr) -> String {
45    match t {
46        TypeExpr::Named { name, args } if args.is_empty() => name.name.clone(),
47        TypeExpr::Named { name, args } => {
48            let args: Vec<String> = args.iter().map(ty).collect();
49            format!("{}[{}]", name.name, args.join(", "))
50        }
51        TypeExpr::Record { fields, .. } => {
52            let fields: Vec<String> = fields.iter().map(|(n, t)| format!("{}: {}", n.name, ty(t))).collect();
53            format!("{{{}}}", fields.join(", "))
54        }
55    }
56}
57
58pub fn block(b: &Block) -> String {
59    let stmts: Vec<String> = b.stmts.iter().map(stmt).collect();
60    if stmts.is_empty() {
61        "(block)".to_string()
62    } else {
63        format!("(block {})", stmts.join(" "))
64    }
65}
66
67pub fn stmt(s: &Stmt) -> String {
68    match &s.kind {
69        StmtKind::Let {
70            mutable,
71            pattern: p,
72            ty: t,
73            op,
74            value,
75        } => {
76            let kw = if *mutable { "var" } else { "let" };
77            let op = match op {
78                BindOp::Assign => "=",
79                BindOp::Draw => "~",
80            };
81            let t = t.as_ref().map(|t| format!(": {}", ty(t))).unwrap_or_default();
82            format!("({kw} {}{t} {op} {})", pattern(p), expr(value))
83        }
84        StmtKind::Assign { target, op, value } => {
85            let op = match op {
86                AssignOp::Set => "=",
87                AssignOp::Draw => "~",
88                AssignOp::Add => "+=",
89                AssignOp::Sub => "-=",
90                AssignOp::Mul => "*=",
91                AssignOp::Div => "/=",
92            };
93            format!("({op} {} {})", expr(target), expr(value))
94        }
95        StmtKind::For { pattern: p, iter, body } => format!("(for {} {} {})", pattern(p), expr(iter), block(body)),
96        StmtKind::While { cond, body } => format!("(while {} {})", expr(cond), block(body)),
97        StmtKind::Repeat { count, body } => format!("(repeat {} {})", expr(count), block(body)),
98        StmtKind::Loop { body } => format!("(loop {})", block(body)),
99        StmtKind::Break => "(break)".to_string(),
100        StmtKind::Continue => "(continue)".to_string(),
101        StmtKind::Return(None) => "(return)".to_string(),
102        StmtKind::Return(Some(e)) => format!("(return {})", expr(e)),
103        StmtKind::Score(value) => format!("(score {})", expr(value)),
104        StmtKind::Observe { value, from } => match from {
105            Some(d) => format!("(observe {} from {})", expr(value), expr(d)),
106            None => format!("(observe {})", expr(value)),
107        },
108        StmtKind::Report { value, by, label } => {
109            let mut out = format!("(report {}", expr(value));
110            if let Some(by) = by {
111                write!(out, " by {}", expr(by)).unwrap();
112            }
113            if let Some((label, _)) = label {
114                write!(out, " as {label:?}").unwrap();
115            }
116            out.push(')');
117            out
118        }
119        StmtKind::Expr(e) => expr(e),
120    }
121}
122
123pub fn pattern(p: &Pattern) -> String {
124    match &p.kind {
125        PatternKind::Wildcard => "_".to_string(),
126        PatternKind::Name(n) => n.clone(),
127        PatternKind::Literal(e) => expr(e),
128        PatternKind::List(items) => {
129            let items: Vec<String> = items.iter().map(pattern).collect();
130            format!("[{}]", items.join(" "))
131        }
132        PatternKind::Or(alts) => {
133            let alts: Vec<String> = alts.iter().map(pattern).collect();
134            format!("(| {})", alts.join(" "))
135        }
136    }
137}
138
139fn args(args: &[Arg]) -> String {
140    args.iter()
141        .map(|a| match &a.name {
142            Some(n) => format!(" {}: {}", n.name, expr(&a.value)),
143            None => format!(" {}", expr(&a.value)),
144        })
145        .collect()
146}
147
148fn fields(fields: &[Field]) -> String {
149    fields
150        .iter()
151        .map(|f| format!(" ({} {})", f.name.name, expr(&f.value)))
152        .collect()
153}
154
155pub fn percent(v: f64) -> String {
156    let pct = (v * 100.0 * 1e9).round() / 1e9;
157    format!("{pct}%")
158}
159
160pub fn expr(e: &Expr) -> String {
161    match &e.kind {
162        ExprKind::Int(v) => v.to_string(),
163        ExprKind::Float(v) => format!("{v:?}"),
164        ExprKind::Percent(v) => percent(*v),
165        ExprKind::Dice { count, sides } => format!("{count}d{sides}"),
166        ExprKind::Bool(b) => b.to_string(),
167        ExprKind::Str(segments) => {
168            if let [StrSegment::Lit(text)] = segments.as_slice() {
169                return format!("{text:?}");
170            }
171            let parts: Vec<String> = segments
172                .iter()
173                .map(|s| match s {
174                    StrSegment::Lit(text) => format!("{text:?}"),
175                    StrSegment::Expr(e) => expr(e),
176                })
177                .collect();
178            format!("(str {})", parts.join(" "))
179        }
180        ExprKind::Name(n) => n.clone(),
181        ExprKind::List(items) => {
182            let items: Vec<String> = items.iter().map(expr).collect();
183            format!("[{}]", items.join(" "))
184        }
185        ExprKind::Map(entries) => {
186            let entries: Vec<String> = entries
187                .iter()
188                .map(|(k, v)| format!("{}: {}", expr(k), expr(v)))
189                .collect();
190            format!(
191                "[{}]",
192                if entries.is_empty() {
193                    ":".to_string()
194                } else {
195                    entries.join(", ")
196                }
197            )
198        }
199        ExprKind::Record { name, fields: fs } => match name {
200            Some(n) => format!("(record {}{})", n.name, fields(fs)),
201            None => format!("(record{})", fields(fs)),
202        },
203        ExprKind::Draw(inner) => format!("(~ {})", expr(inner)),
204        ExprKind::Unary { op, expr: inner } => {
205            let op = match op {
206                UnOp::Neg => "-",
207                UnOp::Not => "not",
208                UnOp::Typeof => "typeof",
209            };
210            format!("({op} {})", expr(inner))
211        }
212        ExprKind::Binary { op, lhs, rhs } => {
213            format!("({} {} {})", op.symbol(), expr(lhs), expr(rhs))
214        }
215        ExprKind::Call { callee, args: a } => format!("(call {}{})", expr(callee), args(a)),
216        ExprKind::Method {
217            receiver,
218            name,
219            args: a,
220        } => format!("(.{} {}{})", name.name, expr(receiver), args(a)),
221        ExprKind::Field { expr: inner, name } => format!("(. {} {})", expr(inner), name.name),
222        ExprKind::Index { expr: inner, index } => {
223            format!("(index {} {})", expr(inner), expr(index))
224        }
225        ExprKind::With {
226            expr: inner,
227            fields: fs,
228        } => format!("(with {}{})", expr(inner), fields(fs)),
229        ExprKind::Lambda { params, body } => {
230            let params: Vec<&str> = params.iter().map(|p| p.name.as_str()).collect();
231            format!("(-> ({}) {})", params.join(" "), expr(body))
232        }
233        ExprKind::If { cond, then, otherwise } => match otherwise {
234            Some(other) => format!("(if {} {} {})", expr(cond), block(then), expr(other)),
235            None => format!("(if {} {})", expr(cond), block(then)),
236        },
237        ExprKind::Chance { arms } => {
238            let arms: Vec<String> = arms
239                .iter()
240                .map(|a| {
241                    let w = a.weight.as_ref().map(expr).unwrap_or_else(|| "else".to_string());
242                    format!("({w} => {})", stmt(&a.body))
243                })
244                .collect();
245            format!("(chance {})", arms.join(" "))
246        }
247        ExprKind::Match { scrutinee, arms } => {
248            let arms: Vec<String> = arms
249                .iter()
250                .map(|a| {
251                    let guard = a.guard.as_ref().map(|g| format!(" if {}", expr(g))).unwrap_or_default();
252                    format!("({}{guard} => {})", pattern(&a.pattern), stmt(&a.body))
253                })
254                .collect();
255            format!("(match {} {})", expr(scrutinee), arms.join(" "))
256        }
257        ExprKind::Simulate(b) => format!("(simulate {})", block(b)),
258        ExprKind::Block(b) => block(b),
259        ExprKind::Try { body, catches } => {
260            let catches: Vec<String> = catches
261                .iter()
262                .map(|c| match &c.fault {
263                    Some(f) => format!("(catch {} {})", f.name, block(&c.body)),
264                    None => format!("(catch {})", block(&c.body)),
265                })
266                .collect();
267            format!("(try {} {})", block(body), catches.join(" "))
268        }
269    }
270}
probl_syntax/sexpr.rs

probl_syntax/
sexpr.rs
