feat(goccy): marshal back to goccy ast

Author: LarsStegman

pkg/yqlib/candidate_node_goccy_yaml.go

@@ -180,82 +180,52 @@ func (o *CandidateNode) goccyProcessMappingValueNode(mappingEntry *ast.MappingVa
 	return nil
 }
 
-func (o *CandidateNode) goccyMarshalChild(path string) (ast.Node, error) {
-	var node ast.Node
-	switch o.Kind {
-	case SequenceNode:
-		values := make([]ast.Node, len(o.Content))
-		for i, content := range o.Content {
-			childNode, err := content.goccyMarshalChild(path + "[" + strconv.Itoa(i) + "]")
-			if err != nil {
-				return nil, err
-			}
-
-			values[i] = childNode
-		}
-
-		seq := ast.Sequence(
-			goccyToken.SequenceStart(o.Value, &goccyToken.Position{}),
-			o.Style == FlowStyle,
-		)
-
-		seq.Values = values
-		seq.ValueHeadComments = []*ast.CommentGroupNode{ast.CommentGroup([]*goccyToken.Token{goccyToken.Comment(o.HeadComment, o.HeadComment, &goccyToken.Position{})})}
-		seq.FootComment = goccyMarshalFootComment(o.FootComment)
-		node = seq
-	case MappingNode:
-		values := make([]*ast.MappingValueNode, len(o.Content)/2)
-		for i := 0; i < len(o.Content); i += 2 {
-			key, err := o.Content[i].goccyMarshalMapKeyNode(path)
-			if err != nil {
-				return nil, err
-			}
+func (o *CandidateNode) MarshalGoccyYAML() (ast.Node, error) {
+	marshaller := &goccyMarshaller{}
+	return marshaller.Marshal("$", o)
+}
 
-			value, err := o.Content[i+1].goccyMarshalChild(key.GetPath())
-			if err != nil {
-				return nil, err
-			}
+type goccyMarshaller struct {
+	currentToken *goccyToken.Token
+}
 
-			childNode := ast.MappingValue(
-				goccyToken.MappingValue(&goccyToken.Position{}),
-				key,
-				value,
-			)
-			childNode.SetPath(key.GetPath())
-			if o.Content[i].FootComment != "" {
-				childNode.FootComment = goccyMarshalFootComment(o.Content[i].FootComment)
-			}
+func (m *goccyMarshaller) PushToken(t *goccyToken.Token) *goccyToken.Token {
+	if m.currentToken == nil {
+		m.currentToken = t
+		return m.currentToken
+	}
 
-			values[i/2] = childNode
-		}
+	m.currentToken.Next = t
+	t.Prev = m.currentToken
+	m.currentToken = t
 
-		mapping := ast.Mapping(goccyToken.MappingStart("", &goccyToken.Position{}), o.Style == FlowStyle, values...)
-		if o.FootComment != "" {
-			mapping.FootComment = goccyMarshalFootComment(o.FootComment)
-		}
+	return m.currentToken
+}
 
-		node = mapping
-	case ScalarNode:
-		v, err := o.GetValueRep()
+func (m *goccyMarshaller) Marshal(path string, o *CandidateNode) (ast.Node, error) {
+	var node ast.Node
+	var err error
+	switch o.Kind {
+	case AliasNode:
+		node, err = m.marshalAlias(path, o)
 		if err != nil {
 			return nil, err
 		}
-
-		node, err = o.goccyMarshalScalar(path, v)
+	case ScalarNode:
+		node, err = m.marshalScalar(path, o)
 		if err != nil {
 			return nil, err
 		}
-	case AliasNode:
-		value, err := o.Content[0].goccyMarshalChild(path)
+	case SequenceNode:
+		node, err = m.marshalSequence(path, o)
 		if err != nil {
 			return nil, err
 		}
-
-		alias := ast.Alias(goccyToken.Alias(o.Value, &goccyToken.Position{}))
-		alias.Value = value
-		node = alias
+	case MappingNode:
+		node, err = m.marshalMapping(path, o)
 	}
 
+	// TODO: Fix Tag and Anchor tokens
 	if o.Tag != "" && !strings.HasPrefix(o.Tag, "!!") {
 		tag := ast.Tag(goccyToken.Tag(o.Tag, o.Tag, &goccyToken.Position{}))
 		tag.Value = node
@@ -273,25 +243,136 @@ func (o *CandidateNode) goccyMarshalChild(path string) (ast.Node, error) {
 	return node, nil
 }
 
-func (o *CandidateNode) MarshalGoccyYAML() (ast.Node, error) {
-	return o.goccyMarshalChild("$")
+func (m *goccyMarshaller) marshalAlias(path string, o *CandidateNode) (ast.Node, error) {
+	t := m.PushToken(goccyToken.Alias(o.Value, &goccyToken.Position{}))
+	var value ast.Node
+	var err error
+	value, err = m.Marshal(path, o.Content[0])
+	if err != nil {
+		return nil, err
+	}
+
+	alias := ast.Alias(t)
+	alias.Value = value
+	return alias, nil
 }
 
-func goccyMarshalFootComment(comment string) *ast.CommentGroupNode {
-	return ast.CommentGroup([]*goccyToken.Token{goccyToken.Comment(comment, comment, &goccyToken.Position{})})
+func (m *goccyMarshaller) marshalSequence(path string, o *CandidateNode) (ast.Node, error) {
+	isFlowStyle := o.Style == FlowStyle
+	seq := ast.Sequence(
+		nil,
+		isFlowStyle,
+	)
+
+	if isFlowStyle {
+		seq.Start = m.PushToken(goccyToken.SequenceStart("[", &goccyToken.Position{}))
+	}
+
+	var err error
+	values := make([]ast.Node, len(o.Content))
+	for i, content := range o.Content {
+		m.PushToken(goccyToken.SequenceEntry("-", &goccyToken.Position{}))
+		values[i], err = m.Marshal(path+"["+strconv.Itoa(i)+"]", content)
+		if err != nil {
+			return nil, err
+		}
+
+		if isFlowStyle && i < len(o.Content)-1 {
+			m.PushToken(goccyToken.CollectEntry(",", &goccyToken.Position{}))
+		}
+	}
+
+	if isFlowStyle {
+		seq.End = m.PushToken(goccyToken.SequenceEnd("]", &goccyToken.Position{}))
+	}
+
+	seq.Values = values
+	seq.ValueHeadComments = []*ast.CommentGroupNode{ast.CommentGroup([]*goccyToken.Token{goccyToken.Comment(o.HeadComment, o.HeadComment, &goccyToken.Position{})})}
+	seq.FootComment = goccyMarshalFootComment(o.FootComment)
+
+	return seq, nil
 }
 
-func (o *CandidateNode) goccyMarshalMapKeyNode(path string) (ast.MapKeyNode, error) {
-	if o.Kind != ScalarNode {
-		return nil, fmt.Errorf("cannot unmarshal non-scalar node to MapKeyNode")
+func (m *goccyMarshaller) marshalMapping(path string, o *CandidateNode) (ast.Node, error) {
+	isFlowStyle := o.Style == FlowStyle
+	mapping := ast.Mapping(nil, isFlowStyle)
+
+	if isFlowStyle {
+		mapping.Start = m.PushToken(goccyToken.MappingStart("{", &goccyToken.Position{}))
 	}
 
-	v, err := o.GetValueRep()
+	values := make([]*ast.MappingValueNode, len(o.Content)/2)
+	var err error
+	for i := 0; i < len(o.Content); i += 2 {
+		var childNode *ast.MappingValueNode
+		childNode, err = m.marshalMappingValueNode(path, o.Content[i], o.Content[i+1])
+		if err != nil {
+			return nil, err
+		}
+
+		if o.Content[i].FootComment != "" {
+			childNode.FootComment = goccyMarshalFootComment(o.Content[i].FootComment)
+		}
+
+		values[i/2] = childNode
+		if isFlowStyle && i < len(o.Content)/2-1 {
+			m.PushToken(goccyToken.CollectEntry(",", &goccyToken.Position{}))
+		}
+	}
+
+	mapping.Values = values
+
+	if !isFlowStyle && len(values) > 0 {
+		mapping.Start = values[0].Start
+	}
+
+	if isFlowStyle {
+		mapping.End = m.PushToken(goccyToken.MappingEnd("}", &goccyToken.Position{}))
+	}
+
+	if o.FootComment != "" {
+		mapping.FootComment = goccyMarshalFootComment(o.FootComment)
+	}
+
+	return mapping, nil
+}
+
+func (m *goccyMarshaller) marshalMappingValueNode(path string, keyNode, valueNode *CandidateNode) (*ast.MappingValueNode, error) {
+	var key ast.MapKeyNode
+	var err error
+	key, err = m.marshalMapKeyNode(path, keyNode)
 	if err != nil {
 		return nil, err
 	}
 
-	key, err := o.goccyMarshalScalar(path, v)
+	separatorToken := m.PushToken(goccyToken.MappingValue(&goccyToken.Position{}))
+
+	value, err := m.Marshal(key.GetPath(), valueNode)
+	if err != nil {
+		return nil, err
+	}
+
+	childNode := ast.MappingValue(
+		separatorToken,
+		key,
+		value,
+	)
+	childNode.SetPath(key.GetPath())
+
+	return childNode, nil
+}
+
+func goccyMarshalFootComment(comment string) *ast.CommentGroupNode {
+	// TODO: Fix Tokens
+	return ast.CommentGroup([]*goccyToken.Token{goccyToken.Comment(comment, comment, &goccyToken.Position{})})
+}
+
+func (m *goccyMarshaller) marshalMapKeyNode(path string, o *CandidateNode) (ast.MapKeyNode, error) {
+	if o.Kind != ScalarNode {
+		return nil, fmt.Errorf("cannot unmarshal non-scalar node to MapKeyNode")
+	}
+
+	key, err := m.marshalScalar(path, o)
 	if err != nil {
 		return nil, err
 	}
@@ -300,30 +381,35 @@ func (o *CandidateNode) goccyMarshalMapKeyNode(path string) (ast.MapKeyNode, err
 	return key, nil
 }
 
-func (o *CandidateNode) goccyMarshalScalar(path string, v any) (ast.ScalarNode, error) {
-	t := goccyToken.Literal(o.Value, o.Value, &goccyToken.Position{})
+func (m *goccyMarshaller) marshalScalar(path string, o *CandidateNode) (ast.ScalarNode, error) {
+	v, err := o.GetValueRep()
+	if err != nil {
+		return nil, err
+	}
+
+	scalarToken := m.PushToken(goccyToken.Literal(o.Value, o.Value, &goccyToken.Position{}))
 
 	if v == nil {
-		n := ast.Null(t)
+		n := ast.Null(scalarToken)
 		n.SetPath(path)
 		return n, nil
 	}
 
 	switch av := v.(type) {
 	case float64:
-		n := ast.Float(t)
+		n := ast.Float(scalarToken)
 		n.Value = av
 		return n, nil
 	case int64:
-		n := ast.Integer(t)
+		n := ast.Integer(scalarToken)
 		n.Value = av
 		return n, nil
 	case bool:
-		n := ast.Bool(t)
+		n := ast.Bool(scalarToken)
 		n.Value = av
 		return n, nil
 	case string:
-		n := ast.String(t)
+		n := ast.String(scalarToken)
 		n.SetPath(path)
 		return n, nil
 	}

pkg/yqlib/candidate_node_goccy_yaml_test.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"testing"
 
+	"github.com/goccy/go-yaml/parser"
 	"gopkg.in/yaml.v3"
 )
 
@@ -17,10 +18,15 @@ a:
 s: 
 - 1
 - 2
+t: [5, five]
+f: [6, {y: true}]
 `
 
+	goccyAst, err := parser.ParseBytes([]byte(input), parser.ParseComments)
+	fmt.Println(goccyAst)
+
 	var yamlNode yaml.Node
-	err := yaml.Unmarshal([]byte(input), &yamlNode)
+	err = yaml.Unmarshal([]byte(input), &yamlNode)
 	if err != nil {
 		t.Error(err)
 		return
@@ -39,5 +45,12 @@ s:
 		return
 	}
 
+	parsed, err := parser.ParseBytes([]byte(input), parser.ParseComments)
+	if err != nil {
+		t.Error(err)
+		return
+	}
+
+	fmt.Println(parsed)
 	fmt.Println(goccyNode)
 }