ucglib/ast/printer/
mod.rs

1// Copyright 2019 Jeremy Wall
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14use 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    // Indexed by line that the comment was on.
27    // We use this to determine when to print a comment in our AstPrinter
28    comment_map: Option<&'a CommentMap>,
29    last_line: usize,
30    comment_group_lines: Vec<usize>,
31}
32
33// TODO(jwall): At some point we probably want to be more aware of line length
34// in our formatting. But not at the moment.
35impl<'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        // TODO(jwall): This is probably inefficient but we'll improve it after
59        // we get it correct.
60        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        // If the field list is just 1 we might be able to collapse the tuple.
161        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            // TODO(jwall): We should maintain precision for floats?
211            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                    // if we'll be rendering a comment then we should
256                    // add a new line here
257                    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                // We print all of the comments we missed before the end of this
498                // expression before the entire range expression.
499                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        // All statements start at the beginning of a line.
544        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;