cons refactor and another attempt at comment alignment with block stmt

5 files changed, 44 insertions, 5 deletions

diff --git a/src/odin/printer/document.odin b/src/odin/printer/document.odin
index 45caba4..8a3f858 100644
--- a/src/odin/printer/document.odin
+++ b/src/odin/printer/document.odin
@@ -215,11 +215,19 @@ group :: proc(grouped_document: ^Document, options := Document_Group_Options{},
 
 cons :: proc(elems: ..^Document, allocator := context.allocator) -> ^Document {
 	document := new(Document, allocator)
-	c := Document_Cons {
-		elements = make([]^Document, len(elems), allocator),
+	elements := make([dynamic]^Document, allocator)
+	for elem in elems {
+		#partial switch e in elem {
+		case Document_Nil:
+			continue
+		case Document_Cons:
+			append(&elements, ..e.elements)
+		case:
+			append(&elements, elem)
+		}
 	}
-	for elem, i in elems {
-		c.elements[i] = elem
+	c := Document_Cons {
+		elements = elements[:],
 	}
 	document^ = c
 	return document
diff --git a/src/odin/printer/printer.odin b/src/odin/printer/printer.odin
index e2fbe8c..183ed75 100644
--- a/src/odin/printer/printer.odin
+++ b/src/odin/printer/printer.odin
@@ -10,6 +10,7 @@ Printer :: struct {
 	string_builder:       strings.Builder,
 	config:               Config,
 	comments:             [dynamic]^ast.Comment_Group,
+	comments_option:      map[int]Line_Suffix_Option,
 	latest_comment_index: int,
 	allocator:            mem.Allocator,
 	file:                 ^ast.File,
@@ -74,6 +75,11 @@ Newline_Style :: enum {
 	LF,
 }
 
+Line_Suffix_Option :: enum {
+	Default,
+	Indent,
+}
+
 
 when ODIN_OS ==  .Windows {
 	default_style := Config {
@@ -138,6 +144,11 @@ build_disabled_lines_info :: proc(p: ^Printer) {
 	}
 }
 
+@private
+set_comment_option :: proc(p: ^Printer, line: int, option: Line_Suffix_Option) {
+	p.comments_option[line] = option
+}
+
 print :: proc {
 	print_file,
 	print_expr,
diff --git a/src/odin/printer/visit.odin b/src/odin/printer/visit.odin
index 653b5be..fdcc661 100644
--- a/src/odin/printer/visit.odin
+++ b/src/odin/printer/visit.odin
@@ -117,7 +117,12 @@ visit_comment :: proc(p: ^Printer, comment: tokenizer.Token) -> (int, ^Document)
 			return 1, empty()
 		} else if comment.pos.line == p.source_position.line && p.source_position.column != 1 {
 			p.source_position = comment.pos
-			return newlines_before_comment, cons_with_nopl(document, line_suffix(comment.text))
+			if comment.pos.line in p.comments_option {
+				delete_key(&p.comments_option, comment.pos.line)
+				return newlines_before_comment, cons_with_nopl(document, cons(text(p.indentation), line_suffix(comment.text)))
+			} else {
+				return newlines_before_comment, cons_with_nopl(document, line_suffix(comment.text))
+			}
 		} else {
 			p.source_position = comment.pos
 			return newlines_before_comment, cons(document, line_suffix(comment.text))
@@ -1460,6 +1465,7 @@ visit_matrix_comp_lit :: proc(p: ^Printer, comp_lit: ^ast.Comp_Lit, matrix_type:
 @(private)
 visit_begin_brace :: proc(p: ^Printer, begin: tokenizer.Pos, type: Block_Type, count := 0, same_line_spaces_before := 1) -> ^Document {
 	set_source_position(p, begin)
+	set_comment_option(p, begin.line, .Indent)
 
 	newline_braced := p.config.brace_style == .Allman
 	newline_braced |= p.config.brace_style == .K_And_R && type == .Proc
diff --git a/tools/odinfmt/tests/.snapshots/comments.odin b/tools/odinfmt/tests/.snapshots/comments.odin
index d872984..262d80e 100644
--- a/tools/odinfmt/tests/.snapshots/comments.odin
+++ b/tools/odinfmt/tests/.snapshots/comments.odin
@@ -33,3 +33,10 @@ line_comments_one_line_seperation :: proc() {
 
 
 //More comments YAY
+
+bracket_comments_alignment :: proc() {
+	{ 	// Describe block
+		a := 10
+		// etc..
+	}
+}
diff --git a/tools/odinfmt/tests/comments.odin b/tools/odinfmt/tests/comments.odin
index f4e5b06..a7bea72 100644
--- a/tools/odinfmt/tests/comments.odin
+++ b/tools/odinfmt/tests/comments.odin
@@ -33,3 +33,10 @@ line_comments_one_line_seperation :: proc() {
 
 
 //More comments YAY
+
+bracket_comments_alignment :: proc() {
+    {   // Describe block
+        a := 10
+        // etc..
+    }
+}
\ No newline at end of file