julefmt: improve comment support

src/comment.jule

@@ -40,15 +40,30 @@ impl CommentMap {
         ret cm
     }
 
-    fn pop(mut self, row: int): &Comment {
+    fn first(mut self, row: int): &Comment {
         if self.map.len == 0 {
             ret nil
         }
         let mut c = self.map[0]
         if c.row > row {
             ret nil
         }
+        ret c
+    }
+
+    fn drop_first(mut self) {
         self.map = self.map[1:]
+    }
+
+    fn pop(mut self, row: int): &Comment {
+        if self.map.len == 0 {
+            ret nil
+        }
+        let mut c = self.map[0]
+        if c.row > row {
+            ret nil
+        }
+        self.drop_first()
         ret c
     }
 }

src/format.jule

@@ -40,18 +40,23 @@ impl Formatter {
         self.indent = self.indent[:self.indent.len - 4]
     }
 
-    fn write_comments(&self, row: int) {
+    fn write_comments_except(&self, row: int) {
         for {
-            let mut c = self.cm.pop(row)
-            if c == nil {
+            let mut c = self.cm.first(row)
+            if c == nil || c.row == row {
                 break
             }
+            self.cm.drop_first()
             self.write(self.indent)
             self.write(c.txt)
             self.write("\n")
         }
     }
 
+    fn write_comments(&self, row: int) {
+        self.write_comments_except(row + 1)
+    }
+
     fn format_expr(&self, &expr: &ast::Expr) {
         let ef = ExprFormatter{
             fmt: self,
@@ -376,27 +381,55 @@ impl Formatter {
     }
 
     fn group_decls[T](&self, nodes: []ast::Node, mut &i: int) {
+        const CAP = 1 << 4
+        let mut lines = make([]str, 0, CAP)
+        let mut comments = make([]&Comment, 0, CAP)
+
         let mut row = -1
+        let mut max = 0
         for i < nodes.len {
             let node = nodes[i]
             match type node.data {
             | T:
                 let decl = T(node.data)
                 if row == -1 || decl.token.row - 1 == row {
                     row = decl.token.row
+                    self.write_comments_except(row)
+                    if self.cm.map.len > 0 && self.cm.map[0].row == row {
+                        comments = append(comments, self.cm.pop(row))
+                    } else {
+                        let mut c: &Comment
+                        comments = append(comments, c)
+                    }
+                    let n = self.buf.len
                     match type T {
                     | &ast::VarDecl:
                         self.var_decl(decl)
                     | &ast::TypeAliasDecl:
                         self.type_alias_decl(decl)
                     }
-                    self.write("\n")
+                    let diff = self.buf.len - n
+                    if max < diff {
+                        max = diff
+                    }
+                    lines = append(lines, self.buf[n:])
+                    self.buf = self.buf[:self.buf.len - diff]
                     i++
                     continue
                 }
             }
             break
         }
+
+        for j, line in lines {
+            self.write(line)
+            let c = comments[j]
+            if c != nil {
+                self.write(strings::repeat(" ", max - line.len + 1))
+                self.write(c.txt)
+            }
+            self.write("\n")
+        }
     }
 
     fn impl_decl(&self, mut d: &ast::Impl) {
@@ -430,6 +463,14 @@ impl Formatter {
     }
 
     fn node(&self, mut &node: ast::Node) {
+        match type node.data {
+        | &ast::TypeAliasDecl:
+            self.group_decls[&ast::TypeAliasDecl](self.f.nodes, self.i)
+        | &ast::VarDecl:
+            self.group_decls[&ast::VarDecl](self.f.nodes, self.i)
+        }
+
+        self.write_comments(node.token.row)
         match type node.data {
         | &ast::EnumDecl:
             self.enum_decl((&ast::EnumDecl)(node.data))
@@ -439,10 +480,6 @@ impl Formatter {
             self.struct_decl((&ast::StructDecl)(node.data))
         | &ast::TraitDecl:
             self.trait_decl((&ast::TraitDecl)(node.data))
-        | &ast::TypeAliasDecl:
-            self.group_decls[&ast::TypeAliasDecl](self.f.nodes, self.i)
-        | &ast::VarDecl:
-            self.group_decls[&ast::VarDecl](self.f.nodes, self.i)
         | &ast::Impl:
             self.impl_decl((&ast::Impl)(node.data))
         }
@@ -453,7 +490,6 @@ impl Formatter {
         let mut row = 0
         for self.i < self.f.nodes.len; self.i++ {
             let mut node = self.f.nodes[self.i]
-            self.write_comments(node.token.row)
             let old = self.i
             self.node(node)
             match {