From 189f71a90fcbc7afc024a6cff63822eb74c405e9 Mon Sep 17 00:00:00 2001
From: nico-bb <64324649+nico-bb@users.noreply.github.com>
Date: Sun, 31 Jul 2022 06:10:32 +0200
Subject: cons refactor and another attempt at comment alignment with block
 stmt

---
 src/odin/printer/document.odin | 16 ++++++++++++----
 src/odin/printer/printer.odin  | 11 +++++++++++
 src/odin/printer/visit.odin    |  8 +++++++-
 3 files changed, 30 insertions(+), 5 deletions(-)

(limited to 'src')

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
-- 
cgit v1.2.3