From f18aa878b9d368848dca85f2301e01395924c148 Mon Sep 17 00:00:00 2001
From: Masaaki Goshima <goccy54@gmail.com>
Date: Thu, 14 Sep 2023 12:22:17 +0900
Subject: [PATCH] Fix comment parsing with null value

---
 parser/parser.go      | 36 +++++++++++++++++++++--
 parser/parser_test.go | 68 +++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 101 insertions(+), 3 deletions(-)

diff --git a/parser/parser.go b/parser/parser.go
index 568e6ad4..6989a378 100644
--- a/parser/parser.go
+++ b/parser/parser.go
@@ -156,15 +156,38 @@ func (p *parser) createMapValueNode(ctx *context, key ast.MapKeyNode, colonToken
 		ctx.insertToken(ctx.idx, nullToken)
 		return ast.Null(nullToken), nil
 	}
-
+	var comment *ast.CommentGroupNode
+	if tk.Type == token.CommentType {
+		comment = p.parseCommentOnly(ctx)
+		if comment != nil {
+			comment.SetPath(ctx.withChild(key.GetToken().Value).path)
+		}
+		tk = ctx.currentToken()
+	}
 	if tk.Position.Column == key.GetToken().Position.Column && tk.Type == token.StringType {
 		// in this case,
 		// ----
 		// key: <value does not defined>
 		// next
+
 		nullToken := p.createNullToken(colonToken)
 		ctx.insertToken(ctx.idx, nullToken)
-		return ast.Null(nullToken), nil
+		nullNode := ast.Null(nullToken)
+
+		if comment != nil {
+			nullNode.SetComment(comment)
+		} else {
+			// If there is a comment, it is already bound to the key node,
+			// so remove the comment from the key to bind it to the null value.
+			keyComment := key.GetComment()
+			if keyComment != nil {
+				if err := key.SetComment(nil); err != nil {
+					return nil, err
+				}
+				nullNode.SetComment(keyComment)
+			}
+		}
+		return nullNode, nil
 	}
 
 	if tk.Position.Column < key.GetToken().Position.Column {
@@ -174,13 +197,20 @@ func (p *parser) createMapValueNode(ctx *context, key ast.MapKeyNode, colonToken
 		// next
 		nullToken := p.createNullToken(colonToken)
 		ctx.insertToken(ctx.idx, nullToken)
-		return ast.Null(nullToken), nil
+		nullNode := ast.Null(nullToken)
+		if comment != nil {
+			nullNode.SetComment(comment)
+		}
+		return nullNode, nil
 	}
 
 	value, err := p.parseToken(ctx, ctx.currentToken())
 	if err != nil {
 		return nil, errors.Wrapf(err, "failed to parse mapping 'value' node")
 	}
+	if comment != nil {
+		value.SetComment(comment)
+	}
 	return value, nil
 }
 
diff --git a/parser/parser_test.go b/parser/parser_test.go
index 6044699e..d24b262a 100644
--- a/parser/parser_test.go
+++ b/parser/parser_test.go
@@ -808,6 +808,74 @@ foo: > # comment
 	}
 }
 
+func TestCommentWithNull(t *testing.T) {
+	t.Run("same line", func(t *testing.T) {
+		content := `
+foo:
+  bar: # comment
+  baz: 1
+`
+		expected := `
+foo:
+  bar: null # comment
+  baz: 1`
+		f, err := parser.ParseBytes([]byte(content), parser.ParseComments)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if len(f.Docs) != 1 {
+			t.Fatal("failed to parse content with same line comment")
+		}
+		if f.Docs[0].String() != strings.TrimPrefix(expected, "\n") {
+			t.Fatal("failed to parse comment")
+		}
+	})
+	t.Run("next line", func(t *testing.T) {
+		content := `
+foo:
+  bar:
+    # comment
+  baz: 1
+`
+		expected := `
+foo:
+  bar: null # comment
+  baz: 1`
+		f, err := parser.ParseBytes([]byte(content), parser.ParseComments)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if len(f.Docs) != 1 {
+			t.Fatal("failed to parse content with next line comment")
+		}
+		if f.Docs[0].String() != strings.TrimPrefix(expected, "\n") {
+			t.Fatal("failed to parse comment")
+		}
+	})
+	t.Run("next line and different indent", func(t *testing.T) {
+		content := `
+foo:
+  bar:
+ # comment
+baz: 1`
+		f, err := parser.ParseBytes([]byte(content), parser.ParseComments)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if len(f.Docs) != 1 {
+			t.Fatal("failed to parse content with next line comment")
+		}
+		expected := `
+foo:
+  bar: null # comment
+baz: 1`
+		if f.Docs[0].String() != strings.TrimPrefix(expected, "\n") {
+			t.Fatal("failed to parse comment")
+		}
+	})
+
+}
+
 func TestNodePath(t *testing.T) {
 	yml := `
 a: # commentA