interpres/orderedmap.go

// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
// SPDX-License-Identifier: MIT

package interpres

import (
	"fmt"
	"maps"
	"reflect"
	"slices"
	"sync"
)

// An OrderedMap is a string-keyed table that remembers the order its keys
// were set in, the shape a map[string]any cannot carry. Marshal writes a
// table of its own kind in that order, and decoding a document into one
// fills it in the order the document wrote the keys, where a map
// destination carries no order at all. The values are untyped, the shape
// the parser produces, so a nested table inside an OrderedMap is a plain
// map[string]any; the order is kept at the level the OrderedMap sits at.
//
// The zero value is an empty table ready for use.
type OrderedMap struct {
	keys   []string
	values map[string]any
}

var orderedMapType = reflect.TypeFor[OrderedMap]()

// NewOrderedMap returns an empty OrderedMap.
func NewOrderedMap() *OrderedMap { return &OrderedMap{} }

// Set stores value under key. A key the table already has keeps its position
// and takes the new value; a new one joins the end.
func (m *OrderedMap) Set(key string, value any) {
	if m.values == nil {
		m.values = make(map[string]any, 4)
	}
	if _, ok := m.values[key]; !ok {
		m.keys = append(m.keys, key)
	}
	m.values[key] = value
}

// Get returns the value under key, and whether the table has one.
func (m *OrderedMap) Get(key string) (any, bool) {
	v, ok := m.values[key]
	return v, ok
}

// Delete removes key. A later Set of the same key appends it to the end
// again.
func (m *OrderedMap) Delete(key string) {
	if _, ok := m.values[key]; !ok {
		return
	}
	delete(m.values, key)
	m.keys = slices.DeleteFunc(m.keys, func(k string) bool { return k == key })
}

// Keys returns the keys in the order they were set.
func (m *OrderedMap) Keys() []string { return m.keys }

// Len returns the number of keys.
func (m *OrderedMap) Len() int { return len(m.keys) }

// Range calls f for every key in order, stopping when f returns false.
func (m *OrderedMap) Range(f func(key string, value any) bool) {
	for _, k := range m.keys {
		if !f(k, m.values[k]) {
			return
		}
	}
}

// Map returns the values as a plain map, which carries no order. It is the
// view Marshal's Document-free callers need.
func (m *OrderedMap) Map() map[string]any { return m.values }

// --- decode: the order the document wrote ----------------------------------

// wantsOrderCache holds whether a destination type mentions OrderedMap
// anywhere a decode can reach. One computed answer per type, the same
// trade-off structSchemaCache makes.
var wantsOrderCache sync.Map // reflect.Type -> bool

// typeWantsOrder reports whether decoding into t can reach an OrderedMap, in
// which case the parse has to build the node tree the key order is read
// from. Structs walk their exported fields, and pointers, slices, arrays and
// maps walk their element; anything else holds no OrderedMap.
func typeWantsOrder(t reflect.Type) bool {
	if t == nil {
		return false
	}
	if v, ok := wantsOrderCache.Load(t); ok {
		return v.(bool)
	}
	r := scanWantsOrder(t, make(map[reflect.Type]bool))
	v, _ := wantsOrderCache.LoadOrStore(t, r)
	return v.(bool)
}

func scanWantsOrder(t reflect.Type, seen map[reflect.Type]bool) bool {
	for {
		if t == orderedMapType {
			return true
		}
		if seen[t] {
			return false
		}
		seen[t] = true
		switch t.Kind() {
		case reflect.Pointer, reflect.Slice, reflect.Array, reflect.Map:
			t = t.Elem()
		case reflect.Struct:
			for f := range t.Fields() {
				if f.PkgPath != "" {
					continue
				}
				if scanWantsOrder(f.Type, seen) {
					return true
				}
			}
			return false
		default:
			return false
		}
	}
}

// nodes maps a table's value map to its node, the index the decoder reads
// the written key order from. The key is the map header's runtime pointer,
// the one identity a map value offers; the nodes share their maps with the
// value tree, so one lookup per table is exact.
type nodeIndex map[uintptr]*Table

// indexNodeIndex walks a document's node tree into an index. A nil tree
// gives a nil index, which every lookup answers with nil.
func indexNodes(t *Table) nodeIndex {
	if t == nil {
		return nil
	}
	idx := nodeIndex{}
	var walk func(t *Table)
	walk = func(t *Table) {
		idx[reflect.ValueOf(t.values).Pointer()] = t
		for _, e := range t.entries {
			if e.child != nil {
				walk(e.child)
			}
			// The elements of a value array carry a node only where an element
			// is an inline table; the rest are nil.
			for _, el := range e.elements {
				if el != nil {
					walk(el)
				}
			}
		}
	}
	walk(t)
	return idx
}

// nodeOf returns the node a value table was parsed into, or nil when the
// parse built no node tree, which is the ordinary decode's shape. A tree
// built by hand carries no nodes either.
func (d *decoder) nodeOf(tbl map[string]any) *Table {
	return d.nodes[reflect.ValueOf(tbl).Pointer()]
}

// fillOrderedMap decodes a parsed table into an OrderedMap destination,
// taking the keys in the order the document wrote them. A table with no
// node, which is what a hand-built tree or a ParseMap result offers, fills
// in sorted key order, the deterministic order a map can offer.
func (d *decoder) fillOrderedMap(tbl map[string]any, dst reflect.Value) error {
	if !dst.CanAddr() {
		return fmt.Errorf("interpres: cannot decode into an OrderedMap that is not addressable")
	}
	om := dst.Addr().Interface().(*OrderedMap)
	if om.values == nil {
		om.values = make(map[string]any, len(tbl))
	}
	keys := slices.Sorted(maps.Keys(tbl))
	if node := d.nodeOf(tbl); node != nil {
		keys = node.Keys()
	}
	for _, key := range keys {
		val, ok := tbl[key]
		if !ok {
			continue
		}
		elem := reflect.New(reflect.TypeFor[any]()).Elem()
		if err := d.assign(val, elem); err != nil {
			return newDecodeError(key, err)
		}
		om.Set(key, elem.Interface())
	}
	return nil
}