From 9dc770e29cd7a174b8a3abf3a4c95e9fab2093e3 Mon Sep 17 00:00:00 2001
From: Lukas Neubert <serkonda@protonmail.com>
Date: Mon, 22 Feb 2021 12:04:48 +0100
Subject: [PATCH] fmt: do not struggle with comments inside maps (#8897)

---
 vlib/v/ast/ast.v                  |  8 +++++---
 vlib/v/fmt/comments.v             |  7 ++++---
 vlib/v/fmt/fmt.v                  |  2 ++
 vlib/v/fmt/tests/comments_keep.vv | 21 ++++++++++++++++++++-
 vlib/v/parser/containers.v        |  7 ++++++-
 5 files changed, 37 insertions(+), 8 deletions(-)

diff --git a/vlib/v/ast/ast.v b/vlib/v/ast/ast.v
index 4ea7175e6d..7ee463dde0 100644
--- a/vlib/v/ast/ast.v
+++ b/vlib/v/ast/ast.v
@@ -979,9 +979,11 @@ pub mut:
 
 pub struct MapInit {
 pub:
-	pos  token.Position
-	keys []Expr
-	vals []Expr
+	pos       token.Position
+	keys      []Expr
+	vals      []Expr
+	comments  [][]Comment // comments after key-value pairs
+	pre_cmnts []Comment   // comments before the first key-value pair
 pub mut:
 	typ        table.Type
 	key_type   table.Type
diff --git a/vlib/v/fmt/comments.v b/vlib/v/fmt/comments.v
index da77004df7..b09db45430 100644
--- a/vlib/v/fmt/comments.v
+++ b/vlib/v/fmt/comments.v
@@ -83,12 +83,13 @@ pub fn (mut f Fmt) comment(node ast.Comment, options CommentsOptions) {
 pub fn (mut f Fmt) comments(comments []ast.Comment, options CommentsOptions) {
 	mut prev_line := options.prev_line
 	for i, c in comments {
+		if options.prev_line > -1 && ((c.pos.line_nr > prev_line && f.out.last_n(1) != '\n')
+			|| (c.pos.line_nr > prev_line + 1 && f.out.last_n(2) != '\n\n')) {
+			f.writeln('')
+		}
 		if !f.out.last_n(1)[0].is_space() {
 			f.write(' ')
 		}
-		if options.prev_line > -1 && c.pos.line_nr > prev_line + 1 {
-			f.writeln('')
-		}
 		f.comment(c, options)
 		if !options.iembed && (i < comments.len - 1 || options.has_nl) {
 			f.writeln('')
diff --git a/vlib/v/fmt/fmt.v b/vlib/v/fmt/fmt.v
index f86007c1cd..894bb3a403 100644
--- a/vlib/v/fmt/fmt.v
+++ b/vlib/v/fmt/fmt.v
@@ -2068,6 +2068,7 @@ pub fn (mut f Fmt) map_init(it ast.MapInit) {
 	}
 	f.writeln('map{')
 	f.indent++
+	f.comments(it.pre_cmnts, {})
 	mut max_field_len := 0
 	for key in it.keys {
 		if key.str().len > max_field_len {
@@ -2079,6 +2080,7 @@ pub fn (mut f Fmt) map_init(it ast.MapInit) {
 		f.write(': ')
 		f.write(strings.repeat(` `, max_field_len - key.str().len))
 		f.expr(it.vals[i])
+		f.comments(it.comments[i], prev_line: it.vals[i].position().last_line, has_nl: false)
 		f.writeln('')
 	}
 	f.indent--
diff --git a/vlib/v/fmt/tests/comments_keep.vv b/vlib/v/fmt/tests/comments_keep.vv
index b923ee34e7..d93309b3df 100644
--- a/vlib/v/fmt/tests/comments_keep.vv
+++ b/vlib/v/fmt/tests/comments_keep.vv
@@ -58,10 +58,29 @@ fn linebreaks_in_block_comments() {
 	*****/
 }
 
+fn map_comments() {
+	mymap := map{
+		// pre
+		`:`: 1
+		`!`: 2 // after
+		// and between
+		`%`: 3
+		// between
+		// between second
+		`$`: 4
+		`&`: 5
+		// post
+	}
+}
+
 fn ifs_comments_and_empty_lines() {
 	if true {
 	}
-	// some comment after an if without else
+	// some comment direct after an if without else
+	if false {
+	} else {
+	}
+	// some comment direct after an else
 	if false {
 	}
 
diff --git a/vlib/v/parser/containers.v b/vlib/v/parser/containers.v
index 036e76aa47..32c41a5951 100644
--- a/vlib/v/parser/containers.v
+++ b/vlib/v/parser/containers.v
@@ -160,7 +160,9 @@ fn (mut p Parser) map_init() ast.MapInit {
 	first_pos := p.prev_tok.position()
 	mut keys := []ast.Expr{}
 	mut vals := []ast.Expr{}
-	for p.tok.kind != .rcbr && p.tok.kind != .eof {
+	mut comments := [][]ast.Comment{}
+	pre_cmnts := p.eat_comments({})
+	for p.tok.kind !in [.rcbr, .eof] {
 		key := p.expr(0)
 		keys << key
 		p.check(.colon)
@@ -169,10 +171,13 @@ fn (mut p Parser) map_init() ast.MapInit {
 		if p.tok.kind == .comma {
 			p.next()
 		}
+		comments << p.eat_comments({})
 	}
 	return ast.MapInit{
 		keys: keys
 		vals: vals
 		pos: first_pos.extend_with_last_line(p.tok.position(), p.tok.line_nr)
+		comments: comments
+		pre_cmnts: pre_cmnts
 	}
 }