Repository: goccy/go-yaml
Pull request #551 (2024-11-28 08:25:50)

Fix parsing of document header option
fix https://github.com/goccy/go-yaml/issues/550


codecov-commenter: :warning: Please install the !['codecov app svg image'](https://github.com/codecov/engineering-team/assets/152432831/e90313f4-9d3a-4b63-8b54-cfe14e7ec20d) to ensure uploads and comments are reliably processed by Codecov.

## [Codecov](https://app.codecov.io/gh/goccy/go-yaml/pull/551?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Masaaki+Goshima) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 77.18%. Comparing base [(`3584ab7`)](https://app.codecov.io/gh/goccy/go-yaml/commit/3584ab7d9c402e6bada57a1490de4442c1755c00?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Masaaki+Goshima) to head [(`a924369`)](https://app.codecov.io/gh/goccy/go-yaml/commit/a92436921e74ca0d65c934c3c2c0b4d7bec26e51?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Masaaki+Goshima).

:exclamation: Your organization needs to install the [Codecov GitHub app](https://github.com/apps/codecov/installations/select_target) to enable full functionality.

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master     #551      +/-   ##
==========================================
- Coverage   77.20%   77.18%   -0.02%     
==========================================
  Files          21       21              
  Lines        7215     7181      -34     
==========================================
- Hits         5570     5543      -27     
+ Misses       1251     1245       -6     
+ Partials      394      393       -1     
```

</details>

Patch:
diff --git a/scanner/scanner.go b/scanner/scanner.go
index 26777be..f702f6d 100644
--- a/scanner/scanner.go
+++ b/scanner/scanner.go
@@ -625,22 +625,6 @@ func (s *Scanner) scanComment(ctx *Context) bool {
 	return true
 }
 
-func (s *Scanner) trimCommentFromDocumentOpt(text string, header rune) (string, error) {
-	idx := strings.Index(text, "#")
-	if idx < 0 {
-		return text, nil
-	}
-	if idx == 0 {
-		return "", ErrInvalidToken(
-			token.Invalid(
-				fmt.Sprintf("invalid document header %s", text),
-				string(header)+text, s.pos(),
-			),
-		)
-	}
-	return text[:idx-1], nil
-}
-
 func (s *Scanner) scanDocument(ctx *Context, c rune) error {
 	ctx.addOriginBuf(c)
 	if ctx.isEOS() {
@@ -982,63 +966,45 @@ func (s *Scanner) scanDocumentHeaderOption(ctx *Context) error {
 		ctx.addOriginBuf(c)
 		switch c {
 		case '\n', '\r':
-			value := ctx.source(ctx.idx, ctx.idx+idx)
-			opt := strings.TrimRight(value, " ")
-			orgOptLen := len(opt)
-			opt, err := s.trimCommentFromDocumentOpt(opt, header)
-			if err != nil {
-				return err
+			value := strings.TrimRight(ctx.source(ctx.idx, ctx.idx+idx), " ")
+			commentValueIndex := strings.Index(value, "#")
+			opt := value
+			if commentValueIndex > 0 {
+				opt = value[:commentValueIndex]
 			}
-			if err := s.validateDocumentHeaderOption(opt); err != nil {
-				invalidTk := token.Invalid(err.Error(), string(ctx.obuf), s.pos())
-				s.progressColumn(ctx, progress)
-				return ErrInvalidToken(invalidTk)
+			opt = strings.TrimRightFunc(opt, func(r rune) bool {
+				return r == ' ' || r == '\t'
+			})
+			if len(opt) != 0 {
+				if err := s.validateDocumentHeaderOption(opt); err != nil {
+					invalidTk := token.Invalid(err.Error(), string(ctx.obuf), s.pos())
+					s.progressColumn(ctx, progress)
+					return ErrInvalidToken(invalidTk)
+				}
 			}
-			hasComment := len(opt) < orgOptLen
 			if s.column == 1 {
 				s.lastDelimColumn = 1
 			}
-			if header == '|' {
-				if hasComment {
-					commentLen := orgOptLen - len(opt)
-					headerPos := strings.Index(string(ctx.obuf), "|")
-					if len(ctx.obuf) < commentLen+headerPos {
-						invalidTk := token.Invalid("found invalid literal header option", string(ctx.obuf), s.pos())
-						s.progressColumn(ctx, progress)
-						return ErrInvalidToken(invalidTk)
-					}
-					litBuf := ctx.obuf[:len(ctx.obuf)-commentLen-headerPos]
-					commentBuf := ctx.obuf[len(litBuf):]
-					ctx.addToken(token.Literal("|"+opt, string(litBuf), s.pos()))
-					s.column += len(litBuf)
-					s.offset += len(litBuf)
-					commentHeader := strings.Index(value, "#")
-					ctx.addToken(token.Comment(string(value[commentHeader+1:]), string(commentBuf), s.pos()))
-				} else {
-					ctx.addToken(token.Literal("|"+opt, string(ctx.obuf), s.pos()))
-				}
+
+			commentIndex := strings.Index(string(ctx.obuf), "#")
+			headerBuf := string(ctx.obuf)
+			if commentIndex > 0 {
+				headerBuf = headerBuf[:commentIndex]
+			}
+			switch header {
+			case '|':
+				ctx.addToken(token.Literal("|"+opt, headerBuf, s.pos()))
 				ctx.isLiteral = true
-			} else if header == '>' {
-				if hasComment {
-					commentLen := orgOptLen - len(opt)
-					headerPos := strings.Index(string(ctx.obuf), ">")
-					if len(ctx.obuf) < commentLen+headerPos {
-						invalidTk := token.Invalid("found invalid folded header option", string(ctx.obuf), s.pos())
-						s.progressColumn(ctx, progress)
-						return ErrInvalidToken(invalidTk)
-					}
-					foldedBuf := ctx.obuf[:len(ctx.obuf)-commentLen-headerPos]
-					commentBuf := ctx.obuf[len(foldedBuf):]
-					ctx.addToken(token.Folded(">"+opt, string(foldedBuf), s.pos()))
-					s.column += len(foldedBuf)
-					s.offset += len(foldedBuf)
-					commentHeader := strings.Index(value, "#")
-					ctx.addToken(token.Comment(string(value[commentHeader+1:]), string(commentBuf), s.pos()))
-				} else {
-					ctx.addToken(token.Folded(">"+opt, string(ctx.obuf), s.pos()))
-				}
+			case '>':
+				ctx.addToken(token.Folded(">"+opt, headerBuf, s.pos()))
 				ctx.isFolded = true
 			}
+			if commentIndex > 0 {
+				comment := string(value[commentValueIndex+1:])
+				s.offset += len(headerBuf)
+				s.column += len(headerBuf)
+				ctx.addToken(token.Comment(comment, string(ctx.obuf[len(headerBuf):]), s.pos()))
+			}
 			s.indentState = IndentStateKeep
 			ctx.resetBuffer()
 			ctx.docOpt = opt


Test patch:
diff --git a/lexer/lexer_test.go b/lexer/lexer_test.go
index b795041..7179517 100644
--- a/lexer/lexer_test.go
+++ b/lexer/lexer_test.go
@@ -2201,7 +2201,7 @@ s: >
 		},
 		{
 			YAML: `
-s: >1
+s: >1        # comment
         1s
 `,
 			Tokens: token.Tokens{
@@ -2224,7 +2224,14 @@ s: >1
 					CharacterType: token.CharacterTypeIndicator,
 					Indicator:     token.BlockScalarIndicator,
 					Value:         ">1",
-					Origin:        " >1\n",
+					Origin:        " >1        ",
+				},
+				{
+					Type:          token.CommentType,
+					CharacterType: token.CharacterTypeIndicator,
+					Indicator:     token.CommentIndicator,
+					Value:         " comment",
+					Origin:        "# comment\n",
 				},
 				{
 					Type:          token.StringType,
@@ -2510,6 +2517,35 @@ s: >-3
 				},
 			},
 		},
+		{
+			YAML: `
+|  		  # comment
+  foo
+`,
+			Tokens: token.Tokens{
+				{
+					Type:          token.LiteralType,
+					CharacterType: token.CharacterTypeIndicator,
+					Indicator:     token.BlockScalarIndicator,
+					Value:         "|",
+					Origin: "\n|  		  ", //nolint:gci,gofmt
+				},
+				{
+					Type:          token.CommentType,
+					CharacterType: token.CharacterTypeIndicator,
+					Indicator:     token.CommentIndicator,
+					Value:         " comment",
+					Origin:        "# comment\n",
+				},
+				{
+					Type:          token.StringType,
+					CharacterType: token.CharacterTypeMiscellaneous,
+					Indicator:     token.NotIndicator,
+					Value:         "foo\n",
+					Origin:        "  foo\n",
+				},
+			},
+		},
 		{
 			YAML: `1x0`,
 			Tokens: token.Tokens{