Skip to main content

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(
159        &mut self,
160        def: &Vec<(Token, Option<Expression>, Expression)>,
161    ) -> std::io::Result<()> {
162        self.w.write(&['{' as u8])?;
163        // If the field list is just 1 we might be able to collapse the tuple.
164        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            // TODO(jwall): We should maintain precision for floats?
219            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                    // if we'll be rendering a comment then we should
264                    // add a new line here
265                    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                // We print all of the comments we missed before the end of this
519                // expression before the entire range expression.
520                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        // All statements start at the beginning of a line.
565        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;