1use std::io::Write;
15
16use crate::ast::*;
17use crate::parse::CommentMap;
18
19pub struct AstPrinter<'a, W>
20where
21 W: Write,
22{
23 indent_size: usize,
24 curr_indent: usize,
25 w: W,
26 comment_map: Option<&'a CommentMap>,
29 last_line: usize,
30 comment_group_lines: Vec<usize>,
31}
32
33impl<'a, W> AstPrinter<'a, W>
36where
37 W: Write,
38{
39 pub fn new(indent: usize, w: W) -> Self {
40 AstPrinter {
41 indent_size: indent,
42 curr_indent: 0,
43 comment_map: None,
44 w: w,
45 last_line: 0,
46 comment_group_lines: Vec::new(),
47 }
48 }
49
50 pub fn with_comment_map(mut self, map: &'a CommentMap) -> Self {
51 self.comment_group_lines = map.keys().cloned().collect();
52 self.comment_group_lines.reverse();
53 self.comment_map = Some(map);
54 self
55 }
56
57 fn make_indent(&self) -> String {
58 let indent: Vec<u8> = std::iter::repeat(' ' as u8)
61 .take(self.curr_indent)
62 .collect();
63 String::from_utf8_lossy(&indent).to_string()
64 }
65
66 fn is_bareword(s: &str) -> bool {
67 match s.chars().nth(0) {
68 Some(c) => {
69 if !(c.is_ascii_alphabetic() || c == '_') {
70 return false;
71 }
72 }
73 None => return false,
74 };
75 for c in s.chars() {
76 if !(c.is_ascii_alphabetic() || c == '_') {
77 return false;
78 }
79 }
80 return true;
81 }
82
83 fn print_comment_group(&mut self, line: usize) -> std::io::Result<()> {
84 if let Some(ref map) = self.comment_map {
85 let empty: Vec<Token> = Vec::new();
86 let cg = map.get(&line).unwrap_or(&empty);
87 let indent = self.make_indent();
88 for c in cg.iter() {
89 let first_char = match c.fragment.chars().nth(0) {
90 Some(c) => c,
91 None => '\0',
92 };
93 if !first_char.is_whitespace() {
94 write!(self.w, "{}// {}\n", indent, c.fragment.trim_end())?;
95 } else {
96 write!(self.w, "{}//{}\n", indent, c.fragment.trim_end())?;
97 }
98 }
99 self.comment_group_lines.pop();
100 }
101 Ok(())
102 }
103
104 fn render_missed_comments(&mut self, line: usize) -> std::io::Result<()> {
105 loop {
106 if let Some(next_comment_line) = self.comment_group_lines.last() {
107 let next_comment_line = *next_comment_line;
108 if next_comment_line <= line {
109 self.print_comment_group(next_comment_line)?;
110 } else {
111 break;
112 }
113 if next_comment_line < line - 1 {
114 write!(self.w, "\n")?;
115 }
116 continue;
117 }
118 break;
119 }
120 Ok(())
121 }
122
123 fn render_comment_if_needed(&mut self, line: usize) -> std::io::Result<()> {
124 self.render_missed_comments(line)?;
125 self.last_line = line;
126 Ok(())
127 }
128
129 fn has_comment(&self, line: usize) -> bool {
130 if let Some(next_comment_line) = self.comment_group_lines.last() {
131 return *next_comment_line < line;
132 }
133 false
134 }
135
136 fn render_list_def(&mut self, def: &ListDef) -> std::io::Result<()> {
137 write!(self.w, "[")?;
138 self.curr_indent += self.indent_size;
139 let indent = self.make_indent();
140 let has_fields = def.elems.len() > 0;
141 if has_fields {
142 write!(self.w, "\n")?;
143 }
144 for e in def.elems.iter() {
145 self.render_comment_if_needed(e.pos().line)?;
146 write!(self.w, "{}", indent)?;
147 self.render_expr(e)?;
148 write!(self.w, ",\n")?;
149 }
150 self.curr_indent -= self.indent_size;
151 if has_fields {
152 write!(self.w, "{}", self.make_indent())?;
153 }
154 self.w.write(&[']' as u8])?;
155 Ok(())
156 }
157
158 fn render_tuple_def(
159 &mut self,
160 def: &Vec<(Token, Option<Expression>, Expression)>,
161 ) -> std::io::Result<()> {
162 self.w.write(&['{' as u8])?;
163 self.curr_indent += self.indent_size;
165 let indent = self.make_indent();
166 let has_fields = def.len() > 0;
167 if has_fields {
168 write!(self.w, "\n")?;
169 }
170 for (ref t, ref constraint, ref expr) in def.iter() {
171 let field_line = t.pos.line;
172 let expr_line = expr.pos().line;
173 self.render_comment_if_needed(field_line)?;
174 if expr_line != field_line {
175 self.render_comment_if_needed(expr_line)?;
176 }
177 write!(self.w, "{}", indent)?;
178 if Self::is_bareword(&t.fragment) {
179 write!(&mut self.w, "{}", t.fragment)?;
180 } else {
181 write!(self.w, "\"{}\"", Self::escape_quotes(&t.fragment))?;
182 }
183 if let Some(c) = constraint {
184 write!(self.w, " :: ")?;
185 self.render_expr(c)?;
186 }
187 write!(self.w, " = ")?;
188 self.render_expr(expr)?;
189 write!(&mut self.w, ",")?;
190 write!(self.w, "\n")?;
191 }
192 self.curr_indent -= self.indent_size;
193 if has_fields {
194 write!(self.w, "{}", self.make_indent())?;
195 }
196 self.w.write(&['}' as u8])?;
197 Ok(())
198 }
199
200 fn escape_quotes(s: &str) -> String {
201 let mut escaped = String::new();
202 for c in s.chars() {
203 if c == '"' {
204 escaped.push_str("\\\"");
205 } else if c == '\\' {
206 escaped.push_str("\\\\");
207 } else {
208 escaped.push(c);
209 }
210 }
211 escaped
212 }
213
214 pub fn render_value(&mut self, v: &Value) -> std::io::Result<()> {
215 match v {
216 Value::Boolean(b) => write!(self.w, "{}", if b.val { "true" } else { "false" })?,
217 Value::Empty(_) => write!(self.w, "NULL")?,
218 Value::Float(f) => write!(self.w, "{}", f.val)?,
220 Value::Int(i) => write!(self.w, "{}", i.val)?,
221 Value::Str(s) => write!(self.w, "\"{}\"", Self::escape_quotes(&s.val))?,
222 Value::Symbol(s) => write!(self.w, "{}", s.val)?,
223 Value::List(l) => self.render_list_def(l)?,
224 Value::Tuple(tpl) => self.render_tuple_def(&tpl.val)?,
225 };
226 Ok(())
227 }
228
229 pub fn render_expr(&mut self, expr: &Expression) -> std::io::Result<()> {
230 let had_comment = self.has_comment(expr.pos().line);
231 self.render_comment_if_needed(expr.pos().line)?;
232 let indent = self.make_indent();
233 if had_comment {
234 write!(self.w, "{}", indent)?;
235 }
236 let mut did_indent = false;
237 match expr {
238 Expression::Binary(_def) => {
239 let op = match _def.kind {
240 BinaryExprType::AND => " && ",
241 BinaryExprType::OR => " || ",
242 BinaryExprType::DOT => ".",
243 BinaryExprType::Equal => " == ",
244 BinaryExprType::NotEqual => " != ",
245 BinaryExprType::GTEqual => " >= ",
246 BinaryExprType::LTEqual => " <= ",
247 BinaryExprType::GT => " > ",
248 BinaryExprType::LT => " < ",
249 BinaryExprType::Add => " + ",
250 BinaryExprType::Sub => " - ",
251 BinaryExprType::Mul => " * ",
252 BinaryExprType::Div => " / ",
253 BinaryExprType::Mod => " %% ",
254 BinaryExprType::IN => " in ",
255 BinaryExprType::IS => " is ",
256 BinaryExprType::REMatch => " ~ ",
257 BinaryExprType::NotREMatch => " !~ ",
258 };
259 let right_line = _def.right.pos().line;
260 self.render_expr(&_def.left)?;
261 self.w.write(op.as_bytes())?;
262 if self.has_comment(right_line) {
263 self.w.write("\n".as_bytes())?;
266 }
267 self.render_expr(&_def.right)?;
268 }
269 Expression::Cast(def) => {
270 self.w.write(format!("{}", def.cast_type).as_bytes())?;
271 self.w.write("(".as_bytes())?;
272 self.render_comment_if_needed(def.target.pos().line)?;
273 self.render_expr(&def.target)?;
274 self.w.write(")".as_bytes())?;
275 }
276 Expression::Call(_def) => {
277 self.render_value(&_def.funcref)?;
278 self.w.write("(".as_bytes())?;
279 self.curr_indent += self.indent_size;
280 let indent = self.make_indent();
281 let has_args = _def.arglist.len() > 1;
282 if has_args {
283 write!(self.w, "\n")?;
284 }
285 for e in _def.arglist.iter() {
286 self.render_comment_if_needed(e.pos().line)?;
287 if has_args {
288 write!(self.w, "{}", indent)?;
289 }
290 self.render_expr(e)?;
291 if has_args {
292 self.w.write(",\n".as_bytes())?;
293 }
294 }
295 self.curr_indent -= self.indent_size;
296 if has_args {
297 write!(self.w, "{}", self.make_indent())?;
298 }
299 self.w.write(")".as_bytes())?;
300 }
301 Expression::Copy(_def) => {
302 self.render_value(&_def.selector)?;
303 self.render_tuple_def(&_def.fields)?;
304 }
305 Expression::Debug(_def) => {
306 self.w.write("TRACE ".as_bytes())?;
307 if self.has_comment(_def.expr.pos().line) {
308 self.w.write("\n".as_bytes())?;
309 }
310 self.render_expr(&_def.expr)?;
311 }
312 Expression::Fail(_def) => {
313 self.w.write("fail ".as_bytes())?;
314 if self.has_comment(_def.message.pos().line) {
315 self.w.write("\n".as_bytes())?;
316 }
317 self.render_expr(&_def.message)?;
318 }
319 Expression::Format(_def) => {
320 write!(self.w, "\"{}\"", Self::escape_quotes(&_def.template))?;
321 write!(self.w, " % ")?;
322 match _def.args {
323 FormatArgs::Single(ref e) => {
324 if self.has_comment(e.pos().line) {
325 self.w.write("\n".as_bytes())?;
326 }
327 self.render_expr(e)?;
328 }
329 FormatArgs::List(ref es) => {
330 self.w.write("(\n".as_bytes())?;
331 self.curr_indent += self.indent_size;
332 let indent = self.make_indent();
333 let mut prefix = if es
334 .first()
335 .and_then(|e| Some(self.has_comment(e.pos().line)))
336 .unwrap_or(false)
337 {
338 "\n"
339 } else {
340 ""
341 };
342 for e in es.iter() {
343 write!(self.w, "{}{}", prefix, indent)?;
344 self.render_expr(e)?;
345 prefix = ",\n";
346 }
347 self.curr_indent -= self.indent_size;
348 self.w.write(")".as_bytes())?;
349 }
350 }
351 }
352 Expression::Func(_def) => {
353 self.w.write("func (".as_bytes())?;
354 if _def.argdefs.len() == 1 {
355 let (ref arg, ref constraint) = _def.argdefs.first().unwrap();
356 write!(self.w, "{}", arg)?;
357 if let Some(c) = constraint {
358 write!(self.w, " :: ")?;
359 self.render_expr(c)?;
360 }
361 } else {
362 let mut prefix = "";
363 for (ref n, ref constraint) in _def.argdefs.iter() {
364 write!(self.w, "{}{}", prefix, n.val)?;
365 if let Some(c) = constraint {
366 write!(self.w, " :: ")?;
367 self.render_expr(c)?;
368 }
369 prefix = ", ";
370 }
371 }
372 self.w.write(") => ".as_bytes())?;
373 self.render_expr(&_def.fields)?;
374 }
375 Expression::FuncOp(_def) => match _def {
376 FuncOpDef::Filter(_def) => {
377 write!(self.w, "filter(")?;
378 if self.has_comment(_def.func.pos().line) {
379 self.curr_indent += self.indent_size;
380 did_indent = true;
381 write!(self.w, "\n")?;
382 }
383 self.render_expr(&_def.func)?;
384 if self.has_comment(_def.target.pos().line) {
385 write!(self.w, ",")?;
386 if !did_indent {
387 self.curr_indent += self.indent_size;
388 }
389 did_indent = true;
390 self.w.write("\n".as_bytes())?;
391 } else {
392 write!(self.w, ", ")?;
393 }
394 self.render_expr(&_def.target)?;
395 write!(self.w, ")")?;
396 }
397 FuncOpDef::Reduce(_def) => {
398 write!(self.w, "reduce(")?;
399 if self.has_comment(_def.func.pos().line) {
400 self.curr_indent += self.indent_size;
401 did_indent = true;
402 write!(self.w, "\n")?;
403 }
404 self.render_expr(&_def.func)?;
405 if self.has_comment(_def.acc.pos().line) {
406 write!(self.w, ",")?;
407 if !did_indent {
408 self.curr_indent += self.indent_size;
409 }
410 did_indent = true;
411 self.w.write("\n".as_bytes())?;
412 } else {
413 write!(self.w, ", ")?;
414 }
415 self.render_expr(&_def.acc)?;
416 if self.has_comment(_def.target.pos().line) {
417 write!(self.w, ",")?;
418 if !did_indent {
419 self.curr_indent += self.indent_size;
420 }
421 did_indent = true;
422 self.w.write("\n".as_bytes())?;
423 } else {
424 write!(self.w, ", ")?;
425 }
426 self.render_expr(&_def.target)?;
427 write!(self.w, ")")?;
428 }
429 FuncOpDef::Map(_def) => {
430 write!(self.w, "map(")?;
431 if self.has_comment(_def.func.pos().line) {
432 self.curr_indent += self.indent_size;
433 did_indent = true;
434 write!(self.w, "\n")?;
435 }
436 self.render_expr(&_def.func)?;
437 if self.has_comment(_def.target.pos().line) {
438 write!(self.w, ",")?;
439 if !did_indent {
440 self.curr_indent += self.indent_size;
441 }
442 did_indent = true;
443 self.w.write("\n".as_bytes())?;
444 } else {
445 write!(self.w, ", ")?;
446 }
447 self.render_expr(&_def.target)?;
448 write!(self.w, ")")?;
449 }
450 },
451 Expression::Grouped(ref expr, _) => {
452 write!(self.w, "(")?;
453 if self.has_comment(expr.pos().line) {
454 self.curr_indent += self.indent_size;
455 did_indent = true;
456 write!(self.w, "\n")?;
457 }
458 self.render_expr(expr)?;
459 if did_indent {
460 write!(self.w, "\n")?;
461 }
462 write!(self.w, ")")?;
463 }
464 Expression::Import(_def) => {
465 if self.has_comment(_def.path.pos.line) {
466 self.render_missed_comments(_def.path.pos.line)?;
467 }
468 write!(
469 self.w,
470 "import \"{}\"",
471 Self::escape_quotes(&_def.path.fragment)
472 )?;
473 }
474 Expression::Include(_def) => {
475 if self.has_comment(_def.path.pos.line) {
476 self.render_missed_comments(_def.path.pos.line)?;
477 }
478 write!(
479 self.w,
480 "include {} \"{}\"",
481 _def.typ.fragment,
482 Self::escape_quotes(&_def.path.fragment)
483 )?;
484 }
485 Expression::Module(_def) => {
486 write!(self.w, "module ")?;
487 self.render_tuple_def(&_def.arg_set)?;
488 write!(self.w, " => ")?;
489 if let Some(ref e) = _def.out_expr {
490 write!(self.w, "(")?;
491 self.render_expr(e)?;
492 if let Some(ref c) = _def.out_constraint {
493 write!(self.w, " :: ")?;
494 self.render_expr(c)?;
495 }
496 write!(self.w, ") ")?;
497 }
498 write!(self.w, "{{\n")?;
499 self.curr_indent += self.indent_size;
500 let indent = self.make_indent();
501 let mut prefix_newline = false;
502 for stmt in _def.statements.iter() {
503 write!(self.w, "{}", indent)?;
504 self.render_stmt(stmt, prefix_newline)?;
505 prefix_newline = true;
506 }
507 self.curr_indent -= self.indent_size;
508 write!(self.w, "}}")?;
509 }
510 Expression::Not(_def) => {
511 if self.has_comment(_def.pos.line) {
512 self.render_missed_comments(_def.pos.line)?;
513 }
514 write!(self.w, "not ")?;
515 self.render_expr(&_def.expr)?;
516 }
517 Expression::Range(_def) => {
518 let end_line = _def.end.pos().line;
521 if self.has_comment(end_line) {
522 self.render_missed_comments(end_line)?;
523 }
524 self.render_expr(&_def.start)?;
525 write!(self.w, ":")?;
526 if let Some(ref e) = _def.step {
527 write!(self.w, ":")?;
528 self.render_expr(e)?;
529 }
530 self.render_expr(&_def.end)?;
531 }
532 Expression::Select(_def) => {
533 let val_line = _def.val.pos().line;
534 if self.has_comment(val_line) {
535 self.render_missed_comments(val_line)?;
536 }
537 if let Some(ref e) = _def.default {
538 let default_line = e.pos().line;
539 if self.has_comment(default_line) {
540 self.render_missed_comments(default_line)?;
541 }
542 }
543 write!(self.w, "select ")?;
544 write!(self.w, "(")?;
545 self.render_expr(&_def.val)?;
546 if let Some(ref e) = _def.default {
547 write!(self.w, ", ")?;
548 self.render_expr(e)?;
549 }
550 write!(self.w, ") => ")?;
551 self.render_tuple_def(&_def.tuple)?;
552 }
553 Expression::Simple(ref _def) => {
554 self.render_value(_def)?;
555 }
556 };
557 if did_indent {
558 self.curr_indent -= self.indent_size;
559 }
560 Ok(())
561 }
562
563 pub fn render_stmt(&mut self, stmt: &Statement, prefix_newline: bool) -> std::io::Result<()> {
564 if prefix_newline {
566 write!(self.w, "\n")?;
567 }
568 let line = stmt.pos().line;
569 self.render_comment_if_needed(line)?;
570 match stmt {
571 Statement::Let(def) => {
572 write!(&mut self.w, "let {}", def.name.fragment)?;
573 if let Some(ref c) = def.constraint {
574 write!(self.w, " :: ")?;
575 self.render_expr(c)?;
576 }
577 write!(self.w, " = ")?;
578 self.render_expr(&def.value)?;
579 }
580 Statement::Expression(_expr) => {
581 self.render_expr(&_expr)?;
582 }
583 Statement::Assert(_, def) => {
584 write!(&mut self.w, "assert ")?;
585 self.render_expr(&def)?;
586 }
587 Statement::Output(_, _tok, _expr) => {
588 write!(&mut self.w, "out {} ", _tok.fragment)?;
589 self.render_expr(&_expr)?;
590 }
591 Statement::Print(_, _tok, _expr) => {
592 write!(&mut self.w, "print {} ", _tok.fragment)?;
593 self.render_expr(&_expr)?;
594 }
595 };
596 write!(self.w, ";\n")?;
597 self.last_line = line;
598 Ok(())
599 }
600
601 pub fn render(&mut self, stmts: &Vec<Statement>) -> std::io::Result<()> {
602 let mut prefix_newline = false;
603 for v in stmts {
604 self.render_stmt(v, prefix_newline)?;
605 prefix_newline = true;
606 }
607 let comment_line = self.comment_group_lines.first().cloned();
608 if let Some(last_comment_line) = comment_line {
609 self.render_missed_comments(last_comment_line + 1)?;
610 }
611 Ok(())
612 }
613}
614
615#[cfg(test)]
616mod test;