addrs: Generic types for maps and sets of addresses

The addrs.Set type previously snuck in accidentally as part of the work
to add addrs.UniqueKey and addrs.UniqueKeyer, because without support for
generic types the addrs.Set type was a bit of a safety hazard due to not
being able to enforce particular address types at compile time.

However, with Go 1.18 adding support for type parameters we can now turn
addrs.Set into a generic type over any specific addrs.UniqueKeyer type,
and complement it with an addrs.Map type which supports addrs.UniqueKeyer
keys as a way to encapsulate the handling of maps with UniqueKey keys that
we currently do inline in various other parts of Terraform.

This doesn't yet introduce any callers of these types, but we'll convert
existing users of addrs.UniqueKeyer gradually in subsequent commits.

internal/addrs/map.go

@@ -0,0 +1,128 @@
+package addrs
+
+// Map represents a mapping whose keys are address types that implement
+// UniqueKeyer.
+//
+// Since not all address types are comparable in the Go language sense, this
+// type cannot work with the typical Go map access syntax, and so instead has
+// a method-based syntax. Use this type only for situations where the key
+// type isn't guaranteed to always be a valid key for a standard Go map.
+type Map[K UniqueKeyer, V any] struct {
+	// Elems is the internal data structure of the map.
+	//
+	// This is exported to allow for comparisons during tests and other similar
+	// careful read operations, but callers MUST NOT modify this map directly.
+	// Use only the methods of Map to modify the contents of this structure,
+	// to ensure that it remains correct and consistent.
+	Elems map[UniqueKey]MapElem[K, V]
+}
+
+type MapElem[K UniqueKeyer, V any] struct {
+	Key   K
+	Value V
+}
+
+func MakeMap[K UniqueKeyer, V any](initialElems ...MapElem[K, V]) Map[K, V] {
+	inner := make(map[UniqueKey]MapElem[K, V], len(initialElems))
+	ret := Map[K, V]{inner}
+	for _, elem := range initialElems {
+		ret.Put(elem.Key, elem.Value)
+	}
+	return ret
+}
+
+func MakeMapElem[K UniqueKeyer, V any](key K, value V) MapElem[K, V] {
+	return MapElem[K, V]{key, value}
+}
+
+// Put inserts a new element into the map, or replaces an existing element
+// which has an equivalent key.
+func (m Map[K, V]) Put(key K, value V) {
+	realKey := key.UniqueKey()
+	m.Elems[realKey] = MapElem[K, V]{key, value}
+}
+
+// PutElement is like Put but takes the key and value from the given MapElement
+// structure instead of as individual arguments.
+func (m Map[K, V]) PutElement(elem MapElem[K, V]) {
+	m.Put(elem.Key, elem.Value)
+}
+
+// Remove deletes the element with the given key from the map, or does nothing
+// if there is no such element.
+func (m Map[K, V]) Remove(key K) {
+	realKey := key.UniqueKey()
+	delete(m.Elems, realKey)
+}
+
+// Get returns the value of the element with the given key, or the zero value
+// of V if there is no such element.
+func (m Map[K, V]) Get(key K) V {
+	realKey := key.UniqueKey()
+	return m.Elems[realKey].Value
+}
+
+// GetOk is like Get but additionally returns a flag for whether there was an
+// element with the given key present in the map.
+func (m Map[K, V]) GetOk(key K) (V, bool) {
+	realKey := key.UniqueKey()
+	elem, ok := m.Elems[realKey]
+	return elem.Value, ok
+}
+
+// Has returns true if and only if there is an element in the map which has the
+// given key.
+func (m Map[K, V]) Has(key K) bool {
+	realKey := key.UniqueKey()
+	_, ok := m.Elems[realKey]
+	return ok
+}
+
+// Len returns the number of elements in the map.
+func (m Map[K, V]) Len() int {
+	return len(m.Elems)
+}
+
+// Elements returns a slice containing a snapshot of the current elements of
+// the map, in an unpredictable order.
+func (m Map[K, V]) Elements() []MapElem[K, V] {
+	if len(m.Elems) == 0 {
+		return nil
+	}
+	ret := make([]MapElem[K, V], 0, len(m.Elems))
+	for _, elem := range m.Elems {
+		ret = append(ret, elem)
+	}
+	return ret
+}
+
+// Keys returns a Set[K] containing a snapshot of the current keys of elements
+// of the map.
+func (m Map[K, V]) Keys() Set[K] {
+	if len(m.Elems) == 0 {
+		return nil
+	}
+	ret := make(Set[K], len(m.Elems))
+
+	// We mess with the internals of Set here, rather than going through its
+	// public interface, because that means we can avoid re-calling UniqueKey
+	// on all of the elements when we know that our own Put method would have
+	// already done the same thing.
+	for realKey, elem := range m.Elems {
+		ret[realKey] = elem.Key
+	}
+	return ret
+}
+
+// Values returns a slice containing a snapshot of the current values of
+// elements of the map, in an unpredictable order.
+func (m Map[K, V]) Values() []V {
+	if len(m.Elems) == 0 {
+		return nil
+	}
+	ret := make([]V, 0, len(m.Elems))
+	for _, elem := range m.Elems {
+		ret = append(ret, elem.Value)
+	}
+	return ret
+}

internal/addrs/map_test.go

@@ -0,0 +1,83 @@
+package addrs
+
+import (
+	"testing"
+)
+
+func TestMap(t *testing.T) {
+	variableName := InputVariable{Name: "name"}
+	localHello := LocalValue{Name: "hello"}
+	pathModule := PathAttr{Name: "module"}
+	moduleBeep := ModuleCall{Name: "beep"}
+	eachKey := ForEachAttr{Name: "key"} // intentionally not in the map
+
+	m := MakeMap(
+		MakeMapElem[Referenceable](variableName, "Aisling"),
+	)
+
+	m.Put(localHello, "hello")
+	m.Put(pathModule, "boop")
+	m.Put(moduleBeep, "unrealistic")
+
+	keySet := m.Keys()
+	if want := variableName; !m.Has(want) {
+		t.Errorf("map does not include %s", want)
+	}
+	if want := variableName; !keySet.Has(want) {
+		t.Errorf("key set does not include %s", want)
+	}
+	if want := localHello; !m.Has(want) {
+		t.Errorf("map does not include %s", want)
+	}
+	if want := localHello; !keySet.Has(want) {
+		t.Errorf("key set does not include %s", want)
+	}
+	if want := pathModule; !keySet.Has(want) {
+		t.Errorf("key set does not include %s", want)
+	}
+	if want := moduleBeep; !keySet.Has(want) {
+		t.Errorf("key set does not include %s", want)
+	}
+	if doNotWant := eachKey; m.Has(doNotWant) {
+		t.Errorf("map includes rogue element %s", doNotWant)
+	}
+	if doNotWant := eachKey; keySet.Has(doNotWant) {
+		t.Errorf("key set includes rogue element %s", doNotWant)
+	}
+
+	if got, want := m.Get(variableName), "Aisling"; got != want {
+		t.Errorf("unexpected value %q for %s; want %q", got, variableName, want)
+	}
+	if got, want := m.Get(localHello), "hello"; got != want {
+		t.Errorf("unexpected value %q for %s; want %q", got, localHello, want)
+	}
+	if got, want := m.Get(pathModule), "boop"; got != want {
+		t.Errorf("unexpected value %q for %s; want %q", got, pathModule, want)
+	}
+	if got, want := m.Get(moduleBeep), "unrealistic"; got != want {
+		t.Errorf("unexpected value %q for %s; want %q", got, moduleBeep, want)
+	}
+	if got, want := m.Get(eachKey), ""; got != want {
+		// eachKey isn't in the map, so Get returns the zero value of string
+		t.Errorf("unexpected value %q for %s; want %q", got, eachKey, want)
+	}
+
+	if v, ok := m.GetOk(variableName); v != "Aisling" || !ok {
+		t.Errorf("GetOk for %q returned incorrect result (%q, %#v)", variableName, v, ok)
+	}
+	if v, ok := m.GetOk(eachKey); v != "" || ok {
+		t.Errorf("GetOk for %q returned incorrect result (%q, %#v)", eachKey, v, ok)
+	}
+
+	m.Remove(moduleBeep)
+	if doNotWant := moduleBeep; m.Has(doNotWant) {
+		t.Errorf("map still includes %s after removing it", doNotWant)
+	}
+	if want := moduleBeep; !keySet.Has(want) {
+		t.Errorf("key set no longer includes %s after removing it from the map; key set is supposed to be a snapshot at the time of call", want)
+	}
+	keySet = m.Keys()
+	if doNotWant := moduleBeep; keySet.Has(doNotWant) {
+		t.Errorf("key set still includes %s after a second call after removing it from the map", doNotWant)
+	}
+}

internal/addrs/set.go

@@ -1,23 +1,37 @@
 package addrs
 
 // Set represents a set of addresses of types that implement UniqueKeyer.
-type Set map[UniqueKey]UniqueKeyer
+//
+// Modify the set only by the methods on this type. This type exposes its
+// internals for convenience during reading, such as iterating over set elements
+// by ranging over the map values, but making direct modifications could
+// potentially make the set data invalid or inconsistent, leading to undefined
+// behavior elsewhere.
+type Set[T UniqueKeyer] map[UniqueKey]UniqueKeyer
 
-func (s Set) Has(addr UniqueKeyer) bool {
+// Has returns true if and only if the set includes the given address.
+func (s Set[T]) Has(addr UniqueKeyer) bool {
 	_, exists := s[addr.UniqueKey()]
 	return exists
 }
 
-func (s Set) Add(addr UniqueKeyer) {
+// Add inserts the given address into the set, if not already present. If
+// an equivalent address is already in the set, this replaces that address
+// with the new value.
+func (s Set[T]) Add(addr UniqueKeyer) {
 	s[addr.UniqueKey()] = addr
 }
 
-func (s Set) Remove(addr UniqueKeyer) {
+// Remove deletes the given address from the set, if present. If not present,
+// this is a no-op.
+func (s Set[T]) Remove(addr UniqueKeyer) {
 	delete(s, addr.UniqueKey())
 }
 
-func (s Set) Union(other Set) Set {
+// Union returns a new set which contains the union of all of the elements
-	ret := make(Set)
+// of both the reciever and the given other set.
+func (s Set[T]) Union(other Set[T]) Set[T] {
+	ret := make(Set[T])
 	for k, addr := range s {
 		ret[k] = addr
 	}
@@ -27,8 +41,10 @@ func (s Set) Union(other Set) Set {
 	return ret
 }
 
-func (s Set) Intersection(other Set) Set {
+// Intersection returns a new set which contains the intersection of all of the
-	ret := make(Set)
+// elements of both the reciever and the given other set.
+func (s Set[T]) Intersection(other Set[T]) Set[T] {
+	ret := make(Set[T])
 	for k, addr := range s {
 		if _, exists := other[k]; exists {
 			ret[k] = addr

internal/addrs/unique_key.go

@@ -21,3 +21,7 @@ type UniqueKey interface {
 type UniqueKeyer interface {
 	UniqueKey() UniqueKey
 }
+
+func Equivalent[T UniqueKeyer](a, b T) bool {
+	return a.UniqueKey() == b.UniqueKey()
+}