telemetry/pdata/pcommon/map.go

// Copyright The OpenTelemetry Authors
// SPDX-License-Identifier: Apache-2.0

package pcommon // import "go.opentelemetry.io/collector/pdata/pcommon"

import (
	"iter"

	"go.uber.org/multierr"

	"go.opentelemetry.io/collector/pdata/internal"
)

// Map stores a map of string keys to elements of Value type.
//
// Must use NewMap function to create new instances.
// Important: zero-initialized instance is not valid for use.
type Map internal.MapWrapper

// NewMap creates a Map with 0 elements.
func NewMap() Map {
	orig := []internal.KeyValue(nil)
	return Map(internal.NewMapWrapper(&orig, internal.NewState()))
}

func (m Map) getOrig() *[]internal.KeyValue {
	return internal.GetMapOrig(internal.MapWrapper(m))
}

func (m Map) getState() *internal.State {
	return internal.GetMapState(internal.MapWrapper(m))
}

func newMap(orig *[]internal.KeyValue, state *internal.State) Map {
	return Map(internal.NewMapWrapper(orig, state))
}

// Clear erases any existing entries in this Map instance.
func (m Map) Clear() {
	m.getState().AssertMutable()
	*m.getOrig() = nil
}

// EnsureCapacity increases the capacity of this Map instance, if necessary,
// to ensure that it can hold at least the number of elements specified by the capacity argument.
func (m Map) EnsureCapacity(capacity int) {
	m.getState().AssertMutable()
	oldOrig := *m.getOrig()
	if capacity <= cap(oldOrig) {
		return
	}
	*m.getOrig() = make([]internal.KeyValue, len(oldOrig), capacity)
	copy(*m.getOrig(), oldOrig)
}

// Get returns the Value associated with the key and true. The returned
// Value is not a copy, it is a reference to the value stored in this map.
// It is allowed to modify the returned value using Value.Set* functions.
// Such modification will be applied to the value stored in this map.
// Accessing the returned value after modifying the underlying map
// (removing or adding new values) is an undefined behavior.
//
// If the key does not exist, returns a zero-initialized KeyValue and false.
// Calling any functions on the returned invalid instance may cause a panic.
func (m Map) Get(key string) (Value, bool) {
	for i := range *m.getOrig() {
		akv := &(*m.getOrig())[i]
		if akv.Key == key {
			return newValue(&akv.Value, m.getState()), true
		}
	}
	return newValue(nil, m.getState()), false
}

// Remove removes the entry associated with the key and returns true if the key
// was present in the map, otherwise returns false.
func (m Map) Remove(key string) bool {
	m.getState().AssertMutable()
	for i := range *m.getOrig() {
		akv := &(*m.getOrig())[i]
		if akv.Key == key {
			*akv = (*m.getOrig())[len(*m.getOrig())-1]
			*m.getOrig() = (*m.getOrig())[:len(*m.getOrig())-1]
			return true
		}
	}
	return false
}

// RemoveIf removes the entries for which the function in question returns true
func (m Map) RemoveIf(f func(string, Value) bool) {
	m.getState().AssertMutable()
	newLen := 0
	for i := 0; i < len(*m.getOrig()); i++ {
		if f((*m.getOrig())[i].Key, newValue(&(*m.getOrig())[i].Value, m.getState())) {
			(*m.getOrig())[i] = internal.KeyValue{}
			continue
		}
		if newLen == i {
			// Nothing to move, element is at the right place.
			newLen++
			continue
		}
		(*m.getOrig())[newLen] = (*m.getOrig())[i]
		(*m.getOrig())[i] = internal.KeyValue{}
		newLen++
	}
	*m.getOrig() = (*m.getOrig())[:newLen]
}

// PutEmpty inserts or updates an empty value to the map under given key
// and return the updated/inserted value.
func (m Map) PutEmpty(k string) Value {
	m.getState().AssertMutable()
	if av, existing := m.Get(k); existing {
		av.getOrig().Value = nil
		return newValue(av.getOrig(), m.getState())
	}
	*m.getOrig() = append(*m.getOrig(), internal.KeyValue{Key: k})
	return newValue(&(*m.getOrig())[len(*m.getOrig())-1].Value, m.getState())
}

// GetOrPutEmpty returns the Value associated with the key and true (loaded) if the key exists in the map,
// otherwise inserts an empty value to the map under the given key and returns the inserted value
// and false (loaded).
func (m Map) GetOrPutEmpty(k string) (Value, bool) {
	m.getState().AssertMutable()
	if av, existing := m.Get(k); existing {
		return av, true
	}
	*m.getOrig() = append(*m.getOrig(), internal.KeyValue{Key: k})
	return newValue(&(*m.getOrig())[len(*m.getOrig())-1].Value, m.getState()), false
}

// PutStr performs the Insert or Update action. The Value is
// inserted to the map that did not originally have the key. The key/value is
// updated to the map where the key already existed.
func (m Map) PutStr(k, v string) {
	m.getState().AssertMutable()
	if av, existing := m.Get(k); existing {
		av.SetStr(v)
		return
	}
	ov := internal.NewAnyValueStringValue()
	ov.StringValue = v
	*m.getOrig() = append(*m.getOrig(), internal.KeyValue{Key: k, Value: internal.AnyValue{Value: ov}})
}

// PutInt performs the Insert or Update action. The int Value is
// inserted to the map that did not originally have the key. The key/value is
// updated to the map where the key already existed.
func (m Map) PutInt(k string, v int64) {
	m.getState().AssertMutable()
	if av, existing := m.Get(k); existing {
		av.SetInt(v)
		return
	}
	ov := internal.NewAnyValueIntValue()
	ov.IntValue = v
	*m.getOrig() = append(*m.getOrig(), internal.KeyValue{Key: k, Value: internal.AnyValue{Value: ov}})
}

// PutDouble performs the Insert or Update action. The double Value is
// inserted to the map that did not originally have the key. The key/value is
// updated to the map where the key already existed.
func (m Map) PutDouble(k string, v float64) {
	m.getState().AssertMutable()
	if av, existing := m.Get(k); existing {
		av.SetDouble(v)
		return
	}
	ov := internal.NewAnyValueDoubleValue()
	ov.DoubleValue = v
	*m.getOrig() = append(*m.getOrig(), internal.KeyValue{Key: k, Value: internal.AnyValue{Value: ov}})
}

// PutBool performs the Insert or Update action. The bool Value is
// inserted to the map that did not originally have the key. The key/value is
// updated to the map where the key already existed.
func (m Map) PutBool(k string, v bool) {
	m.getState().AssertMutable()
	if av, existing := m.Get(k); existing {
		av.SetBool(v)
		return
	}
	ov := internal.NewAnyValueBoolValue()
	ov.BoolValue = v
	*m.getOrig() = append(*m.getOrig(), internal.KeyValue{Key: k, Value: internal.AnyValue{Value: ov}})
}

// PutEmptyBytes inserts or updates an empty byte slice under given key and returns it.
func (m Map) PutEmptyBytes(k string) ByteSlice {
	m.getState().AssertMutable()
	if av, existing := m.Get(k); existing {
		return av.SetEmptyBytes()
	}
	ov := internal.NewAnyValueBytesValue()
	*m.getOrig() = append(*m.getOrig(), internal.KeyValue{Key: k, Value: internal.AnyValue{Value: ov}})
	return ByteSlice(internal.NewByteSliceWrapper(&ov.BytesValue, m.getState()))
}

// PutEmptyMap inserts or updates an empty map under given key and returns it.
func (m Map) PutEmptyMap(k string) Map {
	m.getState().AssertMutable()
	if av, existing := m.Get(k); existing {
		return av.SetEmptyMap()
	}
	ov := internal.NewAnyValueKvlistValue()
	ov.KvlistValue = internal.NewKeyValueList()
	*m.getOrig() = append(*m.getOrig(), internal.KeyValue{Key: k, Value: internal.AnyValue{Value: ov}})
	return Map(internal.NewMapWrapper(&ov.KvlistValue.Values, m.getState()))
}

// PutEmptySlice inserts or updates an empty slice under given key and returns it.
func (m Map) PutEmptySlice(k string) Slice {
	m.getState().AssertMutable()
	if av, existing := m.Get(k); existing {
		return av.SetEmptySlice()
	}
	ov := internal.NewAnyValueArrayValue()
	ov.ArrayValue = internal.NewArrayValue()
	*m.getOrig() = append(*m.getOrig(), internal.KeyValue{Key: k, Value: internal.AnyValue{Value: ov}})
	return Slice(internal.NewSliceWrapper(&ov.ArrayValue.Values, m.getState()))
}

// Len returns the length of this map.
//
// Because the Map is represented internally by a slice of pointers, and the data are comping from the wire,
// it is possible that when iterating using "Range" to get access to fewer elements because nil elements are skipped.
func (m Map) Len() int {
	return len(*m.getOrig())
}

// Range calls f sequentially for each key and value present in the map. If f returns false, range stops the iteration.
//
// Example:
//
//	sm.Range(func(k string, v Value) bool {
//	    ...
//	})
func (m Map) Range(f func(k string, v Value) bool) {
	for i := range *m.getOrig() {
		kv := &(*m.getOrig())[i]
		if !f(kv.Key, Value(internal.NewValueWrapper(&kv.Value, m.getState()))) {
			break
		}
	}
}

// All returns an iterator over key-value pairs in the Map.
//
//	for k, v := range es.All() {
//	    ... // Do something with key-value pair
//	}
func (m Map) All() iter.Seq2[string, Value] {
	return func(yield func(string, Value) bool) {
		for i := range *m.getOrig() {
			kv := &(*m.getOrig())[i]
			if !yield(kv.Key, Value(internal.NewValueWrapper(&kv.Value, m.getState()))) {
				return
			}
		}
	}
}

// MoveTo moves all key/values from the current map overriding the destination and
// resetting the current instance to its zero value
func (m Map) MoveTo(dest Map) {
	m.getState().AssertMutable()
	dest.getState().AssertMutable()
	// If they point to the same data, they are the same, nothing to do.
	if m.getOrig() == dest.getOrig() {
		return
	}
	*dest.getOrig() = *m.getOrig()
	*m.getOrig() = nil
}

// CopyTo copies all elements from the current map overriding the destination.
func (m Map) CopyTo(dest Map) {
	dest.getState().AssertMutable()
	if m.getOrig() == dest.getOrig() {
		return
	}
	*dest.getOrig() = internal.CopyKeyValueSlice(*dest.getOrig(), *m.getOrig())
}

// AsRaw returns a standard go map representation of this Map.
func (m Map) AsRaw() map[string]any {
	rawMap := make(map[string]any, m.Len())
	m.Range(func(k string, v Value) bool {
		rawMap[k] = v.AsRaw()
		return true
	})
	return rawMap
}

// FromRaw overrides this Map instance from a standard go map.
func (m Map) FromRaw(rawMap map[string]any) error {
	m.getState().AssertMutable()
	if len(rawMap) == 0 {
		*m.getOrig() = nil
		return nil
	}

	var errs error
	origs := make([]internal.KeyValue, len(rawMap))
	ix := 0
	for k, iv := range rawMap {
		origs[ix].Key = k
		errs = multierr.Append(errs, newValue(&origs[ix].Value, m.getState()).FromRaw(iv))
		ix++
	}
	*m.getOrig() = origs
	return errs
}

// Equal checks equality with another Map
func (m Map) Equal(val Map) bool {
	if m.Len() != val.Len() {
		return false
	}

	fullEqual := true

	m.Range(func(k string, v Value) bool {
		vv, ok := val.Get(k)
		if !ok {
			fullEqual = false
			return fullEqual
		}

		if !v.Equal(vv) {
			fullEqual = false
		}
		return fullEqual
	})
	return fullEqual
}