collections: Use Go 1.23 iterator patterns

In the first round of collections we exposed the internals of Set and Map
as a pragmatic concession to allow efficiently iterating over the elements.

Go 1.23 stabilizes "range-over-func" and package iter which together allow
us to expose an efficient iteration API _without_ exposing the internal
representations of these collections.

The benefit for Set is pretty marginal, but the benefit for Map is more
substantial: it avoids the oddity of exposing the MapElem values directly
and instead returns the key and value directly so that usage is very
similar to iterating over a built-in map.

internal/backend/local/backend_apply.go

@@ -483,7 +483,7 @@ func applyTimeInputValues(needVars collections.Set[string], decls map[string]*co
 	}
 
 	// We should now have a non-null value for each of the variables in needVars
-	for _, name := range needVars.Elems() {
+	for name := range needVars.All() {
 		val := cty.NullVal(cty.DynamicPseudoType)
 		if defn, ok := ret[name]; ok {
 			val = defn.Value

internal/collections/cmp.go

@@ -40,7 +40,7 @@ func (s Set[T]) transformForCmp() any {
 	ret := make(map[any]any, s.Len())
 	// It's okay to access the keys here because this package is allowed to
 	// depend on its own implementation details.
-	for k, v := range s.Elems() {
+	for k, v := range s.members {
 		ret[k] = v
 	}
 	return ret
@@ -50,7 +50,7 @@ func (m Map[K, V]) transformForCmp() any {
 	ret := make(map[any]any, m.Len())
 	// It's okay to access the keys here because this package is allowed to
 	// depend on its own implementation details.
-	for k, v := range m.Elems() {
+	for k, v := range m.elems {
 		ret[k] = v
 	}
 	return ret

internal/collections/map.go

@@ -3,6 +3,10 @@
 
 package collections
 
+import (
+	"iter"
+)
+
 // Set represents an associative array from keys of type K to values of type V.
 //
 // A caller-provided "key function" defines how to produce a comparable unique
@@ -116,31 +120,24 @@ func (m Map[K, V]) Delete(k K) {
 	delete(m.elems, uniq)
 }
 
-// Elems exposes the internal underlying representation of the map directly,
-// as a pragmatic compromise for efficient iteration.
-//
-// The result of this function is part of the internal state of the receiver
-// and so callers MUST NOT modify it. If a caller is using locks to ensure
-// safe concurrent access then any reads of the resulting map must be
-// guarded by the same lock as would be used for other methods that read
-// data from the reciever.
+// All returns an iterator over the elements of the map, in an unspecified
+// order.
 //
-// The only correct use of this function is as part of a "for ... range"
-// statement using only the values of the resulting map:
+// This is intended for use in a range-over-func statement, like this:
 //
-//	for _, elem := range map.Elems() {
-//	    k := elem.K
-//	    v := elem.V
-//	    // ...
+//	for k, v := range map.All() {
+//		// do something with k and/or v
 //	}
 //
-// Do not access or make any assumptions about the keys of the resulting
-// map. Their exact values are an implementation detail of the receiver.
-func (m Map[K, V]) Elems() map[UniqueKey[K]]MapElem[K, V] {
-	// This is regrettable but the only viable way to support efficient
-	// iteration over map elements until Go gains support for range
-	// loops over custom iterator functions.
-	return m.elems
+// Modifying the map during iteration causes unspecified results. Modifying
+// the map concurrently with advancing the iterator causes undefined behavior
+// including possible memory unsafety.
+func (m Map[K, V]) All() iter.Seq2[K, V] {
+	return func(yield func(K, V) bool) {
+		for _, elem := range m.elems {
+			yield(elem.K, elem.V)
+		}
+	}
 }
 
 // Len returns the number of elements in the map.

internal/collections/set.go

@@ -3,6 +3,10 @@
 
 package collections
 
+import (
+	"iter"
+)
+
 // Set represents an unordered set of values of a particular type.
 //
 // A caller-provided "key function" defines how to produce a comparable unique
@@ -81,29 +85,24 @@ func (s Set[T]) Remove(v T) {
 	delete(s.members, k)
 }
 
-// Elems exposes the internal underlying map representation of the set
-// directly, as a pragmatic compromise for efficient iteration.
-//
-// The result of this function is part of the internal state of the set
-// and so callers MUST NOT modify it. If a caller is using locks to ensure
-// safe concurrent access then any reads of the resulting map must be
-// guarded by the same lock as would be used for other methods that read
-// data from the set.
+// All returns an iterator over the elements of the set, in an unspecified
+// order.
 //
-// The only correct use of this function is as part of a "for ... range"
-// statement using only the values of the resulting map:
+// This is intended for use in a range-over-func statement, like this:
 //
-//	for _, elem := range set.Elems() {
-//	    // ...
+//	for elem := range set.All() {
+//		// do something with elem
 //	}
 //
-// Do not access or make any assumptions about the keys of the resulting
-// map. Their exact values are an implementation detail of the set.
-func (s Set[T]) Elems() map[UniqueKey[T]]T {
-	// This is regrettable but the only viable way to support efficient
-	// iteration over set members until Go gains support for range
-	// loops over custom iterator functions.
-	return s.members
+// Modifying the set during iteration causes unspecified results. Modifying
+// the set concurrently with advancing the iterator causes undefined behavior
+// including possible memory unsafety.
+func (s Set[T]) All() iter.Seq[T] {
+	return func(yield func(T) bool) {
+		for _, v := range s.members {
+			yield(v)
+		}
+	}
 }
 
 // Len returns the number of unique elements in the set.

internal/plans/planfile/tfplan.go

@@ -610,7 +610,7 @@ func writeTfplan(plan *plans.Plan, w io.Writer) error {
 	}
 	if plan.ApplyTimeVariables.Len() != 0 {
 		rawPlan.ApplyTimeVariables = make([]string, 0, plan.ApplyTimeVariables.Len())
-		for _, name := range plan.ApplyTimeVariables.Elems() {
+		for name := range plan.ApplyTimeVariables.All() {
 			rawPlan.ApplyTimeVariables = append(rawPlan.ApplyTimeVariables, name)
 		}
 	}

internal/stacks/stackplan/from_proto.go

@@ -226,27 +226,25 @@ func LoadFromProto(msgs []*anypb.Any) (*Plan, error) {
 
 	// Before we return we'll calculate the reverse dependency information
 	// based on the forward dependency information we loaded above.
-	for _, elem := range ret.Components.Elems() {
-		dependentInstAddr := elem.K
+	for dependentInstAddr, dependencyInst := range ret.Components.All() {
 		dependentAddr := stackaddrs.AbsComponent{
 			Stack: dependentInstAddr.Stack,
 			Item:  dependentInstAddr.Item.Component,
 		}
 
-		for _, dependencyAddr := range elem.V.Dependencies.Elems() {
+		for dependencyAddr := range dependencyInst.Dependencies.All() {
 			// FIXME: This is very inefficient because the current data structure doesn't
 			// allow looking up all of the component instances that have a particular
 			// component. This'll be okay as long as the number of components is
 			// small, but we'll need to improve this if we ever want to support stacks
 			// with a large number of components.
-			for _, elem := range ret.Components.Elems() {
-				maybeDependencyInstAddr := elem.K
+			for maybeDependencyInstAddr, dependencyInst := range ret.Components.All() {
 				maybeDependencyAddr := stackaddrs.AbsComponent{
 					Stack: maybeDependencyInstAddr.Stack,
 					Item:  maybeDependencyInstAddr.Item.Component,
 				}
 				if dependencyAddr.UniqueKey() == maybeDependencyAddr.UniqueKey() {
-					elem.V.Dependents.Add(dependentAddr)
+					dependencyInst.Dependents.Add(dependentAddr)
 				}
 			}
 		}

internal/stacks/stackplan/planned_change.go

@@ -210,7 +210,7 @@ func (pc *PlannedChangeComponentInstance) PlannedChangeProto() (*terraform1.Plan
 	}
 
 	componentAddrsRaw := make([]string, 0, pc.RequiredComponents.Len())
-	for _, componentAddr := range pc.RequiredComponents.Elems() {
+	for componentAddr := range pc.RequiredComponents.All() {
 		componentAddrsRaw = append(componentAddrsRaw, componentAddr.String())
 	}
 

internal/stacks/stackruntime/apply_test.go

@@ -1211,24 +1211,23 @@ func TestApplyWithStateManipulation(t *testing.T) {
 			}
 
 			wantCounts := tc.counts
-			for _, elem := range wantCounts.Elems() {
+			for k, v := range wantCounts.All() {
 				// First, make sure everything we wanted is present.
-				if !gotCounts.HasKey(elem.K) {
-					t.Errorf("wrong counts: wanted %s but didn't get it", elem.K)
+				if !gotCounts.HasKey(k) {
+					t.Errorf("wrong counts: wanted %s but didn't get it", k)
 				}
 
 				// And that the values actually match.
-				got, want := gotCounts.Get(elem.K), elem.V
+				got, want := gotCounts.Get(k), v
 				if diff := cmp.Diff(want, got); diff != "" {
 					t.Errorf("wrong counts for %s: %s", want.Addr, diff)
 				}
-
 			}
 
-			for _, elem := range gotCounts.Elems() {
+			for k, _ := range gotCounts.All() {
 				// Then, make sure we didn't get anything we didn't want.
-				if !wantCounts.HasKey(elem.K) {
-					t.Errorf("wrong counts: got %s but didn't want it", elem.K)
+				if !wantCounts.HasKey(k) {
+					t.Errorf("wrong counts: got %s but didn't want it", k)
 				}
 			}
 		})

internal/stacks/stackruntime/internal/stackeval/change_exec.go

@@ -176,8 +176,8 @@ func (r *ChangeExecResults) AwaitCompletion(ctx context.Context) {
 	// We don't have any single signal that everything is complete here,
 	// but it's sufficient for us to just visit each of our saved promise
 	// getters in turn and read from them.
-	for _, elem := range r.componentInstances.Elems() {
-		elem.V(ctx) // intentionally discards result; we only care that it's complete
+	for _, cb := range r.componentInstances.All() {
+		cb(ctx) // intentionally discards result; we only care that it's complete
 	}
 }
 

internal/stacks/stackruntime/internal/stackeval/component_instance.go

@@ -566,7 +566,7 @@ func (c *ComponentInstance) CheckModuleTreePlan(ctx context.Context) (*plans.Pla
 			// we need to force treating everything in this component as
 			// deferred so we can preserve the correct dependency ordering.
 			deferred := false
-			for _, depAddr := range c.call.RequiredComponents(ctx).Elems() {
+			for depAddr := range c.call.RequiredComponents(ctx).All() {
 				depStack := c.main.Stack(ctx, depAddr.Stack, PlanPhase)
 				if depStack == nil {
 					deferred = true // to be conservative

internal/stacks/stackruntime/internal/stackeval/main_apply.go

@@ -99,9 +99,7 @@ func ApplyPlan(ctx context.Context, config *stackconfig.Config, rawPlan []*anypb
 		// can error rather than deadlock if something goes wrong and causes
 		// us to try to depend on a result that isn't coming.
 		results, begin := ChangeExec(ctx, func(ctx context.Context, reg *ChangeExecRegistry[*Main]) {
-			for _, elem := range plan.Components.Elems() {
-				addr := elem.K
-				componentInstPlan := elem.V
+			for addr, componentInstPlan := range plan.Components.All() {
 				action := componentInstPlan.PlannedAction
 				dependencyAddrs := componentInstPlan.Dependencies
 				dependentAddrs := componentInstPlan.Dependents
@@ -183,7 +181,7 @@ func ApplyPlan(ctx context.Context, config *stackconfig.Config, rawPlan []*anypb
 						if depCount := waitForComponents.Len(); depCount != 0 {
 							log.Printf("[TRACE] stackeval: %s waiting for its predecessors (%d) to complete", addr, depCount)
 						}
-						for _, waitComponentAddr := range waitForComponents.Elems() {
+						for waitComponentAddr := range waitForComponents.All() {
 							if stack := main.Stack(ctx, waitComponentAddr.Stack, ApplyPhase); stack != nil {
 								if component := stack.Component(ctx, waitComponentAddr.Item); component != nil {
 									span.AddEvent("awaiting predecessor", trace.WithAttributes(
@@ -311,7 +309,7 @@ func checkApplyTimeVariables(plan *stackplan.Plan, providedValues map[stackaddrs
 
 	ret := make(map[stackaddrs.InputVariable]cty.Value, len(plan.RootInputValues)+plan.ApplyTimeInputVariables.Len())
 	maps.Copy(ret, plan.RootInputValues)
-	for _, varAddr := range plan.ApplyTimeInputVariables.Elems() {
+	for varAddr := range plan.ApplyTimeInputVariables.All() {
 		applyTimeVal := providedValues[varAddr]
 		if applyTimeVal.Value == cty.NilVal || applyTimeVal.Value.IsNull() {
 			diags = diags.Append(tfdiags.Sourceless(

internal/stacks/stackruntime/plan_test.go

@@ -3688,24 +3688,23 @@ func TestPlanWithStateManipulation(t *testing.T) {
 			}
 
 			wantCounts := tc.counts
-			for _, elem := range wantCounts.Elems() {
+			for k, v := range wantCounts.All() {
 				// First, make sure everything we wanted is present.
-				if !gotCounts.HasKey(elem.K) {
-					t.Errorf("wrong counts: wanted %s but didn't get it", elem.K)
+				if !gotCounts.HasKey(k) {
+					t.Errorf("wrong counts: wanted %s but didn't get it", k)
 				}
 
 				// And that the values actually match.
-				got, want := gotCounts.Get(elem.K), elem.V
+				got, want := gotCounts.Get(k), v
 				if diff := cmp.Diff(want, got); diff != "" {
 					t.Errorf("wrong counts for %s: %s", want.Addr, diff)
 				}
-
 			}
 
-			for _, elem := range gotCounts.Elems() {
+			for k, _ := range gotCounts.All() {
 				// Then, make sure we didn't get anything we didn't want.
-				if !wantCounts.HasKey(elem.K) {
-					t.Errorf("wrong counts: got %s but didn't want it", elem.K)
+				if !wantCounts.HasKey(k) {
+					t.Errorf("wrong counts: got %s but didn't want it", k)
 				}
 			}
 		})
@@ -3955,7 +3954,7 @@ var cmpCollectionsSet = cmp.Comparer(func(x, y collections.Set[stackaddrs.AbsCom
 		return false
 	}
 
-	for _, v := range x.Elems() {
+	for v := range x.All() {
 		if !y.Has(v) {
 			return false
 		}

internal/stacks/stackstate/applied_change.go

@@ -299,13 +299,13 @@ func (ac *AppliedChangeDiscardKeys) AppliedChangeProto() (*terraform1.AppliedCha
 		Raw:          make([]*terraform1.AppliedChange_RawChange, 0, ac.DiscardRawKeys.Len()),
 		Descriptions: make([]*terraform1.AppliedChange_ChangeDescription, 0, ac.DiscardDescKeys.Len()),
 	}
-	for _, key := range ac.DiscardRawKeys.Elems() {
+	for key := range ac.DiscardRawKeys.All() {
 		ret.Raw = append(ret.Raw, &terraform1.AppliedChange_RawChange{
 			Key:   statekeys.String(key),
 			Value: nil, // nil represents deletion
 		})
 	}
-	for _, key := range ac.DiscardDescKeys.Elems() {
+	for key := range ac.DiscardDescKeys.All() {
 		ret.Descriptions = append(ret.Descriptions, &terraform1.AppliedChange_ChangeDescription{
 			Key:         statekeys.String(key),
 			Description: &terraform1.AppliedChange_ChangeDescription_Deleted{

internal/stacks/stackstate/state.go

@@ -61,8 +61,8 @@ func (s *State) AllComponentInstances() collections.Set[stackaddrs.AbsComponentI
 		return ret
 	}
 	ret = collections.NewSet[stackaddrs.AbsComponentInstance]()
-	for _, elem := range s.componentInstances.Elems() {
-		ret.Add(elem.K)
+	for k, _ := range s.componentInstances.All() {
+		ret.Add(k)
 	}
 	return ret
 }
@@ -95,9 +95,8 @@ func (s *State) ComponentInstanceResourceInstanceObjects(addr stackaddrs.AbsComp
 // instance objects that are tracked in the state, across all components.
 func (s *State) AllResourceInstanceObjects() collections.Set[stackaddrs.AbsResourceInstanceObject] {
 	ret := collections.NewSet[stackaddrs.AbsResourceInstanceObject]()
-	for _, elem := range s.componentInstances.Elems() {
-		componentAddr := elem.K
-		for _, elem := range elem.V.resourceInstanceObjects.Elems {
+	for componentAddr, componentState := range s.componentInstances.All() {
+		for _, elem := range componentState.resourceInstanceObjects.Elems {
 			objKey := stackaddrs.AbsResourceInstanceObject{
 				Component: componentAddr,
 				Item:      elem.Key,
@@ -160,7 +159,7 @@ func (s *State) resourceInstanceObjectState(addr stackaddrs.AbsResourceInstanceO
 func (s *State) ComponentInstanceStateForModulesRuntime(addr stackaddrs.AbsComponentInstance) *states.State {
 	return states.BuildState(func(ss *states.SyncState) {
 		objAddrs := s.ComponentInstanceResourceInstanceObjects(addr)
-		for _, objAddr := range objAddrs.Elems() {
+		for objAddr := range objAddrs.All() {
 			rios := s.resourceInstanceObjectState(objAddr)
 
 			if objAddr.Item.IsCurrent() {

internal/terraform/context_apply.go

@@ -241,7 +241,7 @@ Note that the -target option is not suitable for routine use, and is provided on
 
 func checkApplyTimeVariables(needed collections.Set[string], gotValues InputValues, config *configs.Config) tfdiags.Diagnostics {
 	var diags tfdiags.Diagnostics
-	for _, name := range needed.Elems() {
+	for name := range needed.All() {
 		if vv, exists := gotValues[name]; !exists || vv.Value == cty.NilVal || vv.Value.IsNull() {
 			// This error message assumes that the only possible reason for
 			// an apply-time variable is because the variable is ephemeral,