miller/pkg/mlrval/mlrval_yaml.go

// YAML decode/encode for Mlrval and Mlrmap.
// Converts between YAML native types (from gopkg.in/yaml.v3) and Miller's
// record model. YAML maps become Mlrmap; keys are stringified (YAML allows
// non-string keys). Used by the YAML record reader and writer.

package mlrval

import (
	"fmt"
	"io"
	"sort"

	"gopkg.in/yaml.v3"
)

// MlrvalDecodeFromYAML decodes one YAML document from the decoder into an
// *Mlrval. Returns (nil, true, nil) on EOF. The decoded value can be a map
// (one record), array (array of records when elements are maps), or scalar.
func MlrvalDecodeFromYAML(decoder *yaml.Decoder) (*Mlrval, bool, error) {
	var doc interface{}
	err := decoder.Decode(&doc)
	if err == io.EOF {
		return nil, true, nil
	}
	if err != nil {
		return nil, false, err
	}
	if doc == nil {
		return NULL, false, nil
	}
	mv, err := mlrvalFromYAMLNative(doc)
	if err != nil {
		return nil, false, err
	}
	return mv, false, nil
}

// mlrvalFromYAMLNative converts a YAML-decoded value (map[interface{}]interface{},
// []interface{}, or scalar) into *Mlrval.
func mlrvalFromYAMLNative(v interface{}) (*Mlrval, error) {
	if v == nil {
		return NULL, nil
	}
	switch val := v.(type) {
	case map[interface{}]interface{}:
		return mlrvalFromYAMLMap(val)
	case map[string]interface{}:
		return mlrvalFromYAMLStringMap(val)
	case []interface{}:
		return mlrvalFromYAMLArray(val)
	case string:
		return FromString(val), nil
	case bool:
		return FromBool(val), nil
	case int:
		return FromInt(int64(val)), nil
	case int64:
		return FromInt(val), nil
	case uint64:
		if val <= 1<<63-1 {
			return FromInt(int64(val)), nil
		}
		return FromFloat(float64(val)), nil
	case float64:
		return FromFloat(val), nil
	case float32:
		return FromFloat(float64(val)), nil
	default:
		return FromString(fmt.Sprint(val)), nil
	}
}

func mlrvalFromYAMLMap(m map[interface{}]interface{}) (*Mlrval, error) {
	keys := make([]string, 0, len(m))
	for k := range m {
		keys = append(keys, yamlKeyString(k))
	}
	sort.Strings(keys)
	out := FromEmptyMap()
	for _, keyStr := range keys {
		var v interface{}
		for k, val := range m {
			if yamlKeyString(k) == keyStr {
				v = val
				break
			}
		}
		valMv, err := mlrvalFromYAMLNative(v)
		if err != nil {
			return nil, err
		}
		out.MapPut(FromString(keyStr), valMv)
	}
	return out, nil
}

func mlrvalFromYAMLStringMap(m map[string]interface{}) (*Mlrval, error) {
	keys := make([]string, 0, len(m))
	for k := range m {
		keys = append(keys, k)
	}
	sort.Strings(keys)
	out := FromEmptyMap()
	for _, k := range keys {
		valMv, err := mlrvalFromYAMLNative(m[k])
		if err != nil {
			return nil, err
		}
		out.MapPut(FromString(k), valMv)
	}
	return out, nil
}

func yamlKeyString(k interface{}) string {
	switch t := k.(type) {
	case string:
		return t
	case int:
		return fmt.Sprintf("%d", t)
	case int64:
		return fmt.Sprintf("%d", t)
	case float64:
		return fmt.Sprintf("%g", t)
	default:
		return fmt.Sprint(k)
	}
}

func mlrvalFromYAMLArray(a []interface{}) (*Mlrval, error) {
	out := FromEmptyArray()
	for _, elem := range a {
		mv, err := mlrvalFromYAMLNative(elem)
		if err != nil {
			return nil, err
		}
		out.ArrayAppend(mv)
	}
	return out, nil
}

// MlrmapToYAMLNative converts an Mlrmap to a *yaml.Node that preserves
// key insertion order. The returned node is suitable for yaml.Marshal.
func MlrmapToYAMLNative(mlrmap *Mlrmap) (*yaml.Node, error) {
	if mlrmap == nil {
		return &yaml.Node{Kind: yaml.MappingNode, Tag: "!!map"}, nil
	}
	node := &yaml.Node{
		Kind: yaml.MappingNode,
		Tag:  "!!map",
	}
	for pe := mlrmap.Head; pe != nil; pe = pe.Next {
		keyNode, err := encodeScalarNode(pe.Key)
		if err != nil {
			return nil, err
		}
		valNode, err := mlrvalToYAMLNode(pe.Value)
		if err != nil {
			return nil, err
		}
		node.Content = append(node.Content, keyNode, valNode)
	}
	return node, nil
}

// encodeScalarNode creates a yaml.Node by delegating to yaml.v3's Encode,
// which handles edge cases like NaN/Inf floats and non-UTF-8 strings.
func encodeScalarNode(v interface{}) (*yaml.Node, error) {
	node := &yaml.Node{}
	if err := node.Encode(v); err != nil {
		return nil, err
	}
	return node, nil
}

// mlrvalToYAMLNode converts *Mlrval to a *yaml.Node for yaml.Marshal.
func mlrvalToYAMLNode(mv *Mlrval) (*yaml.Node, error) {
	if mv == nil {
		return encodeScalarNode(nil)
	}
	switch mv.Type() {
	case MT_ABSENT, MT_NULL:
		return encodeScalarNode(nil)
	case MT_VOID:
		return encodeScalarNode("")
	case MT_STRING:
		s, _ := mv.GetStringValue()
		return encodeScalarNode(s)
	case MT_INT:
		i, _ := mv.GetIntValue()
		return encodeScalarNode(i)
	case MT_FLOAT:
		f, _ := mv.GetFloatValue()
		return encodeScalarNode(f)
	case MT_BOOL:
		b, _ := mv.GetBoolValue()
		return encodeScalarNode(b)
	case MT_ARRAY:
		arr := mv.GetArray()
		seqNode := &yaml.Node{Kind: yaml.SequenceNode, Tag: "!!seq"}
		for _, elem := range arr {
			v, err := mlrvalToYAMLNode(elem)
			if err != nil {
				return nil, err
			}
			seqNode.Content = append(seqNode.Content, v)
		}
		return seqNode, nil
	case MT_MAP:
		m := mv.GetMap()
		return MlrmapToYAMLNative(m)
	case MT_BYTES:
		return encodeScalarNode(mv.String())
	case MT_ERROR, MT_PENDING:
		return encodeScalarNode(mv.String())
	default:
		return encodeScalarNode(mv.String())
	}
}