Support line comment for flow sequence or flow map (#834)

* support line comment for flow sequence or flow map

* fix lint error

* fix bytes unmarshaler for flow sequence or flow map with line comment

Author: goccy

ast/ast.go

@@ -1623,7 +1623,11 @@ func (n *SequenceNode) flowStyleString() string {
 	for _, value := range n.Values {
 		values = append(values, value.String())
 	}
-	return fmt.Sprintf("[%s]", strings.Join(values, ", "))
+	seqText := fmt.Sprintf("[%s]", strings.Join(values, ", "))
+	if n.Comment != nil {
+		return addCommentString(seqText, n.Comment)
+	}
+	return seqText
 }
 
 func (n *SequenceNode) blockStyleString() string {

decode.go

@@ -288,7 +288,9 @@ func (d *Decoder) addSequenceNodeCommentToMap(node *ast.SequenceNode) {
 			texts = append(texts, comment.Token.Value)
 		}
 		if len(texts) != 0 {
-			d.addCommentToMap(node.Values[0].GetPath(), HeadComment(texts...))
+			if len(node.Values) != 0 {
+				d.addCommentToMap(node.Values[0].GetPath(), HeadComment(texts...))
+			}
 		}
 	}
 }

decode_test.go

@@ -2955,7 +2955,10 @@ func (u *unmarshalList) UnmarshalYAML(b []byte) error {
 
     hello
   f: g
-- h: i`
+- h: i
+- j: [] # comment
+- k: {} # comment
+`
 	actual := "\n" + string(b)
 	if expected != actual {
 		return fmt.Errorf("unexpected bytes: expected [%q] but got [%q]", expected, actual)
@@ -3007,6 +3010,8 @@ a:
      hello
    f: g
  - h: i
+ - j: [] # comment
+ - k: {} # comment
 `
 	var v struct {
 		A unmarshalList
@@ -3015,7 +3020,7 @@ a:
 	if err := yaml.UnmarshalWithOptions([]byte(yml), &v, yaml.CommentToMap(cm)); err != nil {
 		t.Fatal(err)
 	}
-	if len(v.A.v) != 2 {
+	if len(v.A.v) != 4 {
 		t.Fatalf("failed to unmarshal %+v", v)
 	}
 	if len(v.A.v[0]) != 3 {

internal/format/format.go

@@ -351,8 +351,9 @@ func (f *Formatter) formatMapping(n *ast.MappingNode) string {
 	var ret string
 	if n.IsFlowStyle {
 		ret = f.origin(n.Start)
+	} else {
+		ret += f.formatCommentGroup(n.Comment)
 	}
-	ret += f.formatCommentGroup(n.Comment)
 	for _, value := range n.Values {
 		if value.CollectEntry != nil {
 			ret += f.origin(value.CollectEntry)
@@ -361,6 +362,7 @@ func (f *Formatter) formatMapping(n *ast.MappingNode) string {
 	}
 	if n.IsFlowStyle {
 		ret += f.origin(n.End)
+		ret += f.formatCommentGroup(n.Comment)
 	}
 	return ret
 }
@@ -377,8 +379,7 @@ func (f *Formatter) formatSequence(n *ast.SequenceNode) string {
 	var ret string
 	if n.IsFlowStyle {
 		ret = f.origin(n.Start)
-	}
-	if n.Comment != nil {
+	} else {
 		// add head comment.
 		ret += f.formatCommentGroup(n.Comment)
 	}
@@ -387,6 +388,7 @@ func (f *Formatter) formatSequence(n *ast.SequenceNode) string {
 	}
 	if n.IsFlowStyle {
 		ret += f.origin(n.End)
+		ret += f.formatCommentGroup(n.Comment)
 	}
 	ret += f.formatCommentGroup(n.FootComment)
 	return ret

parser/parser.go

@@ -426,6 +426,11 @@ func (p *parser) parseFlowMap(ctx *context) (*ast.MappingNode, error) {
 	if node.End == nil {
 		return nil, errors.ErrSyntax("could not find flow mapping end token '}'", node.Start)
 	}
+
+	// set line comment if exists. e.g.) } # comment
+	if err := setLineComment(ctx, node, ctx.currentToken()); err != nil {
+		return nil, err
+	}
 	ctx.goNext() // skip mapping end token.
 	return node, nil
 }
@@ -1066,6 +1071,11 @@ func (p *parser) parseFlowSequence(ctx *context) (*ast.SequenceNode, error) {
 	if node.End == nil {
 		return nil, errors.ErrSyntax("sequence end token ']' not found", node.Start)
 	}
+
+	// set line comment if exists. e.g.) ] # comment
+	if err := setLineComment(ctx, node, ctx.currentToken()); err != nil {
+		return nil, err
+	}
 	ctx.goNext() // skip sequence end token.
 	return node, nil
 }

parser/parser_test.go

@@ -1570,6 +1570,22 @@ elem1:
 elem1:
   - elem2: # comment
       [a, b, c, d]
+`,
+		},
+		{
+			name: "flow map with inline value comment",
+			yaml: `
+a:
+  b: {} # comment
+c: d
+`,
+		},
+		{
+			name: "flow array with inline value comment",
+			yaml: `
+a:
+  b: [] # comment
+c: d
 `,
 		},
 		{