Switch dyn.Mapping implementation to use map underneath

bundle/config/mutator/python/apply_python_output_test.go

@@ -256,7 +256,7 @@ func TestCreateOverrideVisitor_omitempty(t *testing.T) {
 func mapOf(key string, value dyn.Value) dyn.Value {
 	return dyn.NewValue(map[string]dyn.Value{
  key: value,
-	}, []dyn.Location{})
+	}, nil)
 }
 
 func mapOf2(key1 string, value1 dyn.Value, key2 string, value2 dyn.Value) dyn.Value {
@@ -267,5 +267,5 @@ func mapOf2(key1 string, value1 dyn.Value, key2 string, value2 dyn.Value) dyn.Va
 }
 
 func emptyMap() dyn.Value {
-	return dyn.NewValue(map[string]dyn.Value{}, []dyn.Location{})
+	return dyn.NewValue(map[string]dyn.Value{}, nil)
 }

libs/dyn/convert/from_typed.go

@@ -161,7 +161,7 @@ func fromTypedMap(src reflect.Value, ref dyn.Value) (dyn.Value, error) {
 	}
 
 	refm, _ := ref.AsMap()
-	out := dyn.NewMapping()
+	out := dyn.NewMappingWithSize(src.Len())
 	iter := src.MapRange()
 	for iter.Next() {
  k := iter.Key().String()

libs/dyn/convert/from_typed_test.go

@@ -24,7 +24,7 @@ func TestFromTypedStructZeroFields(t *testing.T) {
 	// For an empty struct with a non-nil reference we expect an empty map.
 	nv, err = FromTyped(src, dyn.V(map[string]dyn.Value{}))
 	require.NoError(t, err)
-	assert.Equal(t, dyn.V(map[string]dyn.Value{}), nv)
+	assert.Equal(t, dyn.V(dyn.Mapping{}), nv)
 }
 
 func TestFromTypedStructPointerZeroFields(t *testing.T) {
@@ -53,13 +53,13 @@ func TestFromTypedStructPointerZeroFields(t *testing.T) {
 	src = &Tmp{}
 	nv, err = FromTyped(src, dyn.NilValue)
 	require.NoError(t, err)
-	assert.Equal(t, dyn.V(map[string]dyn.Value{}), nv)
+	assert.Equal(t, dyn.V(dyn.Mapping{}), nv)
 
 	// For an initialized pointer with a non-nil reference we expect an empty map.
 	src = &Tmp{}
 	nv, err = FromTyped(src, dyn.V(map[string]dyn.Value{}))
 	require.NoError(t, err)
-	assert.Equal(t, dyn.V(map[string]dyn.Value{}), nv)
+	assert.Equal(t, dyn.V(dyn.Mapping{}), nv)
 }
 
 func TestFromTypedStructNilFields(t *testing.T) {

libs/dyn/convert/normalize_test.go

@@ -197,7 +197,7 @@ func TestNormalizeStructIncludeMissingFields(t *testing.T) {
  "ptr": dyn.V(map[string]dyn.Value{
  "string": dyn.V(""),
  }),
- "map": dyn.V(map[string]dyn.Value{}),
+ "map": dyn.V(dyn.Mapping{}),
  "slice": dyn.V([]dyn.Value{}),
  "string": dyn.V(""),
  "bool": dyn.V(false),

libs/dyn/dynassert/assert.go

@@ -1,34 +1,10 @@
 package dynassert
 
 import (
-	"github.com/databricks/cli/libs/dyn"
 	"github.com/stretchr/testify/assert"
 )
 
 func Equal(t assert.TestingT, expected, actual any, msgAndArgs ...any) bool {
-	ev, eok := expected.(dyn.Value)
-	av, aok := actual.(dyn.Value)
-	if eok && aok && ev.IsValid() && av.IsValid() {
- if !assert.Equal(t, ev.AsAny(), av.AsAny(), msgAndArgs...) {
- return false
- }
-
- // The values are equal on contents. Now compare the locations.
- if !assert.Equal(t, ev.Location(), av.Location(), msgAndArgs...) {
- return false
- }
-
- // Walk ev and av and compare the locations of each element.
- _, err := dyn.Walk(ev, func(p dyn.Path, evv dyn.Value) (dyn.Value, error) {
- avv, err := dyn.GetByPath(av, p)
- if assert.NoError(t, err, "unable to get value from actual value at path %v", p.String()) {
- assert.Equal(t, evv.Location(), avv.Location())
- }
- return evv, nil
- })
- return assert.NoError(t, err)
-	}
-
 	return assert.Equal(t, expected, actual, msgAndArgs...)
 }
 

libs/dyn/mapping.go

@@ -4,6 +4,8 @@ import (
 	"fmt"
 	"maps"
 	"slices"
+
+	"github.com/databricks/cli/libs/utils"
 )
 
 // Pair represents a single key-value pair in a Mapping.
@@ -17,44 +19,57 @@ type Pair struct {
 // We need to use dynamic values for keys because it lets us associate metadata
 // with keys (i.e. their definition location). Keys must be strings.
 type Mapping struct {
-pairs []Pair
-index map[string]int
+data map[string]Value
+keyLocations map[string][]Location
 }
 
 // NewMapping creates a new empty Mapping.
 func NewMapping() Mapping {
-	return Mapping{
- pairs: make([]Pair, 0),
- index: make(map[string]int),
-	}
+	return Mapping{}
+}
+
+// NewMappingWithSize creates a new Mapping preallocated to the specified size.
+func NewMappingWithSize(size int) Mapping {
+	return newMappingWithSize(size)
 }
 
 // newMappingWithSize creates a new Mapping preallocated to the specified size.
 func newMappingWithSize(size int) Mapping {
 	return Mapping{
- pairs: make([]Pair, 0, size),
- index: make(map[string]int, size),
+ data: make(map[string]Value, size),
+ // Do not initialize keyLocations, because in many contexts it's not going to be filled
+ // e.g. when constructing Mapping from map[string]Value.
 	}
 }
 
+// NewMappingFromGoMap creates a new Mapping from a Go map of string keys and dynamic values.
+func NewMappingFromGoMap(vin map[string]Value) Mapping {
+	return newMappingFromGoMap(vin)
+}
+
 // newMappingFromGoMap creates a new Mapping from a Go map of string keys and dynamic values.
 func newMappingFromGoMap(vin map[string]Value) Mapping {
-	m := newMappingWithSize(len(vin))
-	for k, v := range vin {
- m.Set(V(k), v) //nolint:errcheck
+	return Mapping{
+ data: maps.Clone(vin),
 	}
-	return m
 }
 
 // Pairs returns all the key-value pairs in the Mapping. The pairs are sorted by
 // their key in lexicographic order.
 func (m Mapping) Pairs() []Pair {
-	return m.pairs
+	pairs := make([]Pair, 0, len(m.data))
+	for _, k := range utils.SortedKeys(m.data) {
+ pairs = append(pairs, Pair{
+ Key: NewValue(k, m.keyLocations[k]),
+ Value: m.data[k],
+ })
+	}
+	return pairs
 }
 
 // Len returns the number of key-value pairs in the Mapping.
 func (m Mapping) Len() int {
-	return len(m.pairs)
+	return len(m.data)
 }
 
 // GetPair returns the key-value pair with the specified key.
@@ -70,10 +85,15 @@ func (m Mapping) GetPair(key Value) (Pair, bool) {
 // GetPairByString returns the key-value pair with the specified string key.
 // It also returns a boolean indicating whether the pair was found.
 func (m Mapping) GetPairByString(skey string) (Pair, bool) {
-	if i, ok := m.index[skey]; ok {
- return m.pairs[i], true
+	val, ok := m.data[skey]
+	if !ok {
+ return Pair{}, false
 	}
-	return Pair{}, false
+
+	return Pair{
+ Key: NewValue(skey, m.keyLocations[skey]),
+ Value: val,
+	}, true
 }
 
 // Get returns the value associated with the specified key.
@@ -100,50 +120,60 @@ func (m *Mapping) Set(key, value Value) error {
  return fmt.Errorf("key must be a string, got %s", key.Kind())
 	}
 
-	// If the key already exists, update the value.
-	if i, ok := m.index[skey]; ok {
- m.pairs[i].Value = value
- return nil
-	}
+	m.set(skey, key.l, value)
+	return nil
+}
 
-	// Otherwise, add a new pair.
-	m.pairs = append(m.pairs, Pair{key, value})
-	if m.index == nil {
- m.index = make(map[string]int)
+func (m *Mapping) set(key string, keyloc []Location, value Value) {
+	if m.data == nil {
+ m.data = make(map[string]Value)
+	}
+	m.data[key] = value
+
+	if len(keyloc) == 0 {
+ delete(m.keyLocations, key)
+	} else {
+ if m.keyLocations == nil {
+ m.keyLocations = make(map[string][]Location)
+ }
+ m.keyLocations[key] = slices.Clone(keyloc)
 	}
-	m.index[skey] = len(m.pairs) - 1
-	return nil
 }
 
 // Keys returns all the keys in the Mapping.
 func (m Mapping) Keys() []Value {
-	keys := make([]Value, 0, len(m.pairs))
-	for _, p := range m.pairs {
- keys = append(keys, p.Key)
+	keys := make([]Value, 0, len(m.data))
+	for _, k := range utils.SortedKeys(m.data) {
+ keys = append(keys, NewValue(k, m.keyLocations[k]))
 	}
 	return keys
 }
 
 // Values returns all the values in the Mapping.
 func (m Mapping) Values() []Value {
-	values := make([]Value, 0, len(m.pairs))
-	for _, p := range m.pairs {
- values = append(values, p.Value)
+	values := make([]Value, 0, len(m.data))
+	for _, k := range utils.SortedKeys(m.data) {
+ values = append(values, m.data[k])
 	}
 	return values
 }
 
 // Clone creates a shallow copy of the Mapping.
 func (m Mapping) Clone() Mapping {
 	return Mapping{
-pairs: slices.Clone(m.pairs),
-index: maps.Clone(m.index),
+data: maps.Clone(m.data),
+keyLocations: maps.Clone(m.keyLocations),
 	}
 }
 
 // Merge merges the key-value pairs from another Mapping into the current Mapping.
 func (m *Mapping) Merge(n Mapping) {
-	for _, p := range n.pairs {
- m.Set(p.Key, p.Value) //nolint:errcheck
+	if len(m.data) == 0 {
+ m.data = maps.Clone(n.data)
+ m.keyLocations = maps.Clone(n.keyLocations)
+ return
+	}
+	for key, value := range n.data {
+ m.set(key, n.keyLocations[key], value)
 	}
 }

libs/dyn/value.go

@@ -32,7 +32,7 @@ var NilValue = Value{
 
 // V constructs a new Value with the given value.
 func V(v any) Value {
-	return NewValue(v, []Location{})
+	return NewValue(v, nil)
 }
 
 // NewValue constructs a new Value with the given value and location.
@@ -103,10 +103,8 @@ func (v Value) AsAny() any {
 	case KindMap:
  m := v.v.(Mapping)
  out := make(map[string]any, m.Len())
- for _, pair := range m.pairs {
- pk := pair.Key
- pv := pair.Value
- out[pk.MustString()] = pv.AsAny()
+ for key, value := range m.data {
+ out[key] = value.AsAny()
  }
  return out
 	case KindSequence:
@@ -214,3 +212,26 @@ func (v Value) eq(w Value) bool {
  return v.v == w.v
 	}
 }
+
+func (v Value) DropKeyLocations() Value {
+	switch v.k {
+	case KindMap:
+ m := v.v.(Mapping)
+ out := make(map[string]Value, m.Len())
+ for _, pair := range m.Pairs() {
+ pk := pair.Key
+ pv := pair.Value.DropKeyLocations()
+ out[pk.MustString()] = pv
+ }
+ return NewValue(NewMappingFromGoMap(out), v.l)
+	case KindSequence:
+ vv := v.v.([]Value)
+ a := make([]Value, len(vv))
+ for i, v := range vv {
+ a[i] = v.DropKeyLocations()
+ }
+ return NewValue(a, v.l)
+	default:
+ return v
+	}
+}