From 2c6a0e723d1d9eec6c64b91731f6659087363f45 Mon Sep 17 00:00:00 2001
From: Masaaki Goshima <goccy54@gmail.com>
Date: Wed, 13 Nov 2024 19:45:38 +0900
Subject: [PATCH] fix comment position of map node has multiple values (#532)

---
 parser/parser.go      |  7 +++++-
 parser/parser_test.go | 57 +++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 63 insertions(+), 1 deletion(-)

diff --git a/parser/parser.go b/parser/parser.go
index ec54ce7..eee8232 100644
--- a/parser/parser.go
+++ b/parser/parser.go
@@ -388,7 +388,12 @@ func (p *parser) createMapValueNode(ctx *context, key ast.MapKeyNode, colonToken
 		return nil, err
 	}
 	if comment != nil {
-		_ = value.SetComment(comment)
+		nextLineComment := key.GetToken().Position.Line < comment.GetToken().Position.Line
+		if n, ok := value.(*ast.MappingNode); ok && nextLineComment && len(n.Values) > 1 {
+			_ = n.Values[0].SetComment(comment)
+		} else {
+			_ = value.SetComment(comment)
+		}
 	}
 	return value, nil
 }
diff --git a/parser/parser_test.go b/parser/parser_test.go
index 8d4e4d1..6221aaf 100644
--- a/parser/parser_test.go
+++ b/parser/parser_test.go
@@ -1560,6 +1560,63 @@ baz:
 	})
 }
 
+func TestCommentWithMap(t *testing.T) {
+	yml := `
+single:
+  # foo comment
+  foo: bar
+
+multiple:
+    # a comment
+    a: b
+    # c comment
+    c: d
+`
+
+	file, err := parser.ParseBytes([]byte(yml), parser.ParseComments)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(file.Docs) == 0 {
+		t.Fatal("cannot get file docs")
+	}
+	if file.Docs[0].Body == nil {
+		t.Fatal("cannot get docs body")
+	}
+	mapNode, ok := file.Docs[0].Body.(*ast.MappingNode)
+	if !ok {
+		t.Fatalf("failed to get map node. got: %T\n", file.Docs[0].Body)
+	}
+	if len(mapNode.Values) != 2 {
+		t.Fatalf("failed to get map values. got %d", len(mapNode.Values))
+	}
+
+	singleNode, ok := mapNode.Values[0].Value.(*ast.MappingValueNode)
+	if !ok {
+		t.Fatalf("failed to get single node. got %T", mapNode.Values[0].Value)
+	}
+	if singleNode.GetComment().GetToken().Value != " foo comment" {
+		t.Fatalf("failed to get comment from single. got %q", singleNode.GetComment().GetToken().Value)
+	}
+
+	multiNode, ok := mapNode.Values[1].Value.(*ast.MappingNode)
+	if !ok {
+		t.Fatalf("failed to get multiple node. got: %T", mapNode.Values[1])
+	}
+	if multiNode.GetComment() != nil {
+		t.Fatalf("found unexpected comment")
+	}
+	if len(multiNode.Values) != 2 {
+		t.Fatalf("failed to get multiple node values. got %d", len(multiNode.Values))
+	}
+	if multiNode.Values[0].GetComment().GetToken().Value != " a comment" {
+		t.Fatalf("failed to get comment from multiple[0]. got %q", multiNode.Values[0].GetComment().GetToken().Value)
+	}
+	if multiNode.Values[1].GetComment().GetToken().Value != " c comment" {
+		t.Fatalf("failed to get comment from multiple[1]. got %q", multiNode.Values[1].GetComment().GetToken().Value)
+	}
+}
+
 func TestNodePath(t *testing.T) {
 	yml := `
 a: # commentA