1use voltaic_ast::BinaryOp;
2
3use voltaic_ast::Block;
4
5use voltaic_ast::CaseArm;
6
7use voltaic_ast::ClassMember;
8
9use voltaic_ast::ClauseBody;
10
11use voltaic_ast::Condition;
12
13use voltaic_ast::ConditionOperand;
14
15use voltaic_ast::Expr;
16
17use voltaic_ast::FieldInit;
18
19use voltaic_ast::Guard;
20
21use voltaic_ast::Item;
22
23use voltaic_ast::Program;
24
25use voltaic_ast::Qualifier;
26
27use voltaic_ast::Section;
28
29use voltaic_ast::Stmt;
30
31use voltaic_ast::StrInterpPart;
32
33#[derive(Debug, Clone, Copy, PartialEq, Eq)]
34pub enum Pass {
35 BlankNumbers,
36 Pipelines,
37}
38#[allow(unused_imports)]
39use Pass::*;
40
41pub fn rewrite_program(program: &mut Program, step: Pass) {
42 for item in program.items.iter_mut() {
43 rewrite_item(item, step);
44 }
45 ();
46}
47
48pub fn rewrite_item(item: &mut Item, step: Pass) {
49 match item {
50 Item::Fn(f) => {
51 for clause in f.clauses.iter_mut() {
52 match &mut clause.body {
53 ClauseBody::Plain(block) => rewrite_block(block, step),
54 ClauseBody::Guarded(rows) => {
55 for row in rows.iter_mut() {
56 rewrite_guard(row, step);
57 }
58 ()
59 }
60 };
61 for stmt in clause.wheres.iter_mut() {
62 rewrite_stmt(stmt, step);
63 }
64 }
65 ()
66 }
67 Item::Instance(block) => {
68 for member in block.members.iter_mut() {
69 if let ClassMember::Method(f) = member {
70 for clause in f.clauses.iter_mut() {
71 match &mut clause.body {
72 ClauseBody::Plain(b) => rewrite_block(b, step),
73 ClauseBody::Guarded(rows) => {
74 for row in rows.iter_mut() {
75 rewrite_guard(row, step);
76 }
77 ()
78 }
79 };
80 for stmt in clause.wheres.iter_mut() {
81 rewrite_stmt(stmt, step);
82 }
83 }
84 }
85 }
86 ()
87 }
88 Item::Test(t) => rewrite_block(&mut t.body, step),
89 Item::Const(c) => rewrite_expr(&mut c.value, step),
90 _ => (),
91 };
92}
93
94pub fn rewrite_block(block: &mut Block, step: Pass) {
95 for stmt in block.stmts.iter_mut() {
96 rewrite_stmt(stmt, step);
97 }
98 ();
99}
100
101pub fn rewrite_stmt(stmt: &mut Stmt, step: Pass) {
102 rewrite_expr(&mut stmt.expr, step);
103 ();
104}
105
106pub fn rewrite_guard(row: &mut Guard, step: Pass) {
107 if let Some(cond) = &mut row.condition {
108 rewrite_condition(cond, step);
109 }
110 rewrite_block(&mut row.body, step);
111 ();
112}
113
114pub fn rewrite_condition(cond: &mut Condition, step: Pass) {
115 for operand in cond.operands.iter_mut() {
116 match operand {
117 ConditionOperand::Expr(e) => rewrite_expr(e, step),
118 ConditionOperand::Pattern { source, .. } => rewrite_expr(source, step),
119 };
120 }
121 ();
122}
123
124pub fn rewrite_arm(arm: &mut CaseArm, step: Pass) {
125 if let Some(cond) = &mut arm.r#guard {
126 rewrite_condition(cond, step);
127 }
128 rewrite_expr(&mut arm.body, step);
129 ();
130}
131
132pub fn rewrite_field(field: &mut FieldInit, step: Pass) {
133 if let Some(e) = &mut field.value {
134 rewrite_expr(e, step);
135 }
136 ();
137}
138
139pub fn rewrite_part(part: &mut StrInterpPart, step: Pass) {
140 match part {
141 StrInterpPart::Interp { expr, .. } => rewrite_expr(expr, step),
142 StrInterpPart::Text(_) => (),
143 };
144}
145
146pub fn rewrite_qualifier(q: &mut Qualifier, step: Pass) {
147 match q {
148 Qualifier::Draw { source, .. } => rewrite_expr(source, step),
149 Qualifier::Guard(e) => rewrite_expr(e, step),
150 };
151}
152
153pub fn rewrite_section(s: &mut Section, step: Pass) {
154 match s {
155 Section::Right { operand, .. } => rewrite_expr(operand, step),
156 Section::Left { operand, .. } => rewrite_expr(operand, step),
157 Section::Field(_) => (),
158 };
159}
160
161pub fn rewrite_expr(e: &mut Expr, step: Pass) {
162 descend(e, step);
163 apply_pass(e, step);
164 ();
165}
166
167pub fn descend(e: &mut Expr, step: Pass) {
168 match e {
169 Expr::Call { callee, args, .. } => {
170 rewrite_expr(callee, step);
171 for a in args.iter_mut() {
172 rewrite_expr(a, step);
173 }
174 ()
175 }
176 Expr::Field { recv, .. } => rewrite_expr(recv, step),
177 Expr::Dot { recv, args, .. } => {
178 rewrite_expr(recv, step);
179 for a in args.iter_mut() {
180 rewrite_expr(a, step);
181 }
182 ()
183 }
184 Expr::Index { recv, index, .. } => {
185 rewrite_expr(recv, step);
186 rewrite_expr(index, step)
187 }
188 Expr::Binary { lhs, rhs, .. } => {
189 rewrite_expr(lhs, step);
190 rewrite_expr(rhs, step)
191 }
192 Expr::UserBinary { lhs, rhs, .. } => {
193 rewrite_expr(lhs, step);
194 rewrite_expr(rhs, step)
195 }
196 Expr::Over {
197 function, place, ..
198 } => {
199 rewrite_expr(function, step);
200 rewrite_expr(place, step)
201 }
202 Expr::Unary { operand, .. } => rewrite_expr(operand, step),
203 Expr::Paren { inner, .. } => rewrite_expr(inner, step),
204 Expr::Lambda { body, .. } => rewrite_expr(body, step),
205 Expr::AsyncBlock { body, .. } => rewrite_block(body, step),
206 Expr::UnsafeBlock { body, .. } => rewrite_block(body, step),
207 Expr::TryBlock { body, .. } => rewrite_block(body, step),
208 Expr::BindBlock { body, .. } => rewrite_block(body, step),
209 Expr::Ascribe { inner, .. } => rewrite_expr(inner, step),
210 Expr::StrInterp { parts, .. } => {
211 for p in parts.iter_mut() {
212 rewrite_part(p, step);
213 }
214 ()
215 }
216 Expr::Tuple { items, .. } => {
217 for i in items.iter_mut() {
218 rewrite_expr(i, step);
219 }
220 ()
221 }
222 Expr::ArrayRepeat { value, count, .. } => {
223 rewrite_expr(value, step);
224 rewrite_expr(count, step)
225 }
226 Expr::Record { head, fields, .. } => {
227 rewrite_expr(head, step);
228 for f in fields.iter_mut() {
229 rewrite_field(f, step);
230 }
231 ()
232 }
233 Expr::List { items, .. } => {
234 for i in items.iter_mut() {
235 rewrite_expr(i, step);
236 }
237 ()
238 }
239 Expr::Comprehension {
240 element,
241 qualifiers,
242 ..
243 } => {
244 rewrite_expr(element, step);
245 for q in qualifiers.iter_mut() {
246 rewrite_qualifier(q, step);
247 }
248 ()
249 }
250 Expr::If {
251 condition,
252 then_block,
253 r#otherwise,
254 ..
255 } => {
256 rewrite_condition(condition, step);
257 rewrite_block(then_block, step);
258 if let Some(block) = r#otherwise {
259 rewrite_block(block, step);
260 }
261 ()
262 }
263 Expr::Block { body, .. } => rewrite_block(body, step),
264 Expr::Try { inner, .. } => rewrite_expr(inner, step),
265 Expr::Section { section, .. } => rewrite_section(section, step),
266 Expr::Case {
267 scrutinee, arms, ..
268 } => {
269 rewrite_expr(scrutinee, step);
270 for a in arms.iter_mut() {
271 rewrite_arm(a, step);
272 }
273 ()
274 }
275 Expr::Ladder { rows, .. } => {
276 for row in rows.iter_mut() {
277 rewrite_guard(row, step);
278 }
279 ()
280 }
281 Expr::Binding(b) => {
282 rewrite_expr(&mut b.value, step);
283 if let Some(block) = &mut b.r#otherwise {
284 rewrite_block(block, step);
285 }
286 ()
287 }
288 Expr::Item { item, .. } => rewrite_item(item, step),
289 Expr::Return { value, .. } => {
290 if let Some(v) = value {
291 rewrite_expr(v, step);
292 }
293 ()
294 }
295 Expr::Draw { value, .. } => rewrite_expr(value, step),
296 Expr::Await { value, .. } => rewrite_expr(value, step),
297 Expr::Assign { target, value, .. } => {
298 rewrite_expr(target, step);
299 rewrite_expr(value, step)
300 }
301 Expr::Guard {
302 source,
303 r#otherwise,
304 ..
305 } => {
306 rewrite_expr(source, step);
307 rewrite_block(r#otherwise, step)
308 }
309 Expr::While {
310 condition, body, ..
311 } => {
312 rewrite_condition(condition, step);
313 rewrite_block(body, step)
314 }
315 Expr::For { source, body, .. } => {
316 rewrite_expr(source, step);
317 rewrite_block(body, step)
318 }
319 Expr::Loop { body, .. } => rewrite_block(body, step),
320 Expr::Break { value, .. } => {
321 if let Some(v) = value {
322 rewrite_expr(v, step);
323 }
324 ()
325 }
326 Expr::MacroCall { args, .. } => {
327 for a in args.iter_mut() {
328 rewrite_expr(a, step);
329 }
330 ()
331 }
332 Expr::Path { .. }
333 | Expr::MethodRef { .. }
334 | Expr::StrLit { .. }
335 | Expr::RawLit { .. }
336 | Expr::Unit { .. }
337 | Expr::Num { .. }
338 | Expr::Bool { .. }
339 | Expr::Char { .. }
340 | Expr::SelfField { .. }
341 | Expr::SelfType { .. }
342 | Expr::Ident(_)
343 | Expr::Continue { .. }
344 | Expr::RustHatch { .. } => (),
345 };
346}
347
348pub fn apply_pass(e: &mut Expr, pass: Pass) {
349 match pass {
350 Pass::BlankNumbers => blank_number(e),
351 Pass::Pipelines => collapse_pipeline(e),
352 };
353}
354
355pub(crate) fn blank_number(e: &mut Expr) {
356 if let Expr::Num { text, .. } = e {
357 *text = "9".to_string();
358 }
359 ();
360}
361
362pub(crate) fn collapse_pipeline(e: &mut Expr) {
363 let Expr::Binary { op, lhs, rhs, span } = e else {
364 return ();
365 };
366 if *op != BinaryOp::Pipeline {
367 return ();
368 }
369 let mut named = None;
370 let mut found: Vec<Expr> = Vec::new();
371 if let Expr::MethodRef { name, .. } = rhs.as_ref() {
372 named = Some(name.clone());
373 }
374 if let Expr::Call { callee, args, .. } = rhs.as_ref()
375 && let Expr::MethodRef { name, .. } = callee.as_ref()
376 {
377 named = Some(name.clone());
378 found = args.as_slice().to_vec();
379 }
380 let Some(name) = named else {
381 return ();
382 };
383 *e = Expr::Dot {
384 recv: lhs.clone(),
385 name,
386 args: found,
387 span: *span,
388 };
389 ();
390}
391
392#[cfg(test)]
393mod __vlt_tests {
394 use super::*;
395
396 #[test]
397 fn a_leaf_has_no_children_so_descending_it_is_a_no_op() {
398 let mut e = Expr::Bool {
399 value: true,
400 span: (0, 1),
401 };
402 let before = e.clone();
403 rewrite_expr(&mut e, Pass::BlankNumbers);
404 assert_eq!(e, before);
405 }
406
407 #[test]
408 fn the_traversal_reaches_a_numeral_in_every_expression_position() {
409 let source = "f : u64 -> u64\nf n =\n m = n + 1\n xs = [2, 3] : Vec<u64>\n case m of\n 4 -> 5\n _ -> m * 6\n";
410 let mut program = voltaic_parser::parse_source(source).expect("should parse");
411 rewrite_program(&mut program, Pass::BlankNumbers);
412 let shown = format!("{:?}", program);
413 for digit in &["\"1\"", "\"2\"", "\"3\"", "\"5\"", "\"6\""] {
414 assert!(!shown.as_str().contains(*digit));
415 }
416 }
417
418 #[test]
419 fn a_numeral_in_a_pattern_is_left_alone_because_a_pass_must_not_change_what_matches() {
420 let source = "f : u64 -> u64\nf n = case n of\n 4 -> 5\n _ -> 6\n";
421 let mut program = voltaic_parser::parse_source(source).expect("should parse");
422 rewrite_program(&mut program, Pass::BlankNumbers);
423 let shown = format!("{:?}", program);
424 assert!(shown.as_str().contains("\"4\""));
425 assert!(!shown.as_str().contains("\"5\""));
426 assert!(!shown.as_str().contains("\"6\""));
427 }
428
429 #[test]
430 fn a_numeral_the_traversal_cannot_reach_would_survive_which_is_what_this_pins() {
431 let source = "f : u64 -> u64\nf n = n + 7\n";
432 let mut program = voltaic_parser::parse_source(source).expect("should parse");
433 let before = format!("{:?}", program);
434 assert!(before.as_str().contains("\"7\""));
435 rewrite_program(&mut program, Pass::BlankNumbers);
436 let after = format!("{:?}", program);
437 assert!(!after.as_str().contains("\"7\""));
438 }
439
440 #[test]
441 fn a_pipeline_with_no_arguments_collapses_now_that_the_printer_writes_the_parentheses() {
442 let source = "f : &str -> usize\nf s = s |> .len\n";
443 let mut program = voltaic_parser::parse_source(source).expect("should parse");
444 rewrite_program(&mut program, Pass::Pipelines);
445 let shown = format!("{:?}", program);
446 assert!(!shown.as_str().contains("Pipeline"));
447 assert!(shown.as_str().contains("Dot"));
448 }
449
450 #[test]
451 fn a_pipeline_with_arguments_keeps_them_receiver_first() {
452 let source = "f : u64 -> u64\nf n = n |> .max 3\n";
453 let mut program = voltaic_parser::parse_source(source).expect("should parse");
454 rewrite_program(&mut program, Pass::Pipelines);
455 let shown = format!("{:?}", program);
456 assert!(shown.as_str().contains("Dot"));
457 assert!(shown.as_str().contains("\"max\""));
458 assert!(shown.as_str().contains("\"3\""));
459 }
460
461 #[test]
462 fn a_chain_collapses_every_stage_arguments_or_not() {
463 let source = "f : u64 -> u64\nf n = n |> .max 3 |> .min 8\n";
464 let mut program = voltaic_parser::parse_source(source).expect("should parse");
465 rewrite_program(&mut program, Pass::Pipelines);
466 let shown = format!("{:?}", program);
467 assert!(!shown.as_str().contains("Pipeline"));
468 }
469
470 #[test]
471 fn a_mixed_chain_collapses_both_stages_the_one_with_arguments_and_the_one_without() {
472 let source = "f : &str -> usize\nf s = s |> .replacen \"a\" \"b\" 1 |> .len\n";
473 let mut program = voltaic_parser::parse_source(source).expect("should parse");
474 rewrite_program(&mut program, Pass::Pipelines);
475 let shown = format!("{:?}", program);
476 assert!(shown.as_str().contains("Dot"));
477 assert!(!shown.as_str().contains("Pipeline"));
478 }
479
480 #[test]
481 fn a_binary_that_is_not_a_pipeline_is_left_alone() {
482 let source = "f : u64 -> u64\nf n = n + 1\n";
483 let mut program = voltaic_parser::parse_source(source).expect("should parse");
484 let before = format!("{:?}", program);
485 rewrite_program(&mut program, Pass::Pipelines);
486 let after = format!("{:?}", program);
487 assert_eq!(after, before);
488 }
489}