support smart anchor option (#662)

Author: goccy

encode.go

@@ -32,8 +32,11 @@ type Encoder struct {
 	isFlowStyle                bool
 	isJSONStyle                bool
 	useJSONMarshaler           bool
+	enableSmartAnchor          bool
+	aliasRefToName             map[uintptr]string
+	anchorRefToName            map[uintptr]string
+	anchorNameMap              map[string]struct{}
 	anchorCallback             func(*ast.AnchorNode, interface{}) error
-	anchorPtrToNameMap         map[uintptr]string
 	customMarshalerMap         map[reflect.Type]func(interface{}) ([]byte, error)
 	useLiteralStyleIfMultiline bool
 	commentMap                 map[*Path][]*Comment
@@ -53,12 +56,14 @@ func NewEncoder(w io.Writer, opts ...EncodeOption) *Encoder {
 	return &Encoder{
 		writer:             w,
 		opts:               opts,
-		anchorPtrToNameMap: map[uintptr]string{},
 		customMarshalerMap: map[reflect.Type]func(interface{}) ([]byte, error){},
 		line:               1,
 		column:             1,
 		offset:             0,
 		indentNum:          DefaultIndentSpaces,
+		anchorRefToName:    make(map[uintptr]string),
+		anchorNameMap:      make(map[string]struct{}),
+		aliasRefToName:     make(map[uintptr]string),
 	}
 }
 
@@ -110,6 +115,13 @@ func (e *Encoder) EncodeToNodeContext(ctx context.Context, v interface{}) (ast.N
 			return nil, err
 		}
 	}
+	if e.enableSmartAnchor {
+		// during the first encoding, store all mappings between alias addresses and their names.
+		if _, err := e.encodeValue(ctx, reflect.ValueOf(v), 1); err != nil {
+			return nil, err
+		}
+		e.clearSmartAnchorRef()
+	}
 	node, err := e.encodeValue(ctx, reflect.ValueOf(v), 1)
 	if err != nil {
 		return nil, err
@@ -445,12 +457,8 @@ func (e *Encoder) encodeValue(ctx context.Context, v reflect.Value, column int)
 	case reflect.Float64:
 		return e.encodeFloat(v.Float(), 64), nil
 	case reflect.Ptr:
-		anchorName := e.anchorPtrToNameMap[v.Pointer()]
-		if anchorName != "" {
-			aliasName := anchorName
-			alias := ast.Alias(token.New("*", "*", e.pos(column)))
-			alias.Value = ast.String(token.New(aliasName, aliasName, e.pos(column)))
-			return alias, nil
+		if value := e.encodePtrAnchor(v, column); value != nil {
+			return value, nil
 		}
 		return e.encodeValue(ctx, v.Elem(), column)
 	case reflect.Interface:
@@ -463,6 +471,9 @@ func (e *Encoder) encodeValue(ctx context.Context, v reflect.Value, column int)
 		if mapSlice, ok := v.Interface().(MapSlice); ok {
 			return e.encodeMapSlice(ctx, mapSlice, column)
 		}
+		if value := e.encodePtrAnchor(v, column); value != nil {
+			return value, nil
+		}
 		return e.encodeSlice(ctx, v)
 	case reflect.Array:
 		return e.encodeArray(ctx, v)
@@ -477,12 +488,27 @@ func (e *Encoder) encodeValue(ctx context.Context, v reflect.Value, column int)
 		}
 		return e.encodeStruct(ctx, v, column)
 	case reflect.Map:
+		if value := e.encodePtrAnchor(v, column); value != nil {
+			return value, nil
+		}
 		return e.encodeMap(ctx, v, column), nil
 	default:
 		return nil, fmt.Errorf("unknown value type %s", v.Type().String())
 	}
 }
 
+func (e *Encoder) encodePtrAnchor(v reflect.Value, column int) ast.Node {
+	anchorName, exists := e.getAnchor(v.Pointer())
+	if !exists {
+		return nil
+	}
+	aliasName := anchorName
+	alias := ast.Alias(token.New("*", "*", e.pos(column)))
+	alias.Value = ast.String(token.New(aliasName, aliasName, e.pos(column)))
+	e.setSmartAlias(aliasName, v.Pointer())
+	return alias
+}
+
 func (e *Encoder) pos(column int) *token.Position {
 	return &token.Position{
 		Line:        e.line,
@@ -676,11 +702,23 @@ func (e *Encoder) encodeMap(ctx context.Context, value reflect.Value, column int
 		if e.isTagAndMapNode(value) {
 			value.AddColumn(e.indentNum)
 		}
+		keyText := fmt.Sprint(key)
+		vRef := e.toPointer(v)
+
+		// during the second encoding, an anchor is assigned if it is found to be used by an alias.
+		if aliasName, exists := e.getSmartAlias(vRef); exists {
+			anchorName := aliasName
+			anchorNode := ast.Anchor(token.New("&", "&", e.pos(column)))
+			anchorNode.Name = ast.String(token.New(anchorName, anchorName, e.pos(column)))
+			anchorNode.Value = value
+			value = anchorNode
+		}
 		node.Values = append(node.Values, ast.MappingValue(
 			nil,
-			e.encodeString(fmt.Sprint(key), column),
+			e.encodeString(keyText, column),
 			value,
 		))
+		e.setSmartAnchor(vRef, keyText)
 	}
 	return node
 }
@@ -761,7 +799,7 @@ func (e *Encoder) encodeAnchor(anchorName string, value ast.Node, fieldValue ref
 		}
 	}
 	if fieldValue.Kind() == reflect.Ptr {
-		e.anchorPtrToNameMap[fieldValue.Pointer()] = anchorName
+		e.setAnchor(fieldValue.Pointer(), anchorName)
 	}
 	return anchorNode, nil
 }
@@ -876,9 +914,87 @@ func (e *Encoder) encodeStruct(ctx context.Context, value reflect.Value, column
 			}
 		}
 		if inlineAnchorValue.Kind() == reflect.Ptr {
-			e.anchorPtrToNameMap[inlineAnchorValue.Pointer()] = anchorName
+			e.setAnchor(inlineAnchorValue.Pointer(), anchorName)
 		}
 		return anchorNode, nil
 	}
 	return node, nil
 }
+
+func (e *Encoder) toPointer(v reflect.Value) uintptr {
+	if e.isInvalidValue(v) {
+		return 0
+	}
+
+	switch v.Type().Kind() {
+	case reflect.Ptr:
+		return v.Pointer()
+	case reflect.Interface:
+		return e.toPointer(v.Elem())
+	case reflect.Slice:
+		return v.Pointer()
+	case reflect.Map:
+		return v.Pointer()
+	}
+	return 0
+}
+
+func (e *Encoder) clearSmartAnchorRef() {
+	if !e.enableSmartAnchor {
+		return
+	}
+	e.anchorRefToName = make(map[uintptr]string)
+	e.anchorNameMap = make(map[string]struct{})
+}
+
+func (e *Encoder) setSmartAnchor(ptr uintptr, name string) {
+	if !e.enableSmartAnchor {
+		return
+	}
+	e.setAnchor(ptr, e.generateAnchorName(name))
+}
+
+func (e *Encoder) setAnchor(ptr uintptr, name string) {
+	if ptr == 0 {
+		return
+	}
+	if name == "" {
+		return
+	}
+	e.anchorRefToName[ptr] = name
+	e.anchorNameMap[name] = struct{}{}
+}
+
+func (e *Encoder) generateAnchorName(base string) string {
+	if _, exists := e.anchorNameMap[base]; !exists {
+		return base
+	}
+	for i := 1; i < 100; i++ {
+		name := base + strconv.Itoa(i)
+		if _, exists := e.anchorNameMap[name]; exists {
+			continue
+		}
+		return name
+	}
+	return ""
+}
+
+func (e *Encoder) getAnchor(ref uintptr) (string, bool) {
+	anchorName, exists := e.anchorRefToName[ref]
+	return anchorName, exists
+}
+
+func (e *Encoder) setSmartAlias(name string, ref uintptr) {
+	if !e.enableSmartAnchor {
+		return
+	}
+	e.aliasRefToName[ref] = name
+}
+
+func (e *Encoder) getSmartAlias(ref uintptr) (string, bool) {
+	if !e.enableSmartAnchor {
+		return "", false
+	}
+	aliasName, exists := e.aliasRefToName[ref]
+	return aliasName, exists
+}

option.go

@@ -143,6 +143,18 @@ func Flow(isFlowStyle bool) EncodeOption {
 	}
 }
 
+// WithSmartAnchor when multiple map values share the same pointer,
+// an anchor is automatically assigned to the first occurrence, and aliases are used for subsequent elements.
+// The map key name is used as the anchor name by default.
+// If key names conflict, a suffix is automatically added to avoid collisions.
+// This is an experimental feature and cannot be used simultaneously with anchor tags.
+func WithSmartAnchor() EncodeOption {
+	return func(e *Encoder) error {
+		e.enableSmartAnchor = true
+		return nil
+	}
+}
+
 // UseLiteralStyleIfMultiline causes encoding multiline strings with a literal syntax,
 // no matter what characters they include
 func UseLiteralStyleIfMultiline(useLiteralStyleIfMultiline bool) EncodeOption {

yaml_test.go

@@ -98,3 +98,134 @@ d: &d [*c,*c]
 		}
 	}
 }
+
+func TestSmartAnchor(t *testing.T) {
+	var data = `
+a: &a [_,_,_,_,_,_,_,_,_,_,_,_,_,_,_]
+b: &b [*a,*a,*a,*a,*a,*a,*a,*a,*a,*a]
+c: &c [*b,*b,*b,*b,*b,*b,*b,*b,*b,*b]
+d: &d [*c,*c,*c,*c,*c,*c,*c,*c,*c,*c]
+e: &e [*d,*d,*d,*d,*d,*d,*d,*d,*d,*d]
+f: &f [*e,*e,*e,*e,*e,*e,*e,*e,*e,*e]
+g: &g [*f,*f,*f,*f,*f,*f,*f,*f,*f,*f]
+h: &h [*g,*g,*g,*g,*g,*g,*g,*g,*g,*g]
+i: &i [*h,*h,*h,*h,*h,*h,*h,*h,*h,*h]
+`
+	var v any
+	if err := yaml.Unmarshal([]byte(data), &v); err != nil {
+		t.Fatal(err)
+	}
+	got, err := yaml.MarshalWithOptions(v, yaml.WithSmartAnchor())
+	if err != nil {
+		t.Fatal(err)
+	}
+	expected := `
+a: &a
+- _
+- _
+- _
+- _
+- _
+- _
+- _
+- _
+- _
+- _
+- _
+- _
+- _
+- _
+- _
+b: &b
+- *a
+- *a
+- *a
+- *a
+- *a
+- *a
+- *a
+- *a
+- *a
+- *a
+c: &c
+- *b
+- *b
+- *b
+- *b
+- *b
+- *b
+- *b
+- *b
+- *b
+- *b
+d: &d
+- *c
+- *c
+- *c
+- *c
+- *c
+- *c
+- *c
+- *c
+- *c
+- *c
+e: &e
+- *d
+- *d
+- *d
+- *d
+- *d
+- *d
+- *d
+- *d
+- *d
+- *d
+f: &f
+- *e
+- *e
+- *e
+- *e
+- *e
+- *e
+- *e
+- *e
+- *e
+- *e
+g: &g
+- *f
+- *f
+- *f
+- *f
+- *f
+- *f
+- *f
+- *f
+- *f
+- *f
+h: &h
+- *g
+- *g
+- *g
+- *g
+- *g
+- *g
+- *g
+- *g
+- *g
+- *g
+i:
+- *h
+- *h
+- *h
+- *h
+- *h
+- *h
+- *h
+- *h
+- *h
+- *h
+`
+	if strings.TrimPrefix(expected, "\n") != string(got) {
+		t.Fatalf("failed to encode: %s", string(got))
+	}
+}