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(&mut self, def: &Vec<(Token, Expression)>) -> std::io::Result<()> {
159 self.w.write(&['{' as u8])?;
160 self.curr_indent += self.indent_size;
162 let indent = self.make_indent();
163 let has_fields = def.len() > 0;
164 if has_fields {
165 write!(self.w, "\n")?;
166 }
167 for &(ref t, ref expr) in def.iter() {
168 let field_line = t.pos.line;
169 let expr_line = expr.pos().line;
170 self.render_comment_if_needed(field_line)?;
171 if expr_line != field_line {
172 self.render_comment_if_needed(expr_line)?;
173 }
174 write!(self.w, "{}", indent)?;
175 if Self::is_bareword(&t.fragment) {
176 write!(&mut self.w, "{} = ", t.fragment)?;
177 } else {
178 write!(self.w, "\"{}\" = ", Self::escape_quotes(&t.fragment))?;
179 }
180 self.render_expr(expr)?;
181 write!(&mut self.w, ",")?;
182 write!(self.w, "\n")?;
183 }
184 self.curr_indent -= self.indent_size;
185 if has_fields {
186 write!(self.w, "{}", self.make_indent())?;
187 }
188 self.w.write(&['}' as u8])?;
189 Ok(())
190 }
191
192 fn escape_quotes(s: &str) -> String {
193 let mut escaped = String::new();
194 for c in s.chars() {
195 if c == '"' {
196 escaped.push_str("\\\"");
197 } else if c == '\\' {
198 escaped.push_str("\\\\");
199 } else {
200 escaped.push(c);
201 }
202 }
203 escaped
204 }
205
206 pub fn render_value(&mut self, v: &Value) -> std::io::Result<()> {
207 match v {
208 Value::Boolean(b) => write!(self.w, "{}", if b.val { "true" } else { "false" })?,
209 Value::Empty(_) => write!(self.w, "NULL")?,
210 Value::Float(f) => write!(self.w, "{}", f.val)?,
212 Value::Int(i) => write!(self.w, "{}", i.val)?,
213 Value::Str(s) => write!(self.w, "\"{}\"", Self::escape_quotes(&s.val))?,
214 Value::Symbol(s) => write!(self.w, "{}", s.val)?,
215 Value::List(l) => self.render_list_def(l)?,
216 Value::Tuple(tpl) => self.render_tuple_def(&tpl.val)?,
217 };
218 Ok(())
219 }
220
221 pub fn render_expr(&mut self, expr: &Expression) -> std::io::Result<()> {
222 let had_comment = self.has_comment(expr.pos().line);
223 self.render_comment_if_needed(expr.pos().line)?;
224 let indent = self.make_indent();
225 if had_comment {
226 write!(self.w, "{}", indent)?;
227 }
228 let mut did_indent = false;
229 match expr {
230 Expression::Binary(_def) => {
231 let op = match _def.kind {
232 BinaryExprType::AND => " && ",
233 BinaryExprType::OR => " || ",
234 BinaryExprType::DOT => ".",
235 BinaryExprType::Equal => " == ",
236 BinaryExprType::NotEqual => " != ",
237 BinaryExprType::GTEqual => " >= ",
238 BinaryExprType::LTEqual => " <= ",
239 BinaryExprType::GT => " > ",
240 BinaryExprType::LT => " < ",
241 BinaryExprType::Add => " + ",
242 BinaryExprType::Sub => " - ",
243 BinaryExprType::Mul => " * ",
244 BinaryExprType::Div => " / ",
245 BinaryExprType::Mod => " %% ",
246 BinaryExprType::IN => " in ",
247 BinaryExprType::IS => " is ",
248 BinaryExprType::REMatch => " ~ ",
249 BinaryExprType::NotREMatch => " !~ ",
250 };
251 let right_line = _def.right.pos().line;
252 self.render_expr(&_def.left)?;
253 self.w.write(op.as_bytes())?;
254 if self.has_comment(right_line) {
255 self.w.write("\n".as_bytes())?;
258 }
259 self.render_expr(&_def.right)?;
260 }
261 Expression::Cast(def) => {
262 self.w.write(format!("{}", def.cast_type).as_bytes())?;
263 self.w.write("(".as_bytes())?;
264 self.render_comment_if_needed(def.target.pos().line)?;
265 self.render_expr(&def.target)?;
266 self.w.write(")".as_bytes())?;
267 }
268 Expression::Call(_def) => {
269 self.render_value(&_def.funcref)?;
270 self.w.write("(".as_bytes())?;
271 self.curr_indent += self.indent_size;
272 let indent = self.make_indent();
273 let has_args = _def.arglist.len() > 1;
274 if has_args {
275 write!(self.w, "\n")?;
276 }
277 for e in _def.arglist.iter() {
278 self.render_comment_if_needed(e.pos().line)?;
279 if has_args {
280 write!(self.w, "{}", indent)?;
281 }
282 self.render_expr(e)?;
283 if has_args {
284 self.w.write(",\n".as_bytes())?;
285 }
286 }
287 self.curr_indent -= self.indent_size;
288 if has_args {
289 write!(self.w, "{}", self.make_indent())?;
290 }
291 self.w.write(")".as_bytes())?;
292 }
293 Expression::Copy(_def) => {
294 self.render_value(&_def.selector)?;
295 self.render_tuple_def(&_def.fields)?;
296 }
297 Expression::Debug(_def) => {
298 self.w.write("TRACE ".as_bytes())?;
299 if self.has_comment(_def.expr.pos().line) {
300 self.w.write("\n".as_bytes())?;
301 }
302 self.render_expr(&_def.expr)?;
303 }
304 Expression::Fail(_def) => {
305 self.w.write("fail ".as_bytes())?;
306 if self.has_comment(_def.message.pos().line) {
307 self.w.write("\n".as_bytes())?;
308 }
309 self.render_expr(&_def.message)?;
310 }
311 Expression::Format(_def) => {
312 write!(self.w, "\"{}\"", Self::escape_quotes(&_def.template))?;
313 write!(self.w, " % ")?;
314 match _def.args {
315 FormatArgs::Single(ref e) => {
316 if self.has_comment(e.pos().line) {
317 self.w.write("\n".as_bytes())?;
318 }
319 self.render_expr(e)?;
320 }
321 FormatArgs::List(ref es) => {
322 self.w.write("(\n".as_bytes())?;
323 self.curr_indent += self.indent_size;
324 let indent = self.make_indent();
325 let mut prefix = if es
326 .first()
327 .and_then(|e| Some(self.has_comment(e.pos().line)))
328 .unwrap_or(false)
329 {
330 "\n"
331 } else {
332 ""
333 };
334 for e in es.iter() {
335 write!(self.w, "{}{}", prefix, indent)?;
336 self.render_expr(e)?;
337 prefix = ",\n";
338 }
339 self.curr_indent -= self.indent_size;
340 self.w.write(")".as_bytes())?;
341 }
342 }
343 }
344 Expression::Func(_def) => {
345 self.w.write("func (".as_bytes())?;
346 if _def.argdefs.len() == 1 {
347 write!(self.w, "{}", _def.argdefs.first().unwrap())?;
348 } else {
349 let mut prefix = "";
350 for n in _def.argdefs.iter() {
351 write!(self.w, "{}{}", prefix, n.val)?;
352 prefix = ", ";
353 }
354 }
355 self.w.write(") => ".as_bytes())?;
356 self.render_expr(&_def.fields)?;
357 }
358 Expression::FuncOp(_def) => match _def {
359 FuncOpDef::Filter(_def) => {
360 write!(self.w, "filter(")?;
361 if self.has_comment(_def.func.pos().line) {
362 self.curr_indent += self.indent_size;
363 did_indent = true;
364 write!(self.w, "\n")?;
365 }
366 self.render_expr(&_def.func)?;
367 if self.has_comment(_def.target.pos().line) {
368 write!(self.w, ",")?;
369 if !did_indent {
370 self.curr_indent += self.indent_size;
371 }
372 did_indent = true;
373 self.w.write("\n".as_bytes())?;
374 } else {
375 write!(self.w, ", ")?;
376 }
377 self.render_expr(&_def.target)?;
378 write!(self.w, ")")?;
379 }
380 FuncOpDef::Reduce(_def) => {
381 write!(self.w, "reduce(")?;
382 if self.has_comment(_def.func.pos().line) {
383 self.curr_indent += self.indent_size;
384 did_indent = true;
385 write!(self.w, "\n")?;
386 }
387 self.render_expr(&_def.func)?;
388 if self.has_comment(_def.acc.pos().line) {
389 write!(self.w, ",")?;
390 if !did_indent {
391 self.curr_indent += self.indent_size;
392 }
393 did_indent = true;
394 self.w.write("\n".as_bytes())?;
395 } else {
396 write!(self.w, ", ")?;
397 }
398 self.render_expr(&_def.acc)?;
399 if self.has_comment(_def.target.pos().line) {
400 write!(self.w, ",")?;
401 if !did_indent {
402 self.curr_indent += self.indent_size;
403 }
404 did_indent = true;
405 self.w.write("\n".as_bytes())?;
406 } else {
407 write!(self.w, ", ")?;
408 }
409 self.render_expr(&_def.target)?;
410 write!(self.w, ")")?;
411 }
412 FuncOpDef::Map(_def) => {
413 write!(self.w, "map(")?;
414 if self.has_comment(_def.func.pos().line) {
415 self.curr_indent += self.indent_size;
416 did_indent = true;
417 write!(self.w, "\n")?;
418 }
419 self.render_expr(&_def.func)?;
420 if self.has_comment(_def.target.pos().line) {
421 write!(self.w, ",")?;
422 if !did_indent {
423 self.curr_indent += self.indent_size;
424 }
425 did_indent = true;
426 self.w.write("\n".as_bytes())?;
427 } else {
428 write!(self.w, ", ")?;
429 }
430 self.render_expr(&_def.target)?;
431 write!(self.w, ")")?;
432 }
433 },
434 Expression::Grouped(ref expr, _) => {
435 write!(self.w, "(")?;
436 if self.has_comment(expr.pos().line) {
437 self.curr_indent += self.indent_size;
438 did_indent = true;
439 write!(self.w, "\n")?;
440 }
441 self.render_expr(expr)?;
442 if did_indent {
443 write!(self.w, "\n")?;
444 }
445 write!(self.w, ")")?;
446 }
447 Expression::Import(_def) => {
448 if self.has_comment(_def.path.pos.line) {
449 self.render_missed_comments(_def.path.pos.line)?;
450 }
451 write!(
452 self.w,
453 "import \"{}\"",
454 Self::escape_quotes(&_def.path.fragment)
455 )?;
456 }
457 Expression::Include(_def) => {
458 if self.has_comment(_def.path.pos.line) {
459 self.render_missed_comments(_def.path.pos.line)?;
460 }
461 write!(
462 self.w,
463 "include {} \"{}\"",
464 _def.typ.fragment,
465 Self::escape_quotes(&_def.path.fragment)
466 )?;
467 }
468 Expression::Module(_def) => {
469 write!(self.w, "module ")?;
470 self.render_tuple_def(&_def.arg_set)?;
471 write!(self.w, " => ")?;
472 if let Some(ref e) = _def.out_expr {
473 write!(self.w, "(")?;
474 self.render_expr(e)?;
475 write!(self.w, ") ")?;
476 }
477 write!(self.w, "{{\n")?;
478 self.curr_indent += self.indent_size;
479 let indent = self.make_indent();
480 let mut prefix_newline = false;
481 for stmt in _def.statements.iter() {
482 write!(self.w, "{}", indent)?;
483 self.render_stmt(stmt, prefix_newline)?;
484 prefix_newline = true;
485 }
486 self.curr_indent -= self.indent_size;
487 write!(self.w, "}}")?;
488 }
489 Expression::Not(_def) => {
490 if self.has_comment(_def.pos.line) {
491 self.render_missed_comments(_def.pos.line)?;
492 }
493 write!(self.w, "not ")?;
494 self.render_expr(&_def.expr)?;
495 }
496 Expression::Range(_def) => {
497 let end_line = _def.end.pos().line;
500 if self.has_comment(end_line) {
501 self.render_missed_comments(end_line)?;
502 }
503 self.render_expr(&_def.start)?;
504 write!(self.w, ":")?;
505 if let Some(ref e) = _def.step {
506 write!(self.w, ":")?;
507 self.render_expr(e)?;
508 }
509 self.render_expr(&_def.end)?;
510 }
511 Expression::Select(_def) => {
512 let val_line = _def.val.pos().line;
513 if self.has_comment(val_line) {
514 self.render_missed_comments(val_line)?;
515 }
516 if let Some(ref e) = _def.default {
517 let default_line = e.pos().line;
518 if self.has_comment(default_line) {
519 self.render_missed_comments(default_line)?;
520 }
521 }
522 write!(self.w, "select ")?;
523 write!(self.w, "(")?;
524 self.render_expr(&_def.val)?;
525 if let Some(ref e) = _def.default {
526 write!(self.w, ", ")?;
527 self.render_expr(e)?;
528 }
529 write!(self.w, ") => ")?;
530 self.render_tuple_def(&_def.tuple)?;
531 }
532 Expression::Simple(ref _def) => {
533 self.render_value(_def)?;
534 }
535 };
536 if did_indent {
537 self.curr_indent -= self.indent_size;
538 }
539 Ok(())
540 }
541
542 pub fn render_stmt(&mut self, stmt: &Statement, prefix_newline: bool) -> std::io::Result<()> {
543 if prefix_newline {
545 write!(self.w, "\n")?;
546 }
547 let line = stmt.pos().line;
548 self.render_comment_if_needed(line)?;
549 match stmt {
550 Statement::Let(def) => {
551 write!(&mut self.w, "let {} = ", def.name.fragment)?;
552 self.render_expr(&def.value)?;
553 }
554 Statement::Expression(_expr) => {
555 self.render_expr(&_expr)?;
556 }
557 Statement::Assert(_, def) => {
558 write!(&mut self.w, "assert ")?;
559 self.render_expr(&def)?;
560 }
561 Statement::Output(_, _tok, _expr) => {
562 write!(&mut self.w, "out {} ", _tok.fragment)?;
563 self.render_expr(&_expr)?;
564 }
565 Statement::Print(_, _tok, _expr) => {
566 write!(&mut self.w, "print {} ", _tok.fragment)?;
567 self.render_expr(&_expr)?;
568 }
569 };
570 write!(self.w, ";\n")?;
571 self.last_line = line;
572 Ok(())
573 }
574
575 pub fn render(&mut self, stmts: &Vec<Statement>) -> std::io::Result<()> {
576 let mut prefix_newline = false;
577 for v in stmts {
578 self.render_stmt(v, prefix_newline)?;
579 prefix_newline = true;
580 }
581 let comment_line = self.comment_group_lines.first().cloned();
582 if let Some(last_comment_line) = comment_line {
583 self.render_missed_comments(last_comment_line + 1)?;
584 }
585 Ok(())
586 }
587}
588
589#[cfg(test)]
590mod test;