From 81c46ed0c726c5ee6028a1a3d0f0056c346efcf5 Mon Sep 17 00:00:00 2001
From: Masaaki Goshima <goccy54@gmail.com>
Date: Tue, 5 Nov 2024 13:24:27 +0900
Subject: [PATCH] fix parsing of flow mapping

---
 parser/context.go     |  9 +++++-
 parser/parser.go      | 65 +++++++++++++++++++++++++++++++++++++++----
 parser/parser_test.go |  4 ++-
 3 files changed, 70 insertions(+), 8 deletions(-)

diff --git a/parser/context.go b/parser/context.go
index 39b23cb..cc7d302 100644
--- a/parser/context.go
+++ b/parser/context.go
@@ -7,7 +7,8 @@ import (
 
 // context context at parsing
 type context struct {
-	path string
+	path   string
+	isFlow bool
 }
 
 var pathSpecialChars = []string{
@@ -42,6 +43,12 @@ func (c *context) withIndex(idx uint) *context {
 	return &ctx
 }
 
+func (c *context) withFlow(isFlow bool) *context {
+	ctx := *c
+	ctx.isFlow = isFlow
+	return &ctx
+}
+
 func newContext() *context {
 	return &context{
 		path: "$",
diff --git a/parser/parser.go b/parser/parser.go
index 5d7bfb3..65bff39 100644
--- a/parser/parser.go
+++ b/parser/parser.go
@@ -140,17 +140,27 @@ func (p *parser) parseMapping(ctx *context) (*ast.MappingNode, error) {
 	node := ast.Mapping(mapTk, true)
 	node.SetPath(ctx.path)
 	p.progress(1) // skip MappingStart token
+
+	isFirst := true
 	for p.next() {
 		tk := p.currentToken()
 		if tk.Type == token.MappingEndType {
 			node.End = tk
-			return node, nil
+			break
 		} else if tk.Type == token.CollectEntryType {
 			p.progress(1)
-			continue
+		} else if !isFirst {
+			return nil, errors.ErrSyntax("',' or '}' must be specified", tk)
+		}
+
+		if tk := p.currentToken(); tk != nil && tk.Type == token.MappingEndType {
+			// this case is here: "{ elem, }".
+			// In this case, ignore the last element and break mapping parsing.
+			node.End = tk
+			break
 		}
 
-		value, err := p.parseMappingValue(ctx)
+		value, err := p.parseMappingValue(ctx.withFlow(true))
 		if err != nil {
 			return nil, err
 		}
@@ -160,8 +170,12 @@ func (p *parser) parseMapping(ctx *context) (*ast.MappingNode, error) {
 		}
 		node.Values = append(node.Values, mvnode)
 		p.progress(1)
+		isFirst = false
+	}
+	if node.End == nil || node.End.Type != token.MappingEndType {
+		return nil, errors.ErrSyntax("could not find flow mapping end token '}'", node.Start)
 	}
-	return nil, errors.ErrSyntax("unterminated flow mapping", node.GetToken())
+	return node, nil
 }
 
 func (p *parser) parseSequence(ctx *context) (*ast.SequenceNode, error) {
@@ -382,12 +396,32 @@ func (p *parser) parseMappingValue(ctx *context) (ast.Node, error) {
 	if err := p.validateMapKey(key.GetToken()); err != nil {
 		return nil, err
 	}
-	p.progress(1)          // progress to mapping value token
-	tk := p.currentToken() // get mapping value token
+	p.progress(1) // progress to mapping value token
+	if ctx.isFlow {
+		// if "{key}" or "{key," style, returns MappingValueNode.
+		node, err := p.parseFlowMapNullValue(ctx, key)
+		if err != nil {
+			return nil, err
+		}
+		if node != nil {
+			return node, nil
+		}
+	}
+	tk := p.currentToken() // get mapping value (':') token.
 	if tk == nil {
 		return nil, errors.ErrSyntax("unexpected map", key.GetToken())
 	}
 	p.progress(1) // progress to value token
+	if ctx.isFlow {
+		// if "{key:}" or "{key:," style, returns MappingValueNode.
+		node, err := p.parseFlowMapNullValue(ctx, key)
+		if err != nil {
+			return nil, err
+		}
+		if node != nil {
+			return node, nil
+		}
+	}
 	if err := p.setSameLineCommentIfExists(ctx.withChild(keyText), key); err != nil {
 		return nil, err
 	}
@@ -474,6 +508,25 @@ func (p *parser) parseMappingValue(ctx *context) (ast.Node, error) {
 	return node, nil
 }
 
+func (p *parser) parseFlowMapNullValue(ctx *context, key ast.MapKeyNode) (*ast.MappingValueNode, error) {
+	tk := p.currentToken()
+	if tk == nil {
+		return nil, errors.ErrSyntax("unexpected map", key.GetToken())
+	}
+	if tk.Type != token.MappingEndType && tk.Type != token.CollectEntryType {
+		return nil, nil
+	}
+	nullTk := p.createNullToken(tk)
+	p.insertToken(p.idx, nullTk)
+	value, err := p.parseToken(ctx, nullTk)
+	if err != nil {
+		return nil, err
+	}
+	node := ast.MappingValue(tk, key, value)
+	node.SetPath(ctx.withChild(key.GetToken().Value).path)
+	return node, nil
+}
+
 func (p *parser) parseSequenceEntry(ctx *context) (*ast.SequenceNode, error) {
 	tk := p.currentToken()
 	sequenceNode := ast.Sequence(tk, false)
diff --git a/parser/parser_test.go b/parser/parser_test.go
index 7adfe77..ae2fe63 100644
--- a/parser/parser_test.go
+++ b/parser/parser_test.go
@@ -94,6 +94,8 @@ func TestParser(t *testing.T) {
 		"value: >\nother:",
 		"value: >\n\nother:",
 		"a:\n-",
+		"a: {foo}",
+		"a: {foo,bar}",
 	}
 	for _, src := range sources {
 		if _, err := parser.Parse(lexer.Tokenize(src), 0); err != nil {
@@ -968,7 +970,7 @@ a
 		{
 			`{ "key": "value" `,
 			`
-[1:1] unterminated flow mapping
+[1:1] could not find flow mapping end token '}'
 >  1 | { "key": "value"
        ^
 `,