Repository: goccy/go-yaml
Pull request #349 (2023-02-28 13:22:00)

Fix comment option
resolve https://github.com/goccy/go-yaml/issues/314

Reversible conversion of comments was not working in various cases, which has been corrected.


Patch:
diff --git a/ast/ast.go b/ast/ast.go
index d5aa1d8..0e3462e 100644
--- a/ast/ast.go
+++ b/ast/ast.go
@@ -655,7 +655,7 @@ func (n *NullNode) GetValue() interface{} {
 // String returns `null` text
 func (n *NullNode) String() string {
 	if n.Comment != nil {
-		return fmt.Sprintf("null %s", n.Comment.String())
+		return addCommentString("null", n.Comment)
 	}
 	return n.stringWithoutComment()
 }
@@ -1154,6 +1154,7 @@ type MappingNode struct {
 	End         *token.Token
 	IsFlowStyle bool
 	Values      []*MappingValueNode
+	FootComment *CommentGroupNode
 }
 
 func (n *MappingNode) startPos() *token.Position {
@@ -1319,9 +1320,10 @@ func (n *MappingKeyNode) MarshalYAML() ([]byte, error) {
 // MappingValueNode type of mapping value
 type MappingValueNode struct {
 	*BaseNode
-	Start *token.Token
-	Key   MapKeyNode
-	Value Node
+	Start       *token.Token
+	Key         MapKeyNode
+	Value       Node
+	FootComment *CommentGroupNode
 }
 
 // Replace replace value node.
@@ -1370,14 +1372,20 @@ func (n *MappingValueNode) SetIsFlowStyle(isFlow bool) {
 
 // String mapping value to text
 func (n *MappingValueNode) String() string {
+	var text string
 	if n.Comment != nil {
-		return fmt.Sprintf(
+		text = fmt.Sprintf(
 			"%s\n%s",
 			n.Comment.StringWithSpace(n.Key.GetToken().Position.Column-1),
 			n.toString(),
 		)
+	} else {
+		text = n.toString()
 	}
-	return n.toString()
+	if n.FootComment != nil {
+		text += fmt.Sprintf("\n%s", n.FootComment.StringWithSpace(n.Key.GetToken().Position.Column-1))
+	}
+	return text
 }
 
 func (n *MappingValueNode) toString() string {
@@ -1472,11 +1480,12 @@ func (m *ArrayNodeIter) Len() int {
 // SequenceNode type of sequence node
 type SequenceNode struct {
 	*BaseNode
-	Start         *token.Token
-	End           *token.Token
-	IsFlowStyle   bool
-	Values        []Node
-	ValueComments []*CommentGroupNode
+	Start             *token.Token
+	End               *token.Token
+	IsFlowStyle       bool
+	Values            []Node
+	ValueHeadComments []*CommentGroupNode
+	FootComment       *CommentGroupNode
 }
 
 // Replace replace value node.
@@ -1576,11 +1585,14 @@ func (n *SequenceNode) blockStyleString() string {
 			newValues = append(newValues, fmt.Sprintf("%s  %s", space, trimmed))
 		}
 		newValue := strings.Join(newValues, "\n")
-		if len(n.ValueComments) == len(n.Values) && n.ValueComments[idx] != nil {
-			values = append(values, n.ValueComments[idx].StringWithSpace(n.Start.Position.Column-1))
+		if len(n.ValueHeadComments) == len(n.Values) && n.ValueHeadComments[idx] != nil {
+			values = append(values, n.ValueHeadComments[idx].StringWithSpace(n.Start.Position.Column-1))
 		}
 		values = append(values, fmt.Sprintf("%s- %s", space, newValue))
 	}
+	if n.FootComment != nil {
+		values = append(values, n.FootComment.StringWithSpace(n.Start.Position.Column-1))
+	}
 	return strings.Join(values, "\n")
 }
 
@@ -1866,8 +1878,8 @@ func (n *CommentGroupNode) String() string {
 }
 
 func (n *CommentGroupNode) StringWithSpace(col int) string {
-	space := strings.Repeat(" ", col)
 	values := []string{}
+	space := strings.Repeat(" ", col)
 	for _, comment := range n.Comments {
 		values = append(values, space+comment.String())
 	}
diff --git a/decode.go b/decode.go
index 8cd7d85..369ee26 100644
--- a/decode.go
+++ b/decode.go
@@ -12,6 +12,7 @@ import (
 	"os"
 	"path/filepath"
 	"reflect"
+	"sort"
 	"strconv"
 	"time"
 
@@ -155,22 +156,106 @@ func (d *Decoder) setPathToCommentMap(node ast.Node) {
 	if d.toCommentMap == nil {
 		return
 	}
+	d.addHeadOrLineCommentToMap(node)
+	d.addFootCommentToMap(node)
+}
+
+func (d *Decoder) addHeadOrLineCommentToMap(node ast.Node) {
+	sequence, ok := node.(*ast.SequenceNode)
+	if ok {
+		d.addSequenceNodeCommentToMap(sequence)
+		return
+	}
 	commentGroup := node.GetComment()
 	if commentGroup == nil {
 		return
 	}
 	texts := []string{}
+	targetLine := node.GetToken().Position.Line
+	minCommentLine := math.MaxInt64
 	for _, comment := range commentGroup.Comments {
+		if minCommentLine > comment.Token.Position.Line {
+			minCommentLine = comment.Token.Position.Line
+		}
 		texts = append(texts, comment.Token.Value)
 	}
 	if len(texts) == 0 {
 		return
 	}
-	if len(texts) == 1 {
-		d.toCommentMap[node.GetPath()] = LineComment(texts[0])
+	commentPath := node.GetPath()
+	if minCommentLine < targetLine {
+		d.addCommentToMap(commentPath, HeadComment(texts...))
 	} else {
-		d.toCommentMap[node.GetPath()] = HeadComment(texts...)
+		d.addCommentToMap(commentPath, LineComment(texts[0]))
+	}
+}
+
+func (d *Decoder) addSequenceNodeCommentToMap(node *ast.SequenceNode) {
+	if len(node.ValueHeadComments) != 0 {
+		for idx, headComment := range node.ValueHeadComments {
+			if headComment == nil {
+				continue
+			}
+			texts := make([]string, 0, len(headComment.Comments))
+			for _, comment := range headComment.Comments {
+				texts = append(texts, comment.Token.Value)
+			}
+			if len(texts) != 0 {
+				d.addCommentToMap(node.Values[idx].GetPath(), HeadComment(texts...))
+			}
+		}
+	}
+	firstElemHeadComment := node.GetComment()
+	if firstElemHeadComment != nil {
+		texts := make([]string, 0, len(firstElemHeadComment.Comments))
+		for _, comment := range firstElemHeadComment.Comments {
+			texts = append(texts, comment.Token.Value)
+		}
+		if len(texts) != 0 {
+			d.addCommentToMap(node.Values[0].GetPath(), HeadComment(texts...))
+		}
+	}
+}
+
+func (d *Decoder) addFootCommentToMap(node ast.Node) {
+	var (
+		footComment     *ast.CommentGroupNode
+		footCommentPath string = node.GetPath()
+	)
+	switch n := node.(type) {
+	case *ast.SequenceNode:
+		if len(n.Values) != 0 {
+			footCommentPath = n.Values[len(n.Values)-1].GetPath()
+		}
+		footComment = n.FootComment
+	case *ast.MappingNode:
+		footComment = n.FootComment
+	case *ast.MappingValueNode:
+		footComment = n.FootComment
+	}
+	if footComment == nil {
+		return
+	}
+	var texts []string
+	for _, comment := range footComment.Comments {
+		texts = append(texts, comment.Token.Value)
+	}
+	if len(texts) != 0 {
+		d.addCommentToMap(footCommentPath, FootComment(texts...))
+	}
+}
+
+func (d *Decoder) addCommentToMap(path string, comment *Comment) {
+	for _, c := range d.toCommentMap[path] {
+		if c.Position == comment.Position {
+			// already added same comment
+			return
+		}
 	}
+	d.toCommentMap[path] = append(d.toCommentMap[path], comment)
+	sort.Slice(d.toCommentMap[path], func(i, j int) bool {
+		return d.toCommentMap[path][i].Position < d.toCommentMap[path][j].Position
+	})
 }
 
 func (d *Decoder) nodeToValue(node ast.Node) interface{} {
diff --git a/encode.go b/encode.go
index 1772be4..4543138 100644
--- a/encode.go
+++ b/encode.go
@@ -38,7 +38,7 @@ type Encoder struct {
 	anchorCallback             func(*ast.AnchorNode, interface{}) error
 	anchorPtrToNameMap         map[uintptr]string
 	useLiteralStyleIfMultiline bool
-	commentMap                 map[*Path]*Comment
+	commentMap                 map[*Path][]*Comment
 	written                    bool
 
 	line        int
@@ -121,41 +121,123 @@ func (e *Encoder) setCommentByCommentMap(node ast.Node) error {
 	if e.commentMap == nil {
 		return nil
 	}
-	for path, comment := range e.commentMap {
+	for path, comments := range e.commentMap {
 		n, err := path.FilterNode(node)
 		if err != nil {
 			return errors.Wrapf(err, "failed to filter node")
 		}
-		comments := []*token.Token{}
-		for _, text := range comment.Texts {
-			comments = append(comments, token.New(text, text, nil))
+		if n == nil {
+			continue
 		}
-		commentGroup := ast.CommentGroup(comments)
-		switch comment.Position {
-		case CommentLinePosition:
-			if err := n.SetComment(commentGroup); err != nil {
-				return errors.Wrapf(err, "failed to set comment")
-			}
-		case CommentHeadPosition:
-			parent := ast.Parent(node, n)
-			if parent == nil {
-				return ErrUnsupportedHeadPositionType(node)
+		for _, comment := range comments {
+			commentTokens := []*token.Token{}
+			for _, text := range comment.Texts {
+				commentTokens = append(commentTokens, token.New(text, text, nil))
 			}
-			switch node := parent.(type) {
-			case *ast.MappingValueNode:
-				if err := node.SetComment(commentGroup); err != nil {
-					return errors.Wrapf(err, "failed to set comment")
+			commentGroup := ast.CommentGroup(commentTokens)
+			switch comment.Position {
+			case CommentHeadPosition:
+				if err := e.setHeadComment(node, n, commentGroup); err != nil {
+					return errors.Wrapf(err, "failed to set head comment")
+				}
+			case CommentLinePosition:
+				if err := e.setLineComment(node, n, commentGroup); err != nil {
+					return errors.Wrapf(err, "failed to set line comment")
 				}
-			case *ast.MappingNode:
-				if err := node.SetComment(commentGroup); err != nil {
-					return errors.Wrapf(err, "failed to set comment")
+			case CommentFootPosition:
+				if err := e.setFootComment(node, n, commentGroup); err != nil {
+					return errors.Wrapf(err, "failed to set foot comment")
 				}
 			default:
-				return ErrUnsupportedHeadPositionType(node)
+				return ErrUnknownCommentPositionType
+			}
+		}
+	}
+	return nil
+}
+
+func (e *Encoder) setHeadComment(node ast.Node, filtered ast.Node, comment *ast.CommentGroupNode) error {
+	parent := ast.Parent(node, filtered)
+	if parent == nil {
+		return ErrUnsupportedHeadPositionType(node)
+	}
+	switch p := parent.(type) {
+	case *ast.MappingValueNode:
+		if err := p.SetComment(comment); err != nil {
+			return errors.Wrapf(err, "failed to set comment")
+		}
+	case *ast.MappingNode:
+		if err := p.SetComment(comment); err != nil {
+			return errors.Wrapf(err, "failed to set comment")
+		}
+	case *ast.SequenceNode:
+		if len(p.ValueHeadComments) == 0 {
+			p.ValueHeadComments = make([]*ast.CommentGroupNode, len(p.Values))
+		}
+		var foundIdx int
+		for idx, v := range p.Values {
+			if v == filtered {
+				foundIdx = idx
+				break
 			}
-		default:
-			return ErrUnknownCommentPositionType
 		}
+		p.ValueHeadComments[foundIdx] = comment
+	default:
+		return ErrUnsupportedHeadPositionType(node)
+	}
+	return nil
+}
+
+func (e *Encoder) setLineComment(node ast.Node, filtered ast.Node, comment *ast.CommentGroupNode) error {
+	switch filtered.(type) {
+	case *ast.MappingValueNode, *ast.SequenceNode:
+		// Line comment cannot be set for mapping value node.
+		// It should probably be set for the parent map node
+		if err := e.setLineCommentToParentMapNode(node, filtered, comment); err != nil {
+			return errors.Wrapf(err, "failed to set line comment to parent node")
+		}
+	default:
+		if err := filtered.SetComment(comment); err != nil {
+			return errors.Wrapf(err, "failed to set comment")
+		}
+	}
+	return nil
+}
+
+func (e *Encoder) setLineCommentToParentMapNode(node ast.Node, filtered ast.Node, comment *ast.CommentGroupNode) error {
+	parent := ast.Parent(node, filtered)
+	if parent == nil {
+		return ErrUnsupportedLinePositionType(node)
+	}
+	switch p := parent.(type) {
+	case *ast.MappingValueNode:
+		if err := p.Key.SetComment(comment); err != nil {
+			return errors.Wrapf(err, "failed to set comment")
+		}
+	case *ast.MappingNode:
+		if err := p.SetComment(comment); err != nil {
+			return errors.Wrapf(err, "failed to set comment")
+		}
+	default:
+		return ErrUnsupportedLinePositionType(parent)
+	}
+	return nil
+}
+
+func (e *Encoder) setFootComment(node ast.Node, filtered ast.Node, comment *ast.CommentGroupNode) error {
+	parent := ast.Parent(node, filtered)
+	if parent == nil {
+		return ErrUnsupportedFootPositionType(node)
+	}
+	switch n := parent.(type) {
+	case *ast.MappingValueNode:
+		n.FootComment = comment
+	case *ast.MappingNode:
+		n.FootComment = comment
+	case *ast.SequenceNode:
+		n.FootComment = comment
+	default:
+		return ErrUnsupportedFootPositionType(n)
 	}
 	return nil
 }
diff --git a/error.go b/error.go
index 028b488..163dcc5 100644
--- a/error.go
+++ b/error.go
@@ -18,6 +18,14 @@ func ErrUnsupportedHeadPositionType(node ast.Node) error {
 	return xerrors.Errorf("unsupported comment head position for %s", node.Type())
 }
 
+func ErrUnsupportedLinePositionType(node ast.Node) error {
+	return xerrors.Errorf("unsupported comment line position for %s", node.Type())
+}
+
+func ErrUnsupportedFootPositionType(node ast.Node) error {
+	return xerrors.Errorf("unsupported comment foot position for %s", node.Type())
+}
+
 // IsInvalidQueryError whether err is ErrInvalidQuery or not.
 func IsInvalidQueryError(err error) bool {
 	return xerrors.Is(err, ErrInvalidQuery)
diff --git a/go.mod b/go.mod
index 3356ece..4f325fe 100644
--- a/go.mod
+++ b/go.mod
@@ -5,6 +5,7 @@ go 1.12
 require (
 	github.com/fatih/color v1.10.0
 	github.com/go-playground/validator/v10 v10.4.1
+	github.com/google/go-cmp v0.5.9
 	github.com/mattn/go-colorable v0.1.8
 	golang.org/x/sys v0.0.0-20220406163625-3f8b81556e12 // indirect
 	golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1
diff --git a/go.sum b/go.sum
index fbb7f81..bc63cea 100644
--- a/go.sum
+++ b/go.sum
@@ -10,6 +10,8 @@ github.com/go-playground/universal-translator v0.17.0 h1:icxd5fm+REJzpZx7ZfpaD87
 github.com/go-playground/universal-translator v0.17.0/go.mod h1:UkSxE5sNxxRwHyU+Scu5vgOQjsIJAF8j9muTVoKLVtA=
 github.com/go-playground/validator/v10 v10.4.1 h1:pH2c5ADXtd66mxoE0Zm9SUhxE20r7aM3F26W0hOn+GE=
 github.com/go-playground/validator/v10 v10.4.1/go.mod h1:nlOn6nFhuKACm19sB/8EGNn9GlaMV7XkbRSipzJ0Ii4=
+github.com/google/go-cmp v0.5.9 h1:O2Tfq5qg4qc4AmwVlvv0oLiVAGB7enBSJ2x2DqQFi38=
+github.com/google/go-cmp v0.5.9/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
 github.com/leodido/go-urn v1.2.0 h1:hpXL4XnriNwQ/ABnpepYM/1vCLWNDfUNts8dX3xTG6Y=
 github.com/leodido/go-urn v1.2.0/go.mod h1:+8+nEpDfqqsY+g338gtMEUOtuK+4dEMhiQEgxpxOKII=
 github.com/mattn/go-colorable v0.1.8 h1:c1ghPdyEDarC70ftn0y+A/Ee++9zz8ljHG1b13eJ0s8=
diff --git a/option.go b/option.go
index a285a45..122a63d 100644
--- a/option.go
+++ b/option.go
@@ -169,16 +169,19 @@ func UseJSONMarshaler() EncodeOption {
 type CommentPosition int
 
 const (
-	CommentLinePosition CommentPosition = iota
-	CommentHeadPosition
+	CommentHeadPosition CommentPosition = CommentPosition(iota)
+	CommentLinePosition
+	CommentFootPosition
 )
 
 func (p CommentPosition) String() string {
 	switch p {
-	case CommentLinePosition:
-		return "Line"
 	case CommentHeadPosition:
 		return "Head"
+	case CommentLinePosition:
+		return "Line"
+	case CommentFootPosition:
+		return "Foot"
 	default:
 		return ""
 	}
@@ -200,6 +203,14 @@ func HeadComment(texts ...string) *Comment {
 	}
 }
 
+// FootComment create a multiline comment for CommentMap.
+func FootComment(texts ...string) *Comment {
+	return &Comment{
+		Texts:    texts,
+		Position: CommentFootPosition,
+	}
+}
+
 // Comment raw data for comment.
 type Comment struct {
 	Texts    []string
@@ -207,12 +218,12 @@ type Comment struct {
 }
 
 // CommentMap map of the position of the comment and the comment information.
-type CommentMap map[string]*Comment
+type CommentMap map[string][]*Comment
 
 // WithComment add a comment using the location and text information given in the CommentMap.
 func WithComment(cm CommentMap) EncodeOption {
 	return func(e *Encoder) error {
-		commentMap := map[*Path]*Comment{}
+		commentMap := map[*Path][]*Comment{}
 		for k, v := range cm {
 			path, err := PathString(k)
 			if err != nil {
diff --git a/parser/parser.go b/parser/parser.go
index 70937c9..568e6ad 100644
--- a/parser/parser.go
+++ b/parser/parser.go
@@ -268,8 +268,24 @@ func (p *parser) parseMappingValue(ctx *context) (ast.Node, error) {
 		antk = ctx.afterNextNotCommentToken()
 	}
 	if len(node.Values) == 1 {
+		mapKeyCol := mvnode.Key.GetToken().Position.Column
+		commentTk := ctx.nextToken()
+		if commentTk != nil && commentTk.Type == token.CommentType && mapKeyCol <= commentTk.Position.Column {
+			// If the comment is in the same or deeper column as the last element column in map value,
+			// treat it as a footer comment for the last element.
+			comment := p.parseFootComment(ctx, mapKeyCol)
+			mvnode.FootComment = comment
+		}
 		return mvnode, nil
 	}
+	mapCol := node.GetToken().Position.Column
+	commentTk := ctx.nextToken()
+	if commentTk != nil && commentTk.Type == token.CommentType && mapCol <= commentTk.Position.Column {
+		// If the comment is in the same or deeper column as the last element column in map value,
+		// treat it as a footer comment for the last element.
+		comment := p.parseFootComment(ctx, mapCol)
+		node.FootComment = comment
+	}
 	return node, nil
 }
 
@@ -299,9 +315,9 @@ func (p *parser) parseSequenceEntry(ctx *context) (*ast.SequenceNode, error) {
 		}
 		if comment != nil {
 			comment.SetPath(ctx.withIndex(uint(len(sequenceNode.Values))).path)
-			sequenceNode.ValueComments = append(sequenceNode.ValueComments, comment)
+			sequenceNode.ValueHeadComments = append(sequenceNode.ValueHeadComments, comment)
 		} else {
-			sequenceNode.ValueComments = append(sequenceNode.ValueComments, nil)
+			sequenceNode.ValueHeadComments = append(sequenceNode.ValueHeadComments, nil)
 		}
 		sequenceNode.Values = append(sequenceNode.Values, value)
 		tk = ctx.nextNotCommentToken()
@@ -316,6 +332,13 @@ func (p *parser) parseSequenceEntry(ctx *context) (*ast.SequenceNode, error) {
 		}
 		ctx.progressIgnoreComment(1)
 	}
+	commentTk := ctx.nextToken()
+	if commentTk != nil && commentTk.Type == token.CommentType && curColumn <= commentTk.Position.Column {
+		// If the comment is in the same or deeper column as the last element column in sequence value,
+		// treat it as a footer comment for the last element.
+		comment := p.parseFootComment(ctx, curColumn)
+		sequenceNode.FootComment = comment
+	}
 	return sequenceNode, nil
 }
 
@@ -527,6 +550,26 @@ func (p *parser) parseCommentOnly(ctx *context) *ast.CommentGroupNode {
 	return ast.CommentGroup(commentTokens)
 }
 
+func (p *parser) parseFootComment(ctx *context, col int) *ast.CommentGroupNode {
+	commentTokens := []*token.Token{}
+	for {
+		ctx.progressIgnoreComment(1)
+		commentTokens = append(commentTokens, ctx.currentToken())
+
+		nextTk := ctx.nextToken()
+		if nextTk == nil {
+			break
+		}
+		if nextTk.Type != token.CommentType {
+			break
+		}
+		if col > nextTk.Position.Column {
+			break
+		}
+	}
+	return ast.CommentGroup(commentTokens)
+}
+
 func (p *parser) parseComment(ctx *context) (ast.Node, error) {
 	group := p.parseCommentOnly(ctx)
 	node, err := p.parseToken(ctx, ctx.currentToken())
diff --git a/scanner/scanner.go b/scanner/scanner.go
index 066b795..ce9c665 100644
--- a/scanner/scanner.go
+++ b/scanner/scanner.go
@@ -484,6 +484,10 @@ func (s *Scanner) scanComment(ctx *Context) (tk *token.Token, pos int) {
 			return
 		}
 	}
+	// document ends with comment.
+	value := string(ctx.src[ctx.idx:])
+	tk = token.Comment(value, string(ctx.obuf), s.pos())
+	pos = len([]rune(value)) + 1
 	return
 }
 


Test patch:
diff --git a/yaml_test.go b/yaml_test.go
index 4e344c8..0404794 100644
--- a/yaml_test.go
+++ b/yaml_test.go
@@ -2,12 +2,13 @@ package yaml_test
 
 import (
 	"bytes"
-	"reflect"
 	"testing"
 
+	"github.com/google/go-cmp/cmp"
+	"golang.org/x/xerrors"
+
 	"github.com/goccy/go-yaml"
 	"github.com/goccy/go-yaml/ast"
-	"golang.org/x/xerrors"
 )
 
 func TestMarshal(t *testing.T) {
@@ -385,26 +386,26 @@ a: 1
 }
 
 func Test_WithCommentOption(t *testing.T) {
-	v := struct {
-		Foo string                 `yaml:"foo"`
-		Bar map[string]interface{} `yaml:"bar"`
-		Baz struct {
-			X int `yaml:"x"`
-		} `yaml:"baz"`
-	}{
-		Foo: "aaa",
-		Bar: map[string]interface{}{"bbb": "ccc"},
-		Baz: struct {
-			X int `yaml:"x"`
-		}{X: 10},
-	}
 	t.Run("line comment", func(t *testing.T) {
+		v := struct {
+			Foo string                 `yaml:"foo"`
+			Bar map[string]interface{} `yaml:"bar"`
+			Baz struct {
+				X int `yaml:"x"`
+			} `yaml:"baz"`
+		}{
+			Foo: "aaa",
+			Bar: map[string]interface{}{"bbb": "ccc"},
+			Baz: struct {
+				X int `yaml:"x"`
+			}{X: 10},
+		}
 		b, err := yaml.MarshalWithOptions(v, yaml.WithComment(
 			yaml.CommentMap{
-				"$.foo":     yaml.LineComment("foo comment"),
-				"$.bar":     yaml.LineComment("bar comment"),
-				"$.bar.bbb": yaml.LineComment("bbb comment"),
-				"$.baz.x":   yaml.LineComment("x comment"),
+				"$.foo":     []*yaml.Comment{yaml.LineComment("foo comment")},
+				"$.bar":     []*yaml.Comment{yaml.LineComment("bar comment")},
+				"$.bar.bbb": []*yaml.Comment{yaml.LineComment("bbb comment")},
+				"$.baz.x":   []*yaml.Comment{yaml.LineComment("x comment")},
 			},
 		))
 		if err != nil {
@@ -422,25 +423,120 @@ baz:
 			t.Fatalf("expected:%s but got %s", expected, actual)
 		}
 	})
-	t.Run("head comment", func(t *testing.T) {
+	t.Run("line comment2", func(t *testing.T) {
+		v := struct {
+			Foo map[string]interface{} `yaml:"foo"`
+		}{
+			Foo: map[string]interface{}{
+				"bar": map[string]interface{}{
+					"baz": true,
+				},
+			},
+		}
+		b, err := yaml.MarshalWithOptions(v, yaml.WithComment(
+			yaml.CommentMap{
+				"$.foo.bar":     []*yaml.Comment{yaml.HeadComment(" bar head comment"), yaml.LineComment(" bar line comment")},
+				"$.foo.bar.baz": []*yaml.Comment{yaml.LineComment(" baz line comment")},
+			},
+		))
+		if err != nil {
+			t.Fatal(err)
+		}
+		expected := `
+foo:
+  # bar head comment
+  bar: # bar line comment
+    baz: true # baz line comment
+`
+		actual := "\n" + string(b)
+		if expected != actual {
+			t.Fatalf("expected:%s but got %s", expected, actual)
+		}
+	})
+	t.Run("single head comment", func(t *testing.T) {
+		v := struct {
+			Foo string                 `yaml:"foo"`
+			Bar map[string]interface{} `yaml:"bar"`
+			Baz struct {
+				X int `yaml:"x"`
+			} `yaml:"baz"`
+		}{
+			Foo: "aaa",
+			Bar: map[string]interface{}{"bbb": "ccc"},
+			Baz: struct {
+				X int `yaml:"x"`
+			}{X: 10},
+		}
+
+		b, err := yaml.MarshalWithOptions(v, yaml.WithComment(
+			yaml.CommentMap{
+				"$.foo":     []*yaml.Comment{yaml.HeadComment("foo comment")},
+				"$.bar":     []*yaml.Comment{yaml.HeadComment("bar comment")},
+				"$.bar.bbb": []*yaml.Comment{yaml.HeadComment("bbb comment")},
+				"$.baz.x":   []*yaml.Comment{yaml.HeadComment("x comment")},
+			},
+		))
+		if err != nil {
+			t.Fatal(err)
+		}
+		expected := `
+#foo comment
+foo: aaa
+#bar comment
+bar:
+  #bbb comment
+  bbb: ccc
+baz:
+  #x comment
+  x: 10
+`
+		actual := "\n" + string(b)
+		if expected != actual {
+			t.Fatalf("expected:%s but got %s", expected, actual)
+		}
+	})
+
+	t.Run("multiple head comment", func(t *testing.T) {
+		v := struct {
+			Foo string                 `yaml:"foo"`
+			Bar map[string]interface{} `yaml:"bar"`
+			Baz struct {
+				X int `yaml:"x"`
+			} `yaml:"baz"`
+		}{
+			Foo: "aaa",
+			Bar: map[string]interface{}{"bbb": "ccc"},
+			Baz: struct {
+				X int `yaml:"x"`
+			}{X: 10},
+		}
+
 		b, err := yaml.MarshalWithOptions(v, yaml.WithComment(
 			yaml.CommentMap{
-				"$.foo": yaml.HeadComment(
-					"foo comment",
-					"foo comment2",
-				),
-				"$.bar": yaml.HeadComment(
-					"bar comment",
-					"bar comment2",
-				),
-				"$.bar.bbb": yaml.HeadComment(
-					"bbb comment",
-					"bbb comment2",
-				),
-				"$.baz.x": yaml.HeadComment(
-					"x comment",
-					"x comment2",
-				),
+				"$.foo": []*yaml.Comment{
+					yaml.HeadComment(
+						"foo comment",
+						"foo comment2",
+					),
+				},
+				"$.bar": []*yaml.Comment{
+					yaml.HeadComment(
+						"bar comment",
+						"bar comment2",
+					),
+				},
+				"$.bar.bbb": []*yaml.Comment{
+					yaml.HeadComment(
+						"bbb comment",
+						"bbb comment2",
+					),
+				},
+				"$.baz.x": []*yaml.Comment{
+					yaml.HeadComment(
+						"x comment",
+						"x comment2",
+					),
+				},
 			},
 		))
 		if err != nil {
@@ -466,6 +562,184 @@ baz:
 			t.Fatalf("expected:%s but got %s", expected, actual)
 		}
 	})
+	t.Run("foot comment", func(t *testing.T) {
+		v := struct {
+			Bar map[string]interface{} `yaml:"bar"`
+			Baz []int                  `yaml:"baz"`
+		}{
+			Bar: map[string]interface{}{"bbb": "ccc"},
+			Baz: []int{1, 2},
+		}
+
+		b, err := yaml.MarshalWithOptions(v, yaml.IndentSequence(true), yaml.WithComment(
+			yaml.CommentMap{
+				"$.bar.bbb": []*yaml.Comment{yaml.FootComment("ccc: ddd")},
+				"$.baz[1]":  []*yaml.Comment{yaml.FootComment("- 3")},
+				"$.baz":     []*yaml.Comment{yaml.FootComment(" foot comment", "foot comment2")},
+			},
+		))
+		if err != nil {
+			t.Fatal(err)
+		}
+		expected := `
+bar:
+  bbb: ccc
+  #ccc: ddd
+baz:
+  - 1
+  - 2
+  #- 3
+# foot comment
+#foot comment2
+`
+		actual := "\n" + string(b)
+		if expected != actual {
+			t.Fatalf("expected:%s but got %s", expected, actual)
+		}
+	})
+
+	t.Run("combination", func(t *testing.T) {
+		v := struct {
+			Foo  map[string]interface{} `yaml:"foo"`
+			O    map[string]interface{} `yaml:"o"`
+			T    map[string]bool        `yaml:"t"`
+			Bar  map[string]interface{} `yaml:"bar"`
+			Baz  []int                  `yaml:"baz"`
+			Hoge map[string]interface{} `yaml:"hoge"`
+		}{
+			Foo: map[string]interface{}{
+				"a": map[string]interface{}{
+					"b": map[string]interface{}{
+						"c": "d",
+					},
+				},
+			},
+			O: map[string]interface{}{
+				"p": map[string]interface{}{
+					"q": map[string]interface{}{
+						"r": "s",
+					},
+				},
+			},
+			T: map[string]bool{
+				"u": true,
+			},
+			Bar: map[string]interface{}{"bbb": "ccc"},
+			Baz: []int{1, 2},
+			Hoge: map[string]interface{}{
+				"moga": true,
+			},
+		}
+
+		b, err := yaml.MarshalWithOptions(v, yaml.IndentSequence(true), yaml.WithComment(
+			yaml.CommentMap{
+				"$.foo": []*yaml.Comment{
+					yaml.HeadComment(" foo head comment", " foo head comment2"),
+					yaml.LineComment(" foo line comment"),
+				},
+				"$.foo.a": []*yaml.Comment{
+					yaml.HeadComment(" a head comment"),
+					yaml.LineComment(" a line comment"),
+				},
+				"$.foo.a.b": []*yaml.Comment{
+					yaml.HeadComment(" b head comment"),
+					yaml.LineComment(" b line comment"),
+				},
+				"$.foo.a.b.c": []*yaml.Comment{
+					yaml.LineComment(" c line comment"),
+				},
+				"$.o": []*yaml.Comment{
+					yaml.LineComment(" o line comment"),
+				},
+				"$.o.p": []*yaml.Comment{
+					yaml.HeadComment(" p head comment", " p head comment2"),
+					yaml.LineComment(" p line comment"),
+				},
+				"$.o.p.q": []*yaml.Comment{
+					yaml.HeadComment(" q head comment", " q head comment2"),
+					yaml.LineComment(" q line comment"),
+				},
+				"$.o.p.q.r": []*yaml.Comment{
+					yaml.LineComment(" r line comment"),
+				},
+				"$.t.u": []*yaml.Comment{
+					yaml.LineComment(" u line comment"),
+				},
+				"$.bar": []*yaml.Comment{
+					yaml.HeadComment(" bar head comment"),
+					yaml.LineComment(" bar line comment"),
+				},
+				"$.bar.bbb": []*yaml.Comment{
+					yaml.HeadComment(" bbb head comment"),
+					yaml.LineComment(" bbb line comment"),
+					yaml.FootComment(" bbb foot comment"),
+				},
+				"$.baz[0]": []*yaml.Comment{
+					yaml.HeadComment(" sequence head comment"),
+					yaml.LineComment(" sequence line comment"),
+				},
+				"$.baz[1]": []*yaml.Comment{
+					yaml.HeadComment(" sequence head comment2"),
+					yaml.LineComment(" sequence line comment2"),
+					yaml.FootComment(" sequence foot comment"),
+				},
+				"$.baz": []*yaml.Comment{
+					yaml.HeadComment(" baz head comment", " baz head comment2"),
+					yaml.LineComment(" baz line comment"),
+					yaml.FootComment(" baz foot comment"),
+				},
+				"$.hoge.moga": []*yaml.Comment{
+					yaml.LineComment(" moga line comment"),
+					yaml.FootComment(" moga foot comment"),
+				},
+			},
+		))
+		if err != nil {
+			t.Fatal(err)
+		}
+		expected := `
+# foo head comment
+# foo head comment2
+foo: # foo line comment
+  # a head comment
+  a: # a line comment
+    # b head comment
+    b: # b line comment
+      c: d # c line comment
+o: # o line comment
+  # p head comment
+  # p head comment2
+  p: # p line comment
+    # q head comment
+    # q head comment2
+    q: # q line comment
+      r: s # r line comment
+t:
+  u: true # u line comment
+# bar head comment
+bar: # bar line comment
+  # bbb head comment
+  bbb: ccc # bbb line comment
+  # bbb foot comment
+# baz head comment
+# baz head comment2
+baz: # baz line comment
+  # sequence head comment
+  - 1 # sequence line comment
+  # sequence head comment2
+  - 2 # sequence line comment2
+  # sequence foot comment
+# baz foot comment
+hoge:
+  moga: true # moga line comment
+  # moga foot comment
+`
+		actual := "\n" + string(b)
+		if expected != actual {
+			t.Fatalf("expected:%s but got %s", expected, actual)
+		}
+	})
+
 }
 
 func Test_CommentToMapOption(t *testing.T) {
@@ -485,37 +759,118 @@ baz:
 			t.Fatal(err)
 		}
 		expected := []struct {
-			path    string
-			comment *yaml.Comment
+			path     string
+			comments []*yaml.Comment
 		}{
 			{
-				path:    "$.foo",
-				comment: yaml.LineComment("foo comment"),
+				path:     "$.foo",
+				comments: []*yaml.Comment{yaml.LineComment("foo comment")},
 			},
 			{
-				path:    "$.bar",
-				comment: yaml.LineComment("bar comment"),
+				path:     "$.bar",
+				comments: []*yaml.Comment{yaml.LineComment("bar comment")},
 			},
 			{
-				path:    "$.bar.bbb",
-				comment: yaml.LineComment("bbb comment"),
+				path:     "$.bar.bbb",
+				comments: []*yaml.Comment{yaml.LineComment("bbb comment")},
 			},
 			{
-				path:    "$.baz.x",
-				comment: yaml.LineComment("x comment"),
+				path:     "$.baz.x",
+				comments: []*yaml.Comment{yaml.LineComment("x comment")},
 			},
 		}
 		for _, exp := range expected {
-			comment := cm[exp.path]
-			if comment == nil {
+			comments := cm[exp.path]
+			if comments == nil {
 				t.Fatalf("failed to get path %s", exp.path)
 			}
-			if !reflect.DeepEqual(exp.comment, comment) {
-				t.Fatalf("failed to get comment. expected:[%+v] but got:[%+v]", exp.comment, comment)
+			if diff := cmp.Diff(exp.comments, comments); diff != "" {
+				t.Errorf("(-got, +want)\n%s", diff)
 			}
 		}
 	})
-	t.Run("head comment", func(t *testing.T) {
+	t.Run("line comment2", func(t *testing.T) {
+		yml := `
+foo:
+  bar: baz # comment`
+		var (
+			v  interface{}
+			cm = yaml.CommentMap{}
+		)
+		if err := yaml.UnmarshalWithOptions([]byte(yml), &v, yaml.CommentToMap(cm)); err != nil {
+			t.Fatal(err)
+		}
+		expected := []struct {
+			path     string
+			comments []*yaml.Comment
+		}{
+			{
+				path:     "$.foo.bar",
+				comments: []*yaml.Comment{yaml.LineComment(" comment")},
+			},
+		}
+		for _, exp := range expected {
+			comments := cm[exp.path]
+			if comments == nil {
+				t.Fatalf("failed to get path %s", exp.path)
+			}
+			if diff := cmp.Diff(exp.comments, comments); diff != "" {
+				t.Errorf("(-got, +want)\n%s", diff)
+			}
+		}
+	})
+
+	t.Run("single head comment", func(t *testing.T) {
+		yml := `
+#foo comment
+foo: aaa
+#bar comment
+bar:
+  #bbb comment
+  bbb: ccc
+baz:
+  #x comment
+  x: 10
+`
+		var (
+			v  interface{}
+			cm = yaml.CommentMap{}
+		)
+		if err := yaml.UnmarshalWithOptions([]byte(yml), &v, yaml.CommentToMap(cm)); err != nil {
+			t.Fatal(err)
+		}
+		expected := []struct {
+			path     string
+			comments []*yaml.Comment
+		}{
+			{
+				path:     "$.foo",
+				comments: []*yaml.Comment{yaml.HeadComment("foo comment")},
+			},
+			{
+				path:     "$.bar",
+				comments: []*yaml.Comment{yaml.HeadComment("bar comment")},
+			},
+			{
+				path:     "$.bar.bbb",
+				comments: []*yaml.Comment{yaml.HeadComment("bbb comment")},
+			},
+			{
+				path:     "$.baz.x",
+				comments: []*yaml.Comment{yaml.HeadComment("x comment")},
+			},
+		}
+		for _, exp := range expected {
+			comments := cm[exp.path]
+			if comments == nil {
+				t.Fatalf("failed to get path %s", exp.path)
+			}
+			if diff := cmp.Diff(exp.comments, comments); diff != "" {
+				t.Errorf("(-got, +want)\n%s", diff)
+			}
+		}
+	})
+	t.Run("multiple head comments", func(t *testing.T) {
 		yml := `
 #foo comment
 #foo comment2
@@ -539,45 +894,268 @@ baz:
 			t.Fatal(err)
 		}
 		expected := []struct {
-			path    string
-			comment *yaml.Comment
+			path     string
+			comments []*yaml.Comment
 		}{
 			{
 				path: "$.foo",
-				comment: yaml.HeadComment(
-					"foo comment",
-					"foo comment2",
-				),
+				comments: []*yaml.Comment{
+					yaml.HeadComment(
+						"foo comment",
+						"foo comment2",
+					),
+				},
 			},
 			{
 				path: "$.bar",
-				comment: yaml.HeadComment(
-					"bar comment",
-					"bar comment2",
-				),
+				comments: []*yaml.Comment{
+					yaml.HeadComment(
+						"bar comment",
+						"bar comment2",
+					),
+				},
 			},
 			{
 				path: "$.bar.bbb",
-				comment: yaml.HeadComment(
-					"bbb comment",
-					"bbb comment2",
-				),
+				comments: []*yaml.Comment{
+					yaml.HeadComment(
+						"bbb comment",
+						"bbb comment2",
+					),
+				},
 			},
 			{
 				path: "$.baz.x",
-				comment: yaml.HeadComment(
-					"x comment",
-					"x comment2",
-				),
+				comments: []*yaml.Comment{
+					yaml.HeadComment(
+						"x comment",
+						"x comment2",
+					),
+				},
+			},
+		}
+		for _, exp := range expected {
+			comments := cm[exp.path]
+			if comments == nil {
+				t.Fatalf("failed to get path %s", exp.path)
+			}
+			if diff := cmp.Diff(exp.comments, comments); diff != "" {
+				t.Errorf("(-got, +want)\n%s", diff)
+			}
+		}
+	})
+	t.Run("foot comment", func(t *testing.T) {
+		yml := `
+bar:
+  bbb: ccc
+  #ccc: ddd
+baz:
+  - 1
+  - 2
+  #- 3
+ # foot comment
+#foot comment2
+`
+		var (
+			v  interface{}
+			cm = yaml.CommentMap{}
+		)
+		if err := yaml.UnmarshalWithOptions([]byte(yml), &v, yaml.CommentToMap(cm)); err != nil {
+			t.Fatal(err)
+		}
+		expected := []struct {
+			path     string
+			comments []*yaml.Comment
+		}{
+			{
+				path:     "$.bar.bbb",
+				comments: []*yaml.Comment{yaml.FootComment("ccc: ddd")},
+			},
+			{
+				path:     "$.baz[1]",
+				comments: []*yaml.Comment{yaml.FootComment("- 3")},
+			},
+			{
+				path:     "$.baz",
+				comments: []*yaml.Comment{yaml.FootComment(" foot comment", "foot comment2")},
+			},
+		}
+		for _, exp := range expected {
+			comments := cm[exp.path]
+			if comments == nil {
+				t.Fatalf("failed to get path %s", exp.path)
+			}
+			if diff := cmp.Diff(exp.comments, comments); diff != "" {
+				t.Errorf("(-got, +want)\n%s", diff)
+			}
+		}
+	})
+	t.Run("combination", func(t *testing.T) {
+		yml := `
+# foo head comment
+# foo head comment2
+foo: # foo line comment
+  # a head comment
+  a: # a line comment
+    # b head comment
+    b: # b line comment
+      c: d # c line comment
+o: # o line comment
+  # p head comment
+  # p head comment2
+  p: # p line comment
+    # q head comment
+    # q head comment2
+    q: # q line comment
+      r: s # r line comment
+t:
+  u: true # u line comment
+# bar head comment
+bar: # bar line comment
+  # bbb head comment
+  bbb: ccc # bbb line comment
+  # bbb foot comment
+# baz head comment
+# baz head comment2
+baz: # baz line comment
+  # sequence head comment
+  - 1 # sequence line comment
+  # sequence head comment2
+  - 2 # sequence line comment2
+  # sequence foot comment
+hoge:
+  moga: true # moga line comment
+  # moga foot comment
+# hoge foot comment
+`
+		var (
+			v  interface{}
+			cm = yaml.CommentMap{}
+		)
+		if err := yaml.UnmarshalWithOptions([]byte(yml), &v, yaml.CommentToMap(cm)); err != nil {
+			t.Fatal(err)
+		}
+		expected := []struct {
+			path     string
+			comments []*yaml.Comment
+		}{
+			{
+				path: "$.foo",
+				comments: []*yaml.Comment{
+					yaml.HeadComment(" foo head comment", " foo head comment2"),
+					yaml.LineComment(" foo line comment"),
+				},
+			},
+			{
+				path: "$.foo.a",
+				comments: []*yaml.Comment{
+					yaml.HeadComment(" a head comment"),
+					yaml.LineComment(" a line comment"),
+				},
+			},
+			{
+				path: "$.foo.a.b",
+				comments: []*yaml.Comment{
+					yaml.HeadComment(" b head comment"),
+					yaml.LineComment(" b line comment"),
+				},
+			},
+			{
+				path: "$.foo.a.b.c",
+				comments: []*yaml.Comment{
+					yaml.LineComment(" c line comment"),
+				},
+			},
+			{
+				path: "$.o",
+				comments: []*yaml.Comment{
+					yaml.LineComment(" o line comment"),
+				},
+			},
+			{
+				path: "$.o.p",
+				comments: []*yaml.Comment{
+					yaml.HeadComment(" p head comment", " p head comment2"),
+					yaml.LineComment(" p line comment"),
+				},
+			},
+			{
+				path: "$.o.p.q",
+				comments: []*yaml.Comment{
+					yaml.HeadComment(" q head comment", " q head comment2"),
+					yaml.LineComment(" q line comment"),
+				},
+			},
+			{
+				path: "$.o.p.q.r",
+				comments: []*yaml.Comment{
+					yaml.LineComment(" r line comment"),
+				},
+			},
+			{
+				path: "$.t.u",
+				comments: []*yaml.Comment{
+					yaml.LineComment(" u line comment"),
+				},
+			},
+			{
+				path: "$.bar",
+				comments: []*yaml.Comment{
+					yaml.HeadComment(" bar head comment"),
+					yaml.LineComment(" bar line comment"),
+				},
+			},
+			{
+				path: "$.bar.bbb",
+				comments: []*yaml.Comment{
+					yaml.HeadComment(" bbb head comment"),
+					yaml.LineComment(" bbb line comment"),
+					yaml.FootComment(" bbb foot comment"),
+				},
+			},
+			{
+				path: "$.baz[0]",
+				comments: []*yaml.Comment{
+					yaml.HeadComment(" sequence head comment"),
+					yaml.LineComment(" sequence line comment"),
+				},
+			},
+			{
+				path: "$.baz[1]",
+				comments: []*yaml.Comment{
+					yaml.HeadComment(" sequence head comment2"),
+					yaml.LineComment(" sequence line comment2"),
+					yaml.FootComment(" sequence foot comment"),
+				},
+			},
+			{
+				path: "$.baz",
+				comments: []*yaml.Comment{
+					yaml.HeadComment(" baz head comment", " baz head comment2"),
+					yaml.LineComment(" baz line comment"),
+				},
+			},
+			{
+				path: "$.hoge",
+				comments: []*yaml.Comment{
+					yaml.FootComment(" hoge foot comment"),
+				},
+			},
+			{
+				path: "$.hoge.moga",
+				comments: []*yaml.Comment{
+					yaml.LineComment(" moga line comment"),
+					yaml.FootComment(" moga foot comment"),
+				},
 			},
 		}
 		for _, exp := range expected {
-			comment := cm[exp.path]
-			if comment == nil {
+			comments := cm[exp.path]
+			if comments == nil {
 				t.Fatalf("failed to get path %s", exp.path)
 			}
-			if !reflect.DeepEqual(exp.comment, comment) {
-				t.Fatalf("failed to get comment. expected:[%+v] but got:[%+v]", exp.comment, comment)
+			if diff := cmp.Diff(exp.comments, comments); diff != "" {
+				t.Errorf("%s: (-got, +want)\n%s", exp.path, diff)
 			}
 		}
 	})