From 3544475f1255026b60696ed7fca4c27fa86d05d2 Mon Sep 17 00:00:00 2001
From: Masaaki Goshima <goccy54@gmail.com>
Date: Thu, 9 Jan 2020 13:33:10 +0900
Subject: [PATCH] Support MarshalAnchor option for encoder

---
 encode.go      | 45 ++++++++++++++++++----------
 encode_test.go | 80 ++++++++++++++++++++++++++++++++++++++++++++++++++
 option.go      | 14 ++++++++-
 3 files changed, 122 insertions(+), 17 deletions(-)

diff --git a/encode.go b/encode.go
index f7c139f4..413942eb 100644
--- a/encode.go
+++ b/encode.go
@@ -29,6 +29,7 @@ type Encoder struct {
 	opts               []EncodeOption
 	indent             int
 	isFlowStyle        bool
+	anchorCallback     func(*ast.AnchorNode, interface{}) error
 	anchorPtrToNameMap map[uintptr]string
 
 	line        int
@@ -347,6 +348,26 @@ func (e *Encoder) encodeTime(v time.Time, column int) ast.Node {
 	return ast.String(token.New(value, value, e.pos(column)))
 }
 
+func (e *Encoder) encodeAnchor(anchorName string, value ast.Node, fieldValue reflect.Value, column int) (ast.Node, error) {
+	anchorNode := &ast.AnchorNode{
+		Start: token.New("&", "&", e.pos(column)),
+		Name:  ast.String(token.New(anchorName, anchorName, e.pos(column))),
+		Value: value,
+	}
+	if e.anchorCallback != nil {
+		if err := e.anchorCallback(anchorNode, fieldValue.Interface()); err != nil {
+			return nil, errors.Wrapf(err, "failed to marshal anchor")
+		}
+		if snode, ok := anchorNode.Name.(*ast.StringNode); ok {
+			anchorName = snode.Value
+		}
+	}
+	if fieldValue.Kind() == reflect.Ptr {
+		e.anchorPtrToNameMap[fieldValue.Pointer()] = anchorName
+	}
+	return anchorNode, nil
+}
+
 func (e *Encoder) encodeStruct(value reflect.Value, column int) (ast.Node, error) {
 	node := ast.Mapping(token.New("", "", e.pos(column)), e.isFlowStyle)
 	structType := value.Type()
@@ -382,25 +403,17 @@ func (e *Encoder) encodeStruct(value reflect.Value, column int) (ast.Node, error
 		key := e.encodeString(structField.RenderName, column)
 		switch {
 		case structField.AnchorName != "":
-			anchorName := structField.AnchorName
-			if fieldValue.Kind() == reflect.Ptr {
-				e.anchorPtrToNameMap[fieldValue.Pointer()] = anchorName
-			}
-			value = &ast.AnchorNode{
-				Start: token.New("&", "&", e.pos(column)),
-				Name:  ast.String(token.New(anchorName, anchorName, e.pos(column))),
-				Value: value,
+			anchorNode, err := e.encodeAnchor(structField.AnchorName, value, fieldValue, column)
+			if err != nil {
+				return nil, errors.Wrapf(err, "failed to encode anchor")
 			}
+			value = anchorNode
 		case structField.IsAutoAnchor:
-			anchorName := structField.RenderName
-			if fieldValue.Kind() == reflect.Ptr {
-				e.anchorPtrToNameMap[fieldValue.Pointer()] = anchorName
-			}
-			value = &ast.AnchorNode{
-				Start: token.New("&", "&", e.pos(column)),
-				Name:  ast.String(token.New(anchorName, anchorName, e.pos(column))),
-				Value: value,
+			anchorNode, err := e.encodeAnchor(structField.RenderName, value, fieldValue, column)
+			if err != nil {
+				return nil, errors.Wrapf(err, "failed to encode anchor")
 			}
+			value = anchorNode
 		case structField.IsAutoAlias:
 			if fieldValue.Kind() != reflect.Ptr {
 				return nil, xerrors.Errorf(
diff --git a/encode_test.go b/encode_test.go
index aab20edd..06fcc26d 100644
--- a/encode_test.go
+++ b/encode_test.go
@@ -9,6 +9,7 @@ import (
 	"time"
 
 	"github.com/goccy/go-yaml"
+	"github.com/goccy/go-yaml/ast"
 )
 
 func TestEncoder(t *testing.T) {
@@ -567,6 +568,85 @@ func TestEncoder_Flow(t *testing.T) {
 	}
 }
 
+func TestEncoder_MarshalAnchor(t *testing.T) {
+	type Host struct {
+		Hostname string
+		Username string
+		Password string
+	}
+	type HostDecl struct {
+		Host *Host `yaml:",anchor"`
+	}
+	type Queue struct {
+		Name  string `yaml:","`
+		*Host `yaml:",alias"`
+	}
+	var doc struct {
+		Hosts  []*HostDecl `yaml:"hosts"`
+		Queues []*Queue    `yaml:"queues"`
+	}
+	host1 := &Host{
+		Hostname: "host1.example.com",
+		Username: "userA",
+		Password: "pass1",
+	}
+	host2 := &Host{
+		Hostname: "host2.example.com",
+		Username: "userB",
+		Password: "pass2",
+	}
+	doc.Hosts = []*HostDecl{
+		{
+			Host: host1,
+		},
+		{
+			Host: host2,
+		},
+	}
+	doc.Queues = []*Queue{
+		{
+			Name: "queue",
+			Host: host1,
+		}, {
+			Name: "queue2",
+			Host: host2,
+		},
+	}
+	hostIdx := 1
+	opt := yaml.MarshalAnchor(func(anchor *ast.AnchorNode, value interface{}) error {
+		if _, ok := value.(*Host); ok {
+			nameNode := anchor.Name.(*ast.StringNode)
+			nameNode.Value = fmt.Sprintf("host%d", hostIdx)
+			hostIdx++
+		}
+		return nil
+	})
+
+	var buf bytes.Buffer
+	if err := yaml.NewEncoder(&buf, opt).Encode(doc); err != nil {
+		t.Fatalf("%+v", err)
+	}
+	expect := `
+hosts:
+- host: &host1
+    hostname: host1.example.com
+    username: userA
+    password: pass1
+- host: &host2
+    hostname: host2.example.com
+    username: userB
+    password: pass2
+queues:
+- name: queue
+  host: *host1
+- name: queue2
+  host: *host2
+`
+	if "\n"+buf.String() != expect {
+		t.Fatalf("unexpected output. %s", buf.String())
+	}
+}
+
 func Example_Marshal_ExplicitAnchorAlias() {
 	type T struct {
 		A int
diff --git a/option.go b/option.go
index 9fe7c3c9..2a98628a 100644
--- a/option.go
+++ b/option.go
@@ -1,6 +1,10 @@
 package yaml
 
-import "io"
+import (
+	"io"
+
+	"github.com/goccy/go-yaml/ast"
+)
 
 // DecodeOption functional option type for Decoder
 type DecodeOption func(d *Decoder) error
@@ -73,3 +77,11 @@ func Flow(isFlowStyle bool) EncodeOption {
 		return nil
 	}
 }
+
+// MarshalAnchor call back if encoder find an anchor during encoding
+func MarshalAnchor(callback func(*ast.AnchorNode, interface{}) error) EncodeOption {
+	return func(e *Encoder) error {
+		e.anchorCallback = callback
+		return nil
+	}
+}