DEV: Handle comments between statements.

src/ast/mod.rs

@@ -756,6 +756,7 @@ impl fmt::Display for Expression {
 /// Encodes a let statement in the UCG AST.
 #[derive(Debug, PartialEq, Clone)]
 pub struct LetDef {
+    pub pos: Position,
     pub name: Token,
     pub value: Expression,
 }
@@ -770,8 +771,19 @@ pub enum Statement {
     Let(LetDef),
 
     // Assert statement
-    Assert(Expression),
+    Assert(Position, Expression),
 
     // Identify an Expression for output.
     Output(Position, Token, Expression),
 }
+
+impl Statement {
+    fn pos(&self) -> &Position {
+        match self {
+            Statement::Expression(ref e) => e.pos(),
+            Statement::Let(ref def) => &def.pos,
+            Statement::Assert(ref pos, _) => pos,
+            Statement::Output(ref pos, _, _) => pos,
+        }
+    }
+}

src/ast/printer/mod.rs

@@ -29,6 +29,8 @@ where
     // Indexed by line that the comment was on.
     // We use this to determine when to print a comment in our AstPrinter
     comment_map: Option<&'a CommentMap>,
+    last_line: usize,
+    comment_group_lines: Vec<usize>,
 }
 
 // TODO(jwall): At some point we probably want to be more aware of line length
@@ -43,10 +45,14 @@ where
             curr_indent: 0,
             comment_map: None,
             w: w,
+            last_line: 0,
+            comment_group_lines: Vec::new(),
         }
     }
 
     pub fn with_comment_map(mut self, map: &'a CommentMap) -> Self {
+        self.comment_group_lines = map.keys().cloned().collect();
+        self.comment_group_lines.reverse();
         self.comment_map = Some(map);
         self
     }
@@ -77,6 +83,47 @@ where
         return true;
     }
 
+    fn print_comment_group(&mut self, line: usize) -> std::io::Result<()> {
+        if let Some(ref map) = self.comment_map {
+            let empty: Vec<Token> = Vec::new();
+            //eprintln!("comment line candidate: {}", line);
+            let cg = map.get(&line).unwrap_or(&empty);
+            //eprintln!("comment_group: {:?}", cg);
+            for c in cg.iter() {
+                write!(self.w, "// {}\n", c.fragment.trim())?;
+            }
+            self.comment_group_lines.pop();
+        }
+        Ok(())
+    }
+
+    fn render_missed_comments(&mut self, line: usize) -> std::io::Result<()> {
+        loop {
+            if let Some(next_comment_line) = self.comment_group_lines.last() {
+                let next_comment_line = *next_comment_line;
+                if next_comment_line < line {
+                    self.print_comment_group(next_comment_line)?;
+                } else {
+                    break;
+                }
+                if next_comment_line < line - 1 {
+                    write!(self.w, "\n")?;
+                }
+                continue;
+            }
+            break;
+        }
+        Ok(())
+    }
+
+    fn render_comment_if_needed(&mut self, line: usize) -> std::io::Result<()> {
+        if line > self.last_line {
+            self.render_missed_comments(line)?;
+            self.last_line = line;
+        }
+        Ok(())
+    }
+
     fn render_list_def(&mut self, def: &ListDef) -> std::io::Result<()> {
         write!(self.w, "[")?;
         self.curr_indent += self.indent_size;
@@ -351,6 +398,8 @@ where
 
     pub fn render_stmt(&mut self, stmt: &Statement) -> std::io::Result<()> {
         // All statements start at the beginning of a line.
+        let line = stmt.pos().line;
+        self.render_comment_if_needed(line)?;
         match stmt {
             Statement::Let(def) => {
                 write!(&mut self.w, "let {} = ", def.name.fragment)?;
@@ -359,7 +408,7 @@ where
             Statement::Expression(_expr) => {
                 self.render_expr(&_expr)?;
             }
-            Statement::Assert(def) => {
+            Statement::Assert(_, def) => {
                 write!(&mut self.w, "assert ")?;
                 self.render_expr(&def)?;
             }
@@ -369,6 +418,7 @@ where
             }
         };
         write!(self.w, ";\n")?;
+        self.last_line = line;
         Ok(())
     }
 
@@ -376,6 +426,12 @@ where
         for v in stmts {
             self.render_stmt(v)?;
         }
+        if let Some(last_comment_line) = self.comment_group_lines.first() {
+            eprintln!("last_comment_line is: {}", last_comment_line);
+            eprintln!("comment_map is: {:?}", self.comment_map);
+            eprintln!("coment_group_lines is: {:?}", self.comment_group_lines);
+            self.render_missed_comments(*last_comment_line + 1)?;
+        }
         Ok(())
     }
 }

src/ast/printer/test.rs

@@ -11,6 +11,7 @@
 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 // See the License for the specific language governing permissions and
 // limitations under the License.
+use std::collections::BTreeMap;
 
 use crate::ast::printer::*;
 use crate::iter::OffsetStrIter;
@@ -351,3 +352,53 @@ fn test_format_expr_list_arg_printing() {
     assert!(printer.render(&stmts).is_ok());
     assert_eq!(String::from_utf8(buffer).unwrap(), format!("{}\n", input));
 }
+
+#[test]
+fn test_statement_with_comment_printing() {
+    let mut comment_map = BTreeMap::new();
+    let input = "// add 1 + 1\n1 + 1;";
+    let stmts = assert_parse(input, Some(&mut comment_map));
+    let mut buffer: Vec<u8> = Vec::new();
+    let mut printer = AstPrinter::new(2, &mut buffer).with_comment_map(&comment_map);
+    assert!(printer.render(&stmts).is_ok());
+    assert_eq!(String::from_utf8(buffer).unwrap(), format!("{}\n", input));
+}
+
+#[test]
+fn test_statement_with_comment_printing_groups() {
+    let mut comment_map = BTreeMap::new();
+    let input = "// add 1\n// and 1\n1 + 1;";
+    let stmts = assert_parse(input, Some(&mut comment_map));
+    let mut buffer: Vec<u8> = Vec::new();
+    let mut printer = AstPrinter::new(2, &mut buffer).with_comment_map(&comment_map);
+    assert!(printer.render(&stmts).is_ok());
+    assert_eq!(String::from_utf8(buffer).unwrap(), format!("{}\n", input));
+}
+
+#[test]
+fn test_statement_with_comment_printing_multiple_groups() {
+    let mut comment_map = BTreeMap::new();
+    let input = "\n// group 1\n// more group 1\n\n// group 2\n// more group 2\n1 + 1;";
+    let stmts = assert_parse(input, Some(&mut comment_map));
+    let mut buffer: Vec<u8> = Vec::new();
+    let mut printer = AstPrinter::new(2, &mut buffer).with_comment_map(&comment_map);
+    assert!(printer.render(&stmts).is_ok());
+    assert_eq!(
+        String::from_utf8(buffer).unwrap(),
+        format!("{}\n", input.trim())
+    );
+}
+
+#[test]
+fn test_statement_with_comment_printing_comments_at_end() {
+    let mut comment_map = BTreeMap::new();
+    let input = "// group 1\n1 + 1;\n// group 2\n\n";
+    let stmts = assert_parse(input, Some(&mut comment_map));
+    let mut buffer: Vec<u8> = Vec::new();
+    let mut printer = AstPrinter::new(2, &mut buffer).with_comment_map(&comment_map);
+    assert!(printer.render(&stmts).is_ok());
+    assert_eq!(
+        String::from_utf8(buffer).unwrap(),
+        format!("{}\n", input.trim())
+    );
+}

src/ast/walk.rs

@@ -16,7 +16,7 @@ pub trait Walker {
             Statement::Expression(ref mut expr) => {
                 self.walk_expression(expr);
             }
-            Statement::Assert(ref mut expr) => {
+            Statement::Assert(_, ref mut expr) => {
                 self.walk_expression(expr);
             }
             Statement::Output(_, _, ref mut expr) => {

src/build/mod.rs

@@ -511,7 +511,7 @@ impl<'a> FileBuilder<'a> {
     fn eval_stmt(&mut self, stmt: &Statement) -> Result<Rc<Val>, Box<dyn Error>> {
         let child_scope = self.scope.clone();
         match stmt {
-            &Statement::Assert(ref expr) => self.eval_assert(&expr, &child_scope),
+            &Statement::Assert(_, ref expr) => self.eval_assert(&expr, &child_scope),
             &Statement::Let(ref def) => self.eval_let(def),
             &Statement::Expression(ref expr) => self.eval_expr(expr, &child_scope),
             // Only one output can be used per file. Right now we enforce this by

src/parse/mod.rs

@@ -793,21 +793,19 @@ make_fn!(
     )
 );
 
-fn tuple_to_let(tok: Token, expr: Expression) -> Statement {
-    Statement::Let(LetDef {
-        name: tok,
-        value: expr,
-    })
-}
-
 make_fn!(
     let_stmt_body<SliceIter<Token>, Statement>,
     do_each!(
+        pos => pos,
         name => wrap_err!(match_type!(BAREWORD), "Expected name for binding"),
         _ => punct!("="),
         val => wrap_err!(trace_parse!(expression), "Expected Expression"),
         _ => punct!(";"),
-        (tuple_to_let(name, val))
+        (Statement::Let(LetDef {
+            pos: pos,
+            name: name,
+            value: val,
+        }))
     )
 );
 
@@ -823,10 +821,11 @@ make_fn!(
 make_fn!(
     assert_statement<SliceIter<Token>, Statement>,
     do_each!(
+        pos => pos,
         _ => word!("assert"),
         expr => wrap_err!(must!(expression), "Expected Tuple {ok=<bool>, desc=<str>}"),
         _ => must!(punct!(";")),
-        (Statement::Assert(expr))
+        (Statement::Assert(pos, expr))
     )
 );
 

src/tokenizer/mod.rs

@@ -24,7 +24,7 @@ use crate::error::StackPrinter;
 use crate::iter::OffsetStrIter;
 
 pub type CommentGroup = Vec<Token>;
-pub type CommentMap = std::collections::HashMap<usize, CommentGroup>;
+pub type CommentMap = std::collections::BTreeMap<usize, CommentGroup>;
 
 fn is_symbol_char<'a>(i: OffsetStrIter<'a>) -> Result<OffsetStrIter<'a>, u8> {
     let mut _i = i.clone();
@@ -353,6 +353,12 @@ fn comment(input: OffsetStrIter) -> Result<OffsetStrIter, Token> {
                 )
             ) {
                 Result::Complete(rest, cmt) => {
+                    // Eat the new lines here before continuing
+                    let rest =
+                        match optional!(rest, either!(text_token!("\r\n"), text_token!("\n"))) {
+                            Result::Complete(next_rest, _) => next_rest,
+                            _ => rest,
+                        };
                     return Result::Complete(rest, make_tok!(CMT => cmt.to_string(), input));
                 }
                 // If we didn't find a new line then we just grab everything.
@@ -504,7 +510,9 @@ pub fn tokenize<'a>(
                         continue;
                     }
                     (&mut Some(ref mut map), _) => {
-                        out.push(tok);
+                        if tok.typ != TokenType::WS {
+                            out.push(tok);
+                        }
                         if let Some(tok) = comment_was_last {
                             map.insert(tok.pos.line, comment_group);
                             comment_group = Vec::new();

src/tokenizer/test.rs

@@ -1,3 +1,5 @@
+use std::collections::BTreeMap;
+
 use super::*;
 
 use abortable_parser::{Offsetable, Result, SliceIter};
@@ -167,6 +169,23 @@ fn test_tokenize_one_of_each() {
     assert_eq!(v[38].typ, TokenType::END);
 }
 
+#[test]
+fn test_tokenize_one_of_each_comment_map_path() {
+    let input = OffsetStrIter::new(
+        "map out filter assert let import func select as => [ ] { } ; = % / * \
+         + - . ( ) , 1 . foo \"bar\" // comment\n ; true false == < > <= >= !=",
+    );
+    let mut comment_map = BTreeMap::new();
+    let result = tokenize(input.clone(), Some(&mut comment_map));
+    assert!(result.is_ok(), format!("result {:?} is not ok", result));
+    let v = result.unwrap();
+    for (i, t) in v.iter().enumerate() {
+        println!("{}: {:?}", i, t);
+    }
+    assert_eq!(v.len(), 39);
+    assert_eq!(v[38].typ, TokenType::END);
+}
+
 #[test]
 fn test_parse_has_end() {
     let input = OffsetStrIter::new("foo");
@@ -327,3 +346,24 @@ fn test_match_type() {
         res => assert!(false, format!("Fail: {:?}", res)),
     }
 }
+
+#[test]
+fn test_tokenize_builds_comment_map() {
+    let input = OffsetStrIter::new("// comment 1\n\n//comment 2");
+    let mut comment_map = BTreeMap::new();
+    let result = tokenize(input.clone(), Some(&mut comment_map));
+    assert!(result.is_ok(), format!("result {:?} is not ok", result));
+
+    assert_eq!(comment_map.len(), 2);
+}
+
+#[test]
+fn test_tokenize_builds_comment_map_groups() {
+    let input = OffsetStrIter::new("// first part\n// comment 1\n\n//comment 2");
+    let mut comment_map = BTreeMap::new();
+    let result = tokenize(input.clone(), Some(&mut comment_map));
+    assert!(result.is_ok(), format!("result {:?} is not ok", result));
+
+    assert_eq!(comment_map.len(), 2);
+    assert_eq!(comment_map[&2].len(), 2);
+}