matrix-go-test/vendor/maunium.net/go/mautrix/util/configupgrade/helper.go

// Copyright (c) 2022 Tulir Asokan
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

package configupgrade

import (
	"fmt"
	"os"
	"strings"

	"gopkg.in/yaml.v3"
)

type YAMLMap map[string]YAMLNode
type YAMLList []YAMLNode

type YAMLNode struct {
	*yaml.Node
	Map  YAMLMap
	List YAMLList
	Key  *yaml.Node
}

type YAMLType uint32

const (
	Null YAMLType = 1 << iota
	Bool
	Str
	Int
	Float
	Timestamp
	List
	Map
	Binary
)

func (t YAMLType) String() string {
	switch t {
	case Null:
		return NullTag
	case Bool:
		return BoolTag
	case Str:
		return StrTag
	case Int:
		return IntTag
	case Float:
		return FloatTag
	case Timestamp:
		return TimestampTag
	case List:
		return SeqTag
	case Map:
		return MapTag
	case Binary:
		return BinaryTag
	default:
		panic(fmt.Errorf("can't convert type %d to string", t))
	}
}

func tagToType(tag string) YAMLType {
	switch tag {
	case NullTag:
		return Null
	case BoolTag:
		return Bool
	case StrTag:
		return Str
	case IntTag:
		return Int
	case FloatTag:
		return Float
	case TimestampTag:
		return Timestamp
	case SeqTag:
		return List
	case MapTag:
		return Map
	case BinaryTag:
		return Binary
	default:
		return 0
	}
}

const (
	NullTag      = "!!null"
	BoolTag      = "!!bool"
	StrTag       = "!!str"
	IntTag       = "!!int"
	FloatTag     = "!!float"
	TimestampTag = "!!timestamp"
	SeqTag       = "!!seq"
	MapTag       = "!!map"
	BinaryTag    = "!!binary"
)

func fromNode(node, key *yaml.Node) YAMLNode {
	switch node.Kind {
	case yaml.DocumentNode:
		return fromNode(node.Content[0], nil)
	case yaml.AliasNode:
		return fromNode(node.Alias, nil)
	case yaml.MappingNode:
		return YAMLNode{
			Node: node,
			Map:  parseYAMLMap(node),
			Key:  key,
		}
	case yaml.SequenceNode:
		return YAMLNode{
			Node: node,
			List: parseYAMLList(node),
		}
	default:
		return YAMLNode{Node: node, Key: key}
	}
}

func (yn *YAMLNode) toNode() *yaml.Node {
	yn.UpdateContent()
	return yn.Node
}

func (yn *YAMLNode) UpdateContent() {
	switch {
	case yn.Map != nil && yn.Node.Kind == yaml.MappingNode:
		yn.Content = yn.Map.toNodes()
	case yn.List != nil && yn.Node.Kind == yaml.SequenceNode:
		yn.Content = yn.List.toNodes()
	}
}

func parseYAMLList(node *yaml.Node) YAMLList {
	data := make(YAMLList, len(node.Content))
	for i, item := range node.Content {
		data[i] = fromNode(item, nil)
	}
	return data
}

func (yl YAMLList) toNodes() []*yaml.Node {
	nodes := make([]*yaml.Node, len(yl))
	for i, item := range yl {
		nodes[i] = item.toNode()
	}
	return nodes
}

func parseYAMLMap(node *yaml.Node) YAMLMap {
	if len(node.Content)%2 != 0 {
		panic(fmt.Errorf("uneven number of items in YAML map (%d)", len(node.Content)))
	}
	data := make(YAMLMap, len(node.Content)/2)
	for i := 0; i < len(node.Content); i += 2 {
		key := node.Content[i]
		value := node.Content[i+1]
		if key.Kind == yaml.ScalarNode {
			data[key.Value] = fromNode(value, key)
		}
	}
	return data
}

func (ym YAMLMap) toNodes() []*yaml.Node {
	nodes := make([]*yaml.Node, len(ym)*2)
	i := 0
	for key, value := range ym {
		nodes[i] = makeStringNode(key)
		nodes[i+1] = value.toNode()
		i += 2
	}
	return nodes
}

func makeStringNode(val string) *yaml.Node {
	var node yaml.Node
	node.SetString(val)
	return &node
}

func StringNode(val string) YAMLNode {
	return YAMLNode{Node: makeStringNode(val)}
}

type Helper struct {
	Base   YAMLNode
	Config YAMLNode
}

func NewHelper(base, cfg *yaml.Node) *Helper {
	return &Helper{
		Base:   fromNode(base, nil),
		Config: fromNode(cfg, nil),
	}
}

func (helper *Helper) AddSpaceBeforeComment(path ...string) {
	node := helper.GetBaseNode(path...)
	if node == nil || node.Key == nil {
		panic(fmt.Errorf("didn't find key at %+v", path))
	}
	node.Key.HeadComment = "\n" + node.Key.HeadComment
}

func (helper *Helper) Copy(allowedTypes YAMLType, path ...string) {
	base, cfg := helper.Base, helper.Config
	var ok bool
	for _, item := range path {
		cfg, ok = cfg.Map[item]
		if !ok {
			return
		}
		base, ok = base.Map[item]
		if !ok {
			_, _ = fmt.Fprintf(os.Stderr, "Ignoring config field %s which is missing in base config\n", strings.Join(path, "->"))
			return
		}
	}
	if allowedTypes&tagToType(cfg.Tag) == 0 {
		_, _ = fmt.Fprintf(os.Stderr, "Ignoring incorrect config field type %s at %s\n", cfg.Tag, strings.Join(path, "->"))
		return
	}
	base.Tag = cfg.Tag
	base.Style = cfg.Style
	switch base.Kind {
	case yaml.ScalarNode:
		base.Value = cfg.Value
	case yaml.SequenceNode, yaml.MappingNode:
		base.Content = cfg.Content
	}
}

func getNode(cfg YAMLNode, path []string) *YAMLNode {
	var ok bool
	for _, item := range path {
		cfg, ok = cfg.Map[item]
		if !ok {
			return nil
		}
	}
	return &cfg
}

func (helper *Helper) GetNode(path ...string) *YAMLNode {
	return getNode(helper.Config, path)
}

func (helper *Helper) GetBaseNode(path ...string) *YAMLNode {
	return getNode(helper.Base, path)
}

func (helper *Helper) Get(tag YAMLType, path ...string) (string, bool) {
	node := helper.GetNode(path...)
	if node == nil || node.Kind != yaml.ScalarNode || tag&tagToType(node.Tag) == 0 {
		return "", false
	}
	return node.Value, true
}

func (helper *Helper) GetBase(path ...string) string {
	return helper.GetBaseNode(path...).Value
}

func (helper *Helper) Set(tag YAMLType, value string, path ...string) {
	base := helper.Base
	for _, item := range path {
		base = base.Map[item]
	}
	base.Tag = tag.String()
	base.Value = value
}

func (helper *Helper) SetMap(value YAMLMap, path ...string) {
	base := helper.Base
	for _, item := range path {
		base = base.Map[item]
	}
	if base.Tag != MapTag || base.Kind != yaml.MappingNode {
		panic(fmt.Errorf("invalid target for SetMap(%+v): tag:%s, kind:%d", path, base.Tag, base.Kind))
	}
	base.Content = value.toNodes()
}