stacks: refactor plan, state, and removed tracking with tree structures for efficient lookups (#36850)

internal/rpcapi/resource_identity.go

@@ -4,15 +4,16 @@
 package rpcapi
 
 import (
+	"github.com/zclconf/go-cty/cty"
+	ctyjson "github.com/zclconf/go-cty/cty/json"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+
 	"github.com/hashicorp/terraform/internal/addrs"
 	"github.com/hashicorp/terraform/internal/plans"
 	"github.com/hashicorp/terraform/internal/providers"
 	"github.com/hashicorp/terraform/internal/rpcapi/terraform1/stacks"
 	"github.com/hashicorp/terraform/internal/stacks/stackstate"
-	"github.com/zclconf/go-cty/cty"
-	ctyjson "github.com/zclconf/go-cty/cty/json"
-	"google.golang.org/grpc/codes"
-	"google.golang.org/grpc/status"
 )
 
 func listResourceIdentities(stackState *stackstate.State, identitySchemas map[addrs.Provider]map[string]providers.IdentitySchema) ([]*stacks.ListResourceIdentities_Resource, error) {
@@ -23,7 +24,7 @@ func listResourceIdentities(stackState *stackstate.State, identitySchemas map[ad
 		return resourceIdentities, nil
 	}
 
-	for ci := range stackState.AllComponentInstances().All() {
+	for ci := range stackState.AllComponentInstances() {
 		componentIdentities := stackState.IdentitiesForComponent(ci)
 		for ri, src := range componentIdentities {
 			// We skip resources without identity JSON

internal/stacks/stackaddrs/stack.go

@@ -69,6 +69,17 @@ func (s Stack) UniqueKey() collections.UniqueKey[Stack] {
 	return stackUniqueKey(s.String())
 }
 
+// ToStackCall converts the stack address into the absolute address of the stack
+// call that would create this stack.
+func (s Stack) ToStackCall() ConfigStackCall {
+	return ConfigStackCall{
+		Stack: s.Parent(),
+		Item: StackCall{
+			Name: s[len(s)-1].Name,
+		},
+	}
+}
+
 type stackUniqueKey string
 
 // IsUniqueKey implements collections.UniqueKey.

internal/stacks/stackconfig/declarations.go

@@ -53,7 +53,7 @@ type Declarations struct {
 
 	// RemovedEmbeddedStacks is the list of embedded stacks that have been removed
 	// from the configuration.
-	RemovedEmbeddedStacks collections.Map[stackaddrs.Stack, []*Removed]
+	RemovedEmbeddedStacks collections.Map[stackaddrs.ConfigStackCall, []*Removed]
 }
 
 func makeDeclarations() Declarations {
@@ -65,7 +65,7 @@ func makeDeclarations() Declarations {
 		OutputValues:          make(map[string]*OutputValue),
 		ProviderConfigs:       make(map[addrs.LocalProviderConfig]*ProviderConfig),
 		RemovedComponents:     collections.NewMap[stackaddrs.ConfigComponent, []*Removed](),
-		RemovedEmbeddedStacks: collections.NewMap[stackaddrs.Stack, []*Removed](),
+		RemovedEmbeddedStacks: collections.NewMap[stackaddrs.ConfigStackCall, []*Removed](),
 	}
 }
 
@@ -142,8 +142,11 @@ func (d *Declarations) addEmbeddedStack(decl *EmbeddedStack) tfdiags.Diagnostics
 		return diags
 	}
 
-	if blocks, exists := d.RemovedEmbeddedStacks.GetOk(stackaddrs.Stack{
-		stackaddrs.StackStep{Name: name},
+	if blocks, exists := d.RemovedEmbeddedStacks.GetOk(stackaddrs.ConfigStackCall{
+		Stack: nil,
+		Item: stackaddrs.StackCall{
+			Name: name,
+		},
 	}); exists {
 		for _, removed := range blocks {
 			if removed.From.Stack[0].Index == nil {
@@ -319,7 +322,7 @@ func (d *Declarations) addRemoved(decl *Removed) tfdiags.Diagnostics {
 
 		d.RemovedComponents.Put(addr, append(d.RemovedComponents.Get(addr), decl))
 	} else {
-		addr := decl.From.TargetStack()
+		addr := decl.From.TargetStack().ToStackCall()
 
 		if len(decl.From.Stack) == 1 && decl.From.Stack[0].Index == nil {
 			// Same logic as for components, we can just error a bit earlier

internal/stacks/stackplan/from_proto.go

@@ -35,11 +35,11 @@ type Loader struct {
 // Constructs a new [Loader], with an initial empty plan.
 func NewLoader() *Loader {
 	ret := &Plan{
+		Root:                    newStackInstance(stackaddrs.RootStackInstance),
 		RootInputValues:         make(map[stackaddrs.InputVariable]cty.Value),
 		ApplyTimeInputVariables: collections.NewSetCmp[stackaddrs.InputVariable](),
 		DeletedInputVariables:   collections.NewSet[stackaddrs.InputVariable](),
 		DeletedOutputValues:     collections.NewSet[stackaddrs.OutputValue](),
-		Components:              collections.NewMap[stackaddrs.AbsComponentInstance, *Component](),
 		DeletedComponents:       collections.NewSet[stackaddrs.AbsComponentInstance](),
 		PrevRunStateRaw:         make(map[string]*anypb.Any),
 	}
@@ -220,27 +220,24 @@ func (l *Loader) AddRaw(rawMsg *anypb.Any) error {
 			})
 		}
 
-		if !l.ret.Components.HasKey(addr) {
-			l.ret.Components.Put(addr, &Component{
-				PlannedAction:          plannedAction,
-				Mode:                   mode,
-				PlanApplyable:          msg.PlanApplyable,
-				PlanComplete:           msg.PlanComplete,
-				Dependencies:           dependencies,
-				Dependents:             collections.NewSet[stackaddrs.AbsComponent](),
-				PlannedInputValues:     inputVals,
-				PlannedInputValueMarks: inputValMarks,
-				PlannedOutputValues:    outputVals,
-				PlannedChecks:          checkResults,
-				PlannedFunctionResults: functionResults,
-
-				ResourceInstancePlanned:         addrs.MakeMap[addrs.AbsResourceInstanceObject, *plans.ResourceInstanceChangeSrc](),
-				ResourceInstancePriorState:      addrs.MakeMap[addrs.AbsResourceInstanceObject, *states.ResourceInstanceObjectSrc](),
-				ResourceInstanceProviderConfig:  addrs.MakeMap[addrs.AbsResourceInstanceObject, addrs.AbsProviderConfig](),
-				DeferredResourceInstanceChanges: addrs.MakeMap[addrs.AbsResourceInstanceObject, *plans.DeferredResourceInstanceChangeSrc](),
-			})
-		}
-		c := l.ret.Components.Get(addr)
+		c := l.ret.GetOrCreate(addr, &Component{
+			PlannedAction:          plannedAction,
+			Mode:                   mode,
+			PlanApplyable:          msg.PlanApplyable,
+			PlanComplete:           msg.PlanComplete,
+			Dependencies:           dependencies,
+			Dependents:             collections.NewSet[stackaddrs.AbsComponent](),
+			PlannedInputValues:     inputVals,
+			PlannedInputValueMarks: inputValMarks,
+			PlannedOutputValues:    outputVals,
+			PlannedChecks:          checkResults,
+			PlannedFunctionResults: functionResults,
+
+			ResourceInstancePlanned:         addrs.MakeMap[addrs.AbsResourceInstanceObject, *plans.ResourceInstanceChangeSrc](),
+			ResourceInstancePriorState:      addrs.MakeMap[addrs.AbsResourceInstanceObject, *states.ResourceInstanceObjectSrc](),
+			ResourceInstanceProviderConfig:  addrs.MakeMap[addrs.AbsResourceInstanceObject, addrs.AbsProviderConfig](),
+			DeferredResourceInstanceChanges: addrs.MakeMap[addrs.AbsResourceInstanceObject, *plans.DeferredResourceInstanceChangeSrc](),
+		})
 		err = c.PlanTimestamp.UnmarshalText([]byte(msg.PlanTimestamp))
 		if err != nil {
 			return fmt.Errorf("invalid plan timestamp %q for %s", msg.PlanTimestamp, addr)
@@ -329,26 +326,15 @@ func (l *Loader) Plan() (*Plan, error) {
 
 	// Before we return we'll calculate the reverse dependency information
 	// based on the forward dependency information we loaded above.
-	for dependentInstAddr, dependencyInst := range l.ret.Components.All() {
+	for dependentInstAddr, dependencyInst := range l.ret.AllComponents() {
 		dependentAddr := stackaddrs.AbsComponent{
 			Stack: dependentInstAddr.Stack,
 			Item:  dependentInstAddr.Item.Component,
 		}
 
 		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 maybeDependencyInstAddr, dependencyInst := range l.ret.Components.All() {
-				maybeDependencyAddr := stackaddrs.AbsComponent{
-					Stack: maybeDependencyInstAddr.Stack,
-					Item:  maybeDependencyInstAddr.Item.Component,
-				}
-				if dependencyAddr.UniqueKey() == maybeDependencyAddr.UniqueKey() {
-					dependencyInst.Dependents.Add(dependentAddr)
-				}
+			for _, dependencyInst := range l.ret.ComponentInstances(dependencyAddr) {
+				dependencyInst.Dependents.Add(dependentAddr)
 			}
 		}
 	}
@@ -439,7 +425,7 @@ func LoadComponentForResourceInstance(plan *Plan, change *tfstackdata1.PlanResou
 		DeposedKey:       deposedKey,
 	}
 
-	c, ok := plan.Components.GetOk(cAddr)
+	c, ok := plan.Root.GetOk(cAddr)
 	if !ok {
 		return nil, addrs.AbsResourceInstanceObject{}, addrs.AbsProviderConfig{}, fmt.Errorf("resource instance change for unannounced component instance %s", cAddr)
 	}
@@ -476,7 +462,7 @@ func LoadComponentForPartialResourceInstance(plan *Plan, change *tfstackdata1.Pl
 		DeposedKey:       deposedKey,
 	}
 
-	c, ok := plan.Components.GetOk(cAddr)
+	c, ok := plan.Root.GetOk(cAddr)
 	if !ok {
 		return nil, addrs.AbsResourceInstanceObject{}, addrs.AbsProviderConfig{}, fmt.Errorf("resource instance change for unannounced component instance %s", cAddr)
 	}

internal/stacks/stackplan/from_proto_test.go

@@ -26,6 +26,7 @@ func TestAddRaw(t *testing.T) {
 		"empty": {
 			Raw: nil,
 			Want: &Plan{
+				Root:            newStackInstance(stackaddrs.RootStackInstance),
 				PrevRunStateRaw: make(map[string]*anypb.Any),
 				RootInputValues: make(map[stackaddrs.InputVariable]cty.Value),
 			},
@@ -48,6 +49,7 @@ func TestAddRaw(t *testing.T) {
 				}),
 			},
 			Want: &Plan{
+				Root:            newStackInstance(stackaddrs.RootStackInstance),
 				PrevRunStateRaw: make(map[string]*anypb.Any),
 				RootInputValues: map[stackaddrs.InputVariable]cty.Value{
 					stackaddrs.InputVariable{Name: "foo"}: cty.StringVal("boop").Mark(marks.Sensitive),
@@ -67,6 +69,7 @@ func TestAddRaw(t *testing.T) {
 				}),
 			},
 			Want: &Plan{
+				Root:            newStackInstance(stackaddrs.RootStackInstance),
 				PrevRunStateRaw: make(map[string]*anypb.Any),
 				RootInputValues: map[stackaddrs.InputVariable]cty.Value{
 					stackaddrs.InputVariable{Name: "foo"}: cty.StringVal("boop"),

internal/stacks/stackplan/plan.go

@@ -4,6 +4,7 @@
 package stackplan
 
 import (
+	"iter"
 	"time"
 
 	"github.com/zclconf/go-cty/cty"
@@ -40,6 +41,11 @@ type Plan struct {
 	// Mode is the original mode of the plan.
 	Mode plans.Mode
 
+	// Root is the root StackInstance for the configuration being planned.
+	// The StackInstance object wraps the specific components for each stack
+	// instance.
+	Root *StackInstance
+
 	// The raw representation of the raw state that was provided in the request
 	// to create the plan. We use this primarily to perform mundane state
 	// data structure maintenence operations, such as discarding keys that
@@ -70,11 +76,6 @@ type Plan struct {
 	// deleted will be recomputed during the apply so are not needed.
 	DeletedOutputValues collections.Set[stackaddrs.OutputValue]
 
-	// Components contains the separate plans for each of the compoonent
-	// instances defined in the overall stack configuration, including any
-	// nested component instances from embedded stacks.
-	Components collections.Map[stackaddrs.AbsComponentInstance, *Component]
-
 	// DeletedComponents are a set of components that are in the state that
 	// should just be removed without any apply operation. This is typically
 	// because they are not referenced in the configuration and have no
@@ -90,42 +91,76 @@ type Plan struct {
 	PlanTimestamp time.Time
 }
 
+func (p *Plan) AllComponents() iter.Seq2[stackaddrs.AbsComponentInstance, *Component] {
+	return func(yield func(stackaddrs.AbsComponentInstance, *Component) bool) {
+		p.Root.iterate(yield)
+	}
+}
+
+func (p *Plan) ComponentInstanceAddresses(addr stackaddrs.AbsComponent) iter.Seq[stackaddrs.ComponentInstance] {
+	return func(yield func(stackaddrs.ComponentInstance) bool) {
+		stack := p.Root.GetDescendentStack(addr.Stack)
+		if stack != nil {
+			components := stack.Components[addr.Item]
+			for key := range components {
+				proceed := yield(stackaddrs.ComponentInstance{
+					Component: addr.Item,
+					Key:       key,
+				})
+				if !proceed {
+					return
+				}
+			}
+		}
+	}
+}
+
 // ComponentInstances returns a set of the component instances that belong to
 // the given component.
-func (p *Plan) ComponentInstances(addr stackaddrs.AbsComponent) collections.Set[stackaddrs.ComponentInstance] {
-	ret := collections.NewSet[stackaddrs.ComponentInstance]()
-	for elem := range p.Components.All() {
-		if elem.Stack.String() != addr.Stack.String() {
-			// Then
-			continue
-		}
-		if elem.Item.Component.Name != addr.Item.Name {
-			continue
+func (p *Plan) ComponentInstances(addr stackaddrs.AbsComponent) iter.Seq2[stackaddrs.ComponentInstance, *Component] {
+	return func(yield func(stackaddrs.ComponentInstance, *Component) bool) {
+		stack := p.Root.GetDescendentStack(addr.Stack)
+		if stack != nil {
+			components := stack.Components[addr.Item]
+			for key, component := range components {
+				proceed := yield(stackaddrs.ComponentInstance{
+					Component: addr.Item,
+					Key:       key,
+				}, component)
+				if !proceed {
+					return
+				}
+			}
 		}
-		ret.Add(elem.Item)
 	}
-	return ret
 }
 
-func (p *Plan) StackInstances(addr stackaddrs.AbsStackCall) collections.Set[stackaddrs.StackInstance] {
-	ret := collections.NewSet[stackaddrs.StackInstance]()
-	for key := range p.Components.All() {
-		if len(key.Stack) == 0 {
-			continue
+func (p *Plan) StackInstances(addr stackaddrs.AbsStackCall) iter.Seq[stackaddrs.StackInstance] {
+	return func(yield func(stackaddrs.StackInstance) bool) {
+		stack := p.Root.GetDescendentStack(addr.Stack)
+		if stack != nil {
+			stacks := stack.Children[addr.Item.Name]
+			for key := range stacks {
+				proceed := yield(append(addr.Stack, stackaddrs.StackInstanceStep{
+					Name: addr.Item.Name,
+					Key:  key,
+				}))
+				if !proceed {
+					return
+				}
+			}
 		}
+	}
+}
 
-		last := key.Stack[len(key.Stack)-1]
-		path := key.Stack[:len(key.Stack)-1]
+func (p *Plan) GetOrCreate(addr stackaddrs.AbsComponentInstance, component *Component) *Component {
+	targetStackInstance := p.Root.GetOrCreateDescendentStack(addr.Stack)
+	return targetStackInstance.GetOrCreateComponent(addr.Item, component)
+}
 
-		if path.String() != addr.Stack.String() {
-			continue
-		}
-		if last.Name != addr.Item.Name {
-			continue
-		}
-		ret.Add(key.Stack)
-	}
-	return ret
+func (p *Plan) GetComponent(addr stackaddrs.AbsComponentInstance) *Component {
+	targetStackInstance := p.Root.GetDescendentStack(addr.Stack)
+	return targetStackInstance.GetComponent(addr.Item)
 }
 
 // RequiredProviderInstances returns a description of all of the provider
@@ -136,9 +171,156 @@ func (p *Plan) StackInstances(addr stackaddrs.AbsStackCall) collections.Set[stac
 // function that operates on the configuration of a component instance rather
 // than the plan of one.
 func (p *Plan) RequiredProviderInstances(addr stackaddrs.AbsComponentInstance) addrs.Set[addrs.RootProviderConfig] {
-	component, ok := p.Components.GetOk(addr)
+	stack := p.Root.GetDescendentStack(addr.Stack)
+	if stack == nil {
+		return addrs.MakeSet[addrs.RootProviderConfig]()
+	}
+
+	components, ok := stack.Components[addr.Item.Component]
+	if !ok {
+		return addrs.MakeSet[addrs.RootProviderConfig]()
+	}
+
+	component, ok := components[addr.Item.Key]
 	if !ok {
 		return addrs.MakeSet[addrs.RootProviderConfig]()
 	}
 	return component.RequiredProviderInstances()
 }
+
+// StackInstance stores the components and embedded stacks for a single stack
+// instance.
+type StackInstance struct {
+	Address    stackaddrs.StackInstance
+	Children   map[string]map[addrs.InstanceKey]*StackInstance
+	Components map[stackaddrs.Component]map[addrs.InstanceKey]*Component
+}
+
+func newStackInstance(address stackaddrs.StackInstance) *StackInstance {
+	return &StackInstance{
+		Address:    address,
+		Components: make(map[stackaddrs.Component]map[addrs.InstanceKey]*Component),
+		Children:   make(map[string]map[addrs.InstanceKey]*StackInstance),
+	}
+}
+
+func (stack *StackInstance) GetComponent(addr stackaddrs.ComponentInstance) *Component {
+	components, ok := stack.Components[addr.Component]
+	if !ok {
+		return nil
+	}
+	return components[addr.Key]
+}
+
+func (stack *StackInstance) GetOrCreateComponent(addr stackaddrs.ComponentInstance, component *Component) *Component {
+	components, ok := stack.Components[addr.Component]
+	if !ok {
+		components = make(map[addrs.InstanceKey]*Component)
+	}
+	existing, ok := components[addr.Key]
+	if ok {
+		return existing
+	}
+	components[addr.Key] = component
+	stack.Components[addr.Component] = components
+	return component
+}
+
+func (stack *StackInstance) GetOrCreateDescendentStack(addr stackaddrs.StackInstance) *StackInstance {
+	if len(addr) == 0 {
+		return stack
+	}
+	next := stack.GetOrCreateChildStack(addr[0])
+	return next.GetOrCreateDescendentStack(addr[1:])
+}
+
+func (stack *StackInstance) GetOrCreateChildStack(step stackaddrs.StackInstanceStep) *StackInstance {
+	child := stack.GetChildStack(step)
+	if child == nil {
+		child = stack.CreateChildStack(step)
+	}
+	return child
+}
+
+func (stack *StackInstance) GetDescendentStack(addr stackaddrs.StackInstance) *StackInstance {
+	if len(addr) == 0 {
+		return stack
+	}
+
+	next := stack.GetChildStack(addr[0])
+	if next == nil {
+		return nil
+	}
+	return next.GetDescendentStack(addr[1:])
+}
+
+func (stack *StackInstance) GetChildStack(step stackaddrs.StackInstanceStep) *StackInstance {
+	insts, ok := stack.Children[step.Name]
+	if !ok {
+		return nil
+	}
+	return insts[step.Key]
+}
+
+func (stack *StackInstance) CreateChildStack(step stackaddrs.StackInstanceStep) *StackInstance {
+	stacks, ok := stack.Children[step.Name]
+	if !ok {
+		stacks = make(map[addrs.InstanceKey]*StackInstance)
+	}
+	stacks[step.Key] = newStackInstance(append(stack.Address, step))
+	stack.Children[step.Name] = stacks
+	return stacks[step.Key]
+}
+
+func (stack *StackInstance) GetOk(addr stackaddrs.AbsComponentInstance) (*Component, bool) {
+	if len(addr.Stack) == 0 {
+		component, ok := stack.Components[addr.Item.Component]
+		if !ok {
+			return nil, false
+		}
+
+		instance, ok := component[addr.Item.Key]
+		return instance, ok
+	}
+
+	stacks, ok := stack.Children[addr.Stack[0].Name]
+	if !ok {
+		return nil, false
+	}
+	next, ok := stacks[addr.Stack[0].Key]
+	if !ok {
+		return nil, false
+	}
+	return next.GetOk(stackaddrs.AbsComponentInstance{
+		Stack: addr.Stack[1:],
+		Item:  addr.Item,
+	})
+}
+
+func (stack *StackInstance) iterate(yield func(stackaddrs.AbsComponentInstance, *Component) bool) bool {
+	for name, components := range stack.Components {
+		for key, component := range components {
+			proceed := yield(stackaddrs.AbsComponentInstance{
+				Stack: stack.Address,
+				Item: stackaddrs.ComponentInstance{
+					Component: name,
+					Key:       key,
+				},
+			}, component)
+			if !proceed {
+				return false
+			}
+		}
+	}
+
+	for _, stacks := range stack.Children {
+		for _, inst := range stacks {
+			proceed := inst.iterate(yield)
+			if !proceed {
+				return false
+			}
+		}
+	}
+
+	return true
+}

internal/stacks/stackruntime/apply_test.go

@@ -4297,7 +4297,7 @@ func TestApply_WithProviderFunctions(t *testing.T) {
 	if len(plan.ProviderFunctionResults) == 0 {
 		t.Errorf("expected provider function results, got none")
 
-		if len(plan.Components.Get(mustAbsComponentInstance("component.self")).PlannedFunctionResults) == 0 {
+		if len(plan.GetComponent(mustAbsComponentInstance("component.self")).PlannedFunctionResults) == 0 {
 			t.Errorf("expected component function results, got none")
 		}
 	}

internal/stacks/stackruntime/helper_test.go

@@ -84,8 +84,6 @@ func (tc TestContext) Validate(t *testing.T, ctx context.Context, cycle TestCycl
 }
 
 func (tc TestContext) Plan(t *testing.T, ctx context.Context, state *stackstate.State, cycle TestCycle) *stackplan.Plan {
-	t.Helper()
-
 	request := PlanRequest{
 		PlanMode:  cycle.planMode,
 		Config:    tc.config,
@@ -163,8 +161,6 @@ func (tc TestContext) Plan(t *testing.T, ctx context.Context, state *stackstate.
 }
 
 func (tc TestContext) Apply(t *testing.T, ctx context.Context, plan *stackplan.Plan, cycle TestCycle) *stackstate.State {
-	t.Helper()
-
 	request := ApplyRequest{
 		Config: tc.config,
 		Plan:   plan,

internal/stacks/stackruntime/internal/stackeval/applying.go

@@ -114,7 +114,7 @@ func ApplyComponentPlan(ctx context.Context, main *Main, plan *plans.Plan, requi
 	// changed at all.
 	noOpResult := inst.PlaceholderApplyResultForSkippedApply(plan)
 
-	stackPlan := main.PlanBeingApplied().Components.Get(inst.Addr())
+	stackPlan := main.PlanBeingApplied().GetComponent(inst.Addr())
 
 	// We'll gather up our set of potentially-affected objects before we do
 	// anything else, because the modules runtime tends to mutate the objects

internal/stacks/stackruntime/internal/stackeval/applying_test.go

@@ -165,7 +165,7 @@ func TestApply_componentOrdering(t *testing.T) {
 		t.Fatalf("plan is not applyable")
 	}
 	{
-		cmpPlan := plan.Components.Get(cmpCInstAddr)
+		cmpPlan := plan.GetComponent(cmpCInstAddr)
 		gotDeps := cmpPlan.Dependencies
 		wantDeps := collections.NewSet[stackaddrs.AbsComponent]()
 		wantDeps.Add(cmpBAddr)
@@ -174,7 +174,7 @@ func TestApply_componentOrdering(t *testing.T) {
 		}
 	}
 	{
-		cmpPlan := plan.Components.Get(cmpBInst1Addr)
+		cmpPlan := plan.GetComponent(cmpBInst1Addr)
 		gotDeps := cmpPlan.Dependencies
 		wantDeps := collections.NewSet[stackaddrs.AbsComponent]()
 		wantDeps.Add(cmpAAddr)

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

@@ -333,7 +333,7 @@ func (c *ComponentInstance) ApplyModuleTreePlan(ctx context.Context, plan *plans
 	}
 
 	if plan.UIMode == plans.DestroyMode && plan.Changes.Empty() {
-		stackPlan := c.main.PlanBeingApplied().Components.Get(c.Addr())
+		stackPlan := c.main.PlanBeingApplied().GetComponent(c.Addr())
 
 		// If we're destroying and there's nothing to destroy, then we can
 		// consider this a no-op.
@@ -526,7 +526,7 @@ func (c *ComponentInstance) ResultValue(ctx context.Context, phase EvalPhase) ct
 		// begin their own destroy phases before we start ours.
 		if phase == ApplyPhase {
 			fullPlan := c.main.PlanBeingApplied()
-			ourPlan := fullPlan.Components.Get(c.Addr())
+			ourPlan := fullPlan.GetComponent(c.Addr())
 			if ourPlan == nil {
 				// Weird, but we'll tolerate it.
 				return cty.DynamicVal
@@ -727,7 +727,7 @@ func (c *ComponentInstance) CheckApply(ctx context.Context) ([]stackstate.Applie
 
 	var changes []stackstate.AppliedChange
 	if applyResult != nil {
-		changes, moreDiags = stackstate.FromState(ctx, applyResult.FinalState, c.main.PlanBeingApplied().Components.Get(c.Addr()), inputs, applyResult.AffectedResourceInstanceObjects, c)
+		changes, moreDiags = stackstate.FromState(ctx, applyResult.FinalState, c.main.PlanBeingApplied().GetComponent(c.Addr()), inputs, applyResult.AffectedResourceInstanceObjects, c)
 		diags = diags.Append(moreDiags)
 	}
 	return changes, diags

internal/stacks/stackruntime/internal/stackeval/main.go

@@ -18,7 +18,6 @@ import (
 	"github.com/hashicorp/terraform/internal/addrs"
 	fileProvisioner "github.com/hashicorp/terraform/internal/builtin/provisioners/file"
 	remoteExecProvisioner "github.com/hashicorp/terraform/internal/builtin/provisioners/remote-exec"
-	"github.com/hashicorp/terraform/internal/collections"
 	"github.com/hashicorp/terraform/internal/depsfile"
 	"github.com/hashicorp/terraform/internal/lang"
 	"github.com/hashicorp/terraform/internal/plans"
@@ -317,7 +316,7 @@ func (m *Main) MainStack() *Stack {
 			mode = m.PlanningOpts().PlanningMode
 		}
 
-		m.mainStack = newStack(m, stackaddrs.RootStackInstance, nil, config, collections.NewMap[stackaddrs.ConfigComponent, []*RemovedComponent](), collections.NewMap[stackaddrs.Stack, []*RemovedStackCall](), mode, false)
+		m.mainStack = newStack(m, stackaddrs.RootStackInstance, nil, config, newRemoved(), mode, false)
 	}
 	return m.mainStack
 }

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

@@ -77,7 +77,7 @@ func ApplyPlan(ctx context.Context, config *stackconfig.Config, plan *stackplan.
 		// 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 key, elem := range plan.Components.All() {
+			for key, elem := range plan.AllComponents() {
 				addr := key
 				componentInstPlan := elem
 				action := componentInstPlan.PlannedAction

internal/stacks/stackruntime/internal/stackeval/planning_test.go

@@ -250,8 +250,8 @@ func TestPlanning_DestroyMode(t *testing.T) {
 	plan, diags := testPlan(t, main)
 	assertNoDiagnostics(t, diags)
 
-	aCmpPlan := plan.Components.Get(aComponentInstAddr)
-	bCmpPlan := plan.Components.Get(bComponentInstAddr)
+	aCmpPlan := plan.GetComponent(aComponentInstAddr)
+	bCmpPlan := plan.GetComponent(bComponentInstAddr)
 	if aCmpPlan == nil || bCmpPlan == nil {
 		t.Fatalf(
 			"incomplete plan\n%s: %#v\n%s: %#v",
@@ -412,8 +412,6 @@ func TestPlanning_RequiredComponents(t *testing.T) {
 		plan, diags := testPlan(t, main)
 		assertNoDiagnostics(t, diags)
 
-		componentPlans := plan.Components
-
 		tests := []struct {
 			component        stackaddrs.AbsComponent
 			wantDependencies []stackaddrs.AbsComponent
@@ -453,7 +451,7 @@ func TestPlanning_RequiredComponents(t *testing.T) {
 						Component: test.component.Item,
 					},
 				}
-				cp := componentPlans.Get(instAddr)
+				cp := plan.GetComponent(instAddr)
 				{
 					got := cp.Dependencies
 					want := collections.NewSet[stackaddrs.AbsComponent]()
@@ -591,11 +589,11 @@ func TestPlanning_DeferredChangesPropagation(t *testing.T) {
 		plan, diags := testPlan(t, main)
 		assertNoErrors(t, diags)
 
-		firstPlan := plan.Components.Get(componentFirstInstAddr)
+		firstPlan := plan.GetComponent(componentFirstInstAddr)
 		if firstPlan.PlanComplete {
 			t.Error("first component has a complete plan; should be incomplete because it has deferred actions")
 		}
-		secondPlan := plan.Components.Get(componentSecondInstAddr)
+		secondPlan := plan.GetComponent(componentSecondInstAddr)
 		if secondPlan.PlanComplete {
 			t.Error("second component has a complete plan; should be incomplete because everything in it should've been deferred")
 		}
@@ -772,7 +770,7 @@ func TestPlanning_RemoveDataResource(t *testing.T) {
 		// address for this data resource, but no planned action because
 		// dropping a data resource from the state is not an "action" in the
 		// usual sense (it doesn't cause any calls to the provider).
-		mainPlan := plan.Components.Get(stackaddrs.AbsComponentInstance{
+		mainPlan := plan.GetComponent(stackaddrs.AbsComponentInstance{
 			Stack: stackaddrs.RootStackInstance,
 			Item: stackaddrs.ComponentInstance{
 				Component: stackaddrs.Component{Name: "main"},
@@ -836,7 +834,7 @@ func TestPlanning_PathValues(t *testing.T) {
 			t.Fatalf("unexpected diagnostics: %s", diags)
 		}
 
-		component, ok := plan.Components.GetOk(stackaddrs.AbsComponentInstance{
+		component := plan.GetComponent(stackaddrs.AbsComponentInstance{
 			Stack: stackaddrs.RootStackInstance,
 			Item: stackaddrs.ComponentInstance{
 				Component: stackaddrs.Component{
@@ -845,7 +843,7 @@ func TestPlanning_PathValues(t *testing.T) {
 				Key: addrs.NoKey,
 			},
 		})
-		if !ok {
+		if component == nil {
 			t.Fatalf("component not found in plan")
 		}
 

internal/stacks/stackruntime/internal/stackeval/removed.go

@@ -4,61 +4,70 @@
 package stackeval
 
 import (
-	"github.com/hashicorp/terraform/internal/collections"
+	"sync"
+
 	"github.com/hashicorp/terraform/internal/stacks/stackaddrs"
 )
 
 // Removed encapsulates the somewhat complicated logic for tracking and
 // managing the removed block instances in a given stack.
-//
-// All addresses within Removed are relative to the current stack.
 type Removed struct {
-	stackCallComponents collections.Map[stackaddrs.ConfigComponent, []*RemovedComponent]
-	localComponents     map[stackaddrs.Component][]*RemovedComponent
-	embeddedStackCalls  collections.Map[stackaddrs.Stack, []*RemovedStackCall]
-	localStackCalls     map[stackaddrs.StackCall][]*RemovedStackCall
+	sync.Mutex
+
+	components map[stackaddrs.Component][]*RemovedComponent
+	stackCalls map[stackaddrs.StackCall][]*RemovedStackCall
+
+	children map[string]*Removed
 }
 
-func newRemoved(localComponents map[stackaddrs.Component][]*RemovedComponent,
-	stackCallComponents collections.Map[stackaddrs.ConfigComponent, []*RemovedComponent],
-	localStackCalls map[stackaddrs.StackCall][]*RemovedStackCall,
-	embeddedStackCalls collections.Map[stackaddrs.Stack, []*RemovedStackCall]) *Removed {
+func newRemoved() *Removed {
 	return &Removed{
-		stackCallComponents: stackCallComponents,
-		localComponents:     localComponents,
-		localStackCalls:     localStackCalls,
-		embeddedStackCalls:  embeddedStackCalls,
+		components: make(map[stackaddrs.Component][]*RemovedComponent),
+		stackCalls: make(map[stackaddrs.StackCall][]*RemovedStackCall),
+		children:   make(map[string]*Removed),
 	}
 }
 
-// ForStackCall returns all removed component blocks that target the given
-// stack call. The addresses are transformed to be relative to the stack
-// created by the stack call.
-func (r *Removed) ForStackCall(addr stackaddrs.StackCall) (collections.Map[stackaddrs.ConfigComponent, []*RemovedComponent], collections.Map[stackaddrs.Stack, []*RemovedStackCall]) {
-	components := collections.NewMap[stackaddrs.ConfigComponent, []*RemovedComponent]()
-	for target, blocks := range r.stackCallComponents.All() {
-		step := target.Stack[0]
-		rest := target.Stack[1:]
+func (removed *Removed) Get(addr stackaddrs.ConfigStackCall) *Removed {
+	if len(addr.Stack) == 0 {
+		return removed.Next(addr.Item.Name)
+	}
+	return removed.Next(addr.Stack[0].Name).Get(stackaddrs.ConfigStackCall{
+		Stack: addr.Stack[1:],
+		Item:  addr.Item,
+	})
+}
 
-		if step.Name != addr.Name {
-			continue
-		}
+func (removed *Removed) Next(step string) *Removed {
+	removed.Lock()
+	defer removed.Unlock()
 
-		components.Put(stackaddrs.ConfigComponent{
-			Stack: rest,
-			Item:  target.Item,
-		}, blocks)
+	next := removed.children[step]
+	if next == nil {
+		next = newRemoved()
+		removed.children[step] = next
 	}
-	stackCalls := collections.NewMap[stackaddrs.Stack, []*RemovedStackCall]()
-	for target, blocks := range r.embeddedStackCalls.All() {
-		step := target[0]
-		rest := target[1:]
+	return next
+}
 
-		if step.Name != addr.Name {
-			continue
-		}
+func (removed *Removed) AddComponent(addr stackaddrs.ConfigComponent, components []*RemovedComponent) {
+	if len(addr.Stack) == 0 {
+		removed.components[addr.Item] = append(removed.components[addr.Item], components...)
+		return
+	}
+	removed.Next(addr.Stack[0].Name).AddComponent(stackaddrs.ConfigComponent{
+		Stack: addr.Stack[1:],
+		Item:  addr.Item,
+	}, components)
+}
 
-		stackCalls.Put(rest, blocks)
+func (removed *Removed) AddStackCall(addr stackaddrs.ConfigStackCall, stackCalls []*RemovedStackCall) {
+	if len(addr.Stack) == 0 {
+		removed.stackCalls[addr.Item] = append(removed.stackCalls[addr.Item], stackCalls...)
+		return
 	}
-	return components, stackCalls
+	removed.Next(addr.Stack[0].Name).AddStackCall(stackaddrs.ConfigStackCall{
+		Stack: addr.Stack[1:],
+		Item:  addr.Item,
+	}, stackCalls)
 }

internal/stacks/stackruntime/internal/stackeval/removed_component.go

@@ -163,7 +163,7 @@ func (r *RemovedComponent) Instances(ctx context.Context, phase EvalPhase) (map[
 					continue
 				}
 			case ApplyPhase:
-				if _, ok := r.main.PlanBeingApplied().Components.GetOk(ci.Addr()); ok {
+				if component := r.main.PlanBeingApplied().GetComponent(ci.Addr()); component != nil {
 					knownInstances[key] = ci
 					knownAddrs = append(knownAddrs, ci.Addr())
 					continue

internal/stacks/stackruntime/internal/stackeval/removed_component_instance.go

@@ -173,7 +173,7 @@ func (r *RemovedComponentInstance) PlanPrevInputs() terraform.InputValues {
 func (r *RemovedComponentInstance) PlanCurrentInputs() (cty.Value, tfdiags.Diagnostics) {
 	var diags tfdiags.Diagnostics
 
-	plan := r.main.PlanBeingApplied().Components.Get(r.Addr())
+	plan := r.main.PlanBeingApplied().GetComponent(r.Addr())
 	inputs := make(map[string]cty.Value, len(plan.PlannedInputValues))
 	for name, input := range plan.PlannedInputValues {
 		value, err := input.Decode(cty.DynamicPseudoType)
@@ -271,7 +271,7 @@ func (r *RemovedComponentInstance) CheckApply(ctx context.Context) ([]stackstate
 
 	var changes []stackstate.AppliedChange
 	if result != nil {
-		changes, moreDiags = stackstate.FromState(ctx, result.FinalState, r.main.PlanBeingApplied().Components.Get(r.Addr()), inputs, result.AffectedResourceInstanceObjects, r)
+		changes, moreDiags = stackstate.FromState(ctx, result.FinalState, r.main.PlanBeingApplied().GetComponent(r.Addr()), inputs, result.AffectedResourceInstanceObjects, r)
 		diags = diags.Append(moreDiags)
 	}
 	return changes, diags

internal/stacks/stackruntime/internal/stackeval/removed_stack_call.go

@@ -12,7 +12,6 @@ import (
 	"github.com/zclconf/go-cty/cty"
 
 	"github.com/hashicorp/terraform/internal/addrs"
-	"github.com/hashicorp/terraform/internal/collections"
 	"github.com/hashicorp/terraform/internal/instances"
 	"github.com/hashicorp/terraform/internal/lang"
 	"github.com/hashicorp/terraform/internal/promising"
@@ -27,7 +26,7 @@ var _ Applyable = (*RemovedStackCall)(nil)
 
 type RemovedStackCall struct {
 	stack  *Stack
-	target stackaddrs.Stack // relative to stack
+	target stackaddrs.ConfigStackCall // relative to stack
 
 	config *RemovedStackCallConfig
 
@@ -37,7 +36,7 @@ type RemovedStackCall struct {
 	instances    perEvalPhase[promising.Once[withDiagnostics[instancesResult[*RemovedStackCallInstance]]]]
 }
 
-func newRemovedStackCall(main *Main, target stackaddrs.Stack, stack *Stack, config *RemovedStackCallConfig) *RemovedStackCall {
+func newRemovedStackCall(main *Main, target stackaddrs.ConfigStackCall, stack *Stack, config *RemovedStackCallConfig) *RemovedStackCall {
 	return &RemovedStackCall{
 		stack:  stack,
 		target: target,
@@ -48,10 +47,8 @@ func newRemovedStackCall(main *Main, target stackaddrs.Stack, stack *Stack, conf
 
 // GetExternalRemovedBlocks fetches the removed blocks that target the stack
 // instances being created by this stack call.
-func (r *RemovedStackCall) GetExternalRemovedBlocks() (collections.Map[stackaddrs.ConfigComponent, []*RemovedComponent], collections.Map[stackaddrs.Stack, []*RemovedStackCall]) {
-	return r.stack.Removed().ForStackCall(stackaddrs.StackCall{
-		Name: r.target[0].Name,
-	})
+func (r *RemovedStackCall) GetExternalRemovedBlocks() *Removed {
+	return r.stack.Removed().Get(r.target)
 }
 
 func (r *RemovedStackCall) ForEachValue(ctx context.Context, phase EvalPhase) (cty.Value, tfdiags.Diagnostics) {

internal/stacks/stackruntime/internal/stackeval/removed_stack_call_config.go

@@ -27,7 +27,7 @@ var (
 )
 
 type RemovedStackCallConfig struct {
-	target stackaddrs.Stack // relative to stack
+	target stackaddrs.ConfigStackCall // relative to stack
 	config *stackconfig.Removed
 	stack  *StackConfig
 
@@ -37,7 +37,7 @@ type RemovedStackCallConfig struct {
 	inputVariableValues perEvalPhase[promising.Once[withDiagnostics[map[stackaddrs.InputVariable]cty.Value]]]
 }
 
-func newRemovedStackCallConfig(main *Main, target stackaddrs.Stack, stack *StackConfig, config *stackconfig.Removed) *RemovedStackCallConfig {
+func newRemovedStackCallConfig(main *Main, target stackaddrs.ConfigStackCall, stack *StackConfig, config *stackconfig.Removed) *RemovedStackCallConfig {
 	return &RemovedStackCallConfig{
 		target: target,
 		config: config,
@@ -48,10 +48,10 @@ func newRemovedStackCallConfig(main *Main, target stackaddrs.Stack, stack *Stack
 
 func (r *RemovedStackCallConfig) TargetConfig() *StackConfig {
 	current := r.stack
-	for _, step := range r.target {
+	for _, step := range r.target.Stack {
 		current = current.ChildConfig(step)
 	}
-	return current
+	return current.ChildConfig(stackaddrs.StackStep{Name: r.target.Item.Name})
 }
 
 func (r *RemovedStackCallConfig) InputVariableValues(ctx context.Context, phase EvalPhase) (map[stackaddrs.InputVariable]cty.Value, tfdiags.Diagnostics) {

internal/stacks/stackruntime/internal/stackeval/removed_stack_call_instance.go

@@ -48,8 +48,7 @@ func newRemovedStackCallInstance(call *RemovedStackCall, from stackaddrs.StackIn
 
 func (r *RemovedStackCallInstance) Stack(ctx context.Context, phase EvalPhase) *Stack {
 	stack, err := r.stack.For(phase).Do(ctx, r.from.String()+" create", func(ctx context.Context) (*Stack, error) {
-		components, embeddedStackCalls := r.call.GetExternalRemovedBlocks()
-		return newStack(r.main, r.from, r.call.stack, r.call.config.TargetConfig(), components, embeddedStackCalls, plans.DestroyMode, r.deferred), nil
+		return newStack(r.main, r.from, r.call.stack, r.call.config.TargetConfig(), r.call.GetExternalRemovedBlocks(), plans.DestroyMode, r.deferred), nil
 	})
 	if err != nil {
 		// we never return an error from within the once call, so this shouldn't

internal/stacks/stackruntime/internal/stackeval/stack.go

@@ -6,6 +6,7 @@ package stackeval
 import (
 	"context"
 	"fmt"
+	"iter"
 	"sync"
 	"time"
 
@@ -36,21 +37,18 @@ type Stack struct {
 	deferred bool
 	mode     plans.Mode
 
-	// externalRemovedComponents are the set of removed blocks that might
-	// target components within either this stack or a child of this stack.
-	externalRemovedComponents collections.Map[stackaddrs.ConfigComponent, []*RemovedComponent]
-	externalRemovedStackCalls collections.Map[stackaddrs.Stack, []*RemovedStackCall]
+	removed *Removed // contains removed logic
 
 	// The remaining fields memoize other objects we might create in response
 	// to method calls. Must lock "mu" before interacting with them.
-	mu             sync.Mutex
-	inputVariables map[stackaddrs.InputVariable]*InputVariable
-	localValues    map[stackaddrs.LocalValue]*LocalValue
-	stackCalls     map[stackaddrs.StackCall]*StackCall
-	outputValues   map[stackaddrs.OutputValue]*OutputValue
-	components     map[stackaddrs.Component]*Component
-	removed        *Removed
-	providers      map[stackaddrs.ProviderConfigRef]*Provider
+	mu                 sync.Mutex
+	inputVariables     map[stackaddrs.InputVariable]*InputVariable
+	localValues        map[stackaddrs.LocalValue]*LocalValue
+	stackCalls         map[stackaddrs.StackCall]*StackCall
+	outputValues       map[stackaddrs.OutputValue]*OutputValue
+	components         map[stackaddrs.Component]*Component
+	providers          map[stackaddrs.ProviderConfigRef]*Provider
+	removedInitialised bool
 }
 
 var (
@@ -64,19 +62,17 @@ func newStack(
 	addr stackaddrs.StackInstance,
 	parent *Stack,
 	config *StackConfig,
-	removedComponents collections.Map[stackaddrs.ConfigComponent, []*RemovedComponent],
-	removedStackCalls collections.Map[stackaddrs.Stack, []*RemovedStackCall],
+	removed *Removed,
 	mode plans.Mode,
 	deferred bool) *Stack {
 	return &Stack{
-		parent:                    parent,
-		config:                    config,
-		addr:                      addr,
-		deferred:                  deferred,
-		mode:                      mode,
-		main:                      main,
-		externalRemovedComponents: removedComponents,
-		externalRemovedStackCalls: removedStackCalls,
+		parent:   parent,
+		config:   config,
+		addr:     addr,
+		deferred: deferred,
+		mode:     mode,
+		main:     main,
+		removed:  removed,
 	}
 }
 
@@ -95,7 +91,7 @@ func (s *Stack) ChildStack(ctx context.Context, addr stackaddrs.StackInstanceSte
 		}
 	}
 
-	calls := s.Removed().localStackCalls[callAddr]
+	calls := s.Removed().stackCalls[callAddr]
 	for _, call := range calls {
 		instances, _ := call.InstancesFor(ctx, s.addr, phase)
 		if instance, exists := instances[addr.Key]; exists {
@@ -221,10 +217,10 @@ func (s *Stack) EmbeddedStackCall(addr stackaddrs.StackCall) *StackCall {
 }
 
 func (s *Stack) RemovedEmbeddedStackCall(addr stackaddrs.StackCall) []*RemovedStackCall {
-	return s.Removed().localStackCalls[addr]
+	return s.Removed().stackCalls[addr]
 }
 
-func (s *Stack) KnownEmbeddedStacks(addr stackaddrs.StackCall, phase EvalPhase) collections.Set[stackaddrs.StackInstance] {
+func (s *Stack) KnownEmbeddedStacks(addr stackaddrs.StackCall, phase EvalPhase) iter.Seq[stackaddrs.StackInstance] {
 	switch phase {
 	case PlanPhase:
 		return s.main.PlanPrevState().StackInstances(stackaddrs.AbsStackCall{
@@ -239,7 +235,7 @@ func (s *Stack) KnownEmbeddedStacks(addr stackaddrs.StackCall, phase EvalPhase)
 	default:
 		// We're not executing with an existing state in the other phases, so
 		// we have no known instances.
-		return collections.NewSet[stackaddrs.StackInstance]()
+		return func(yield func(stackaddrs.StackInstance) bool) {}
 	}
 }
 
@@ -274,29 +270,19 @@ func (s *Stack) Removed() *Removed {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
-	if s.removed != nil {
+	if s.removedInitialised {
 		return s.removed
 	}
 
 	// otherwise we're going to initialise removed.
 
-	stackCallComponents := collections.NewMap[stackaddrs.ConfigComponent, []*RemovedComponent]()
-	localComponents := make(map[stackaddrs.Component][]*RemovedComponent)
-	embeddedStackCalls := collections.NewMap[stackaddrs.Stack, []*RemovedStackCall]()
-	localStackCalls := make(map[stackaddrs.StackCall][]*RemovedStackCall)
-
 	for addr, configs := range s.config.RemovedComponents().All() {
 		blocks := make([]*RemovedComponent, 0, len(configs))
 		for _, config := range configs {
 			blocks = append(blocks, newRemovedComponent(s.main, addr, s, config))
 		}
 
-		if addr.Stack.IsRoot() {
-			localComponents[addr.Item] = blocks
-			continue
-		}
-
-		stackCallComponents.Put(addr, blocks)
+		s.removed.AddComponent(addr, blocks)
 	}
 
 	for addr, configs := range s.config.RemovedStackCalls().All() {
@@ -305,35 +291,15 @@ func (s *Stack) Removed() *Removed {
 			blocks = append(blocks, newRemovedStackCall(s.main, addr, s, config))
 		}
 
-		if len(addr) == 1 {
-			localStackCalls[stackaddrs.StackCall{Name: addr[0].Name}] = blocks
-			continue
-		}
-		embeddedStackCalls.Put(addr, blocks)
-	}
-
-	for addr, configs := range s.externalRemovedComponents.All() {
-		if addr.Stack.IsRoot() {
-			localComponents[addr.Item] = append(localComponents[addr.Item], configs...)
-			continue
-		}
-		stackCallComponents.Put(addr, append(stackCallComponents.Get(addr), configs...))
-	}
-
-	for addr, configs := range s.externalRemovedStackCalls.All() {
-		if len(addr) == 1 {
-			localStackCalls[stackaddrs.StackCall{Name: addr[0].Name}] = append(localStackCalls[stackaddrs.StackCall{Name: addr[0].Name}], configs...)
-			continue
-		}
-		embeddedStackCalls.Put(addr, append(embeddedStackCalls.Get(addr), configs...))
+		s.removed.AddStackCall(addr, blocks)
 	}
 
-	s.removed = newRemoved(localComponents, stackCallComponents, localStackCalls, embeddedStackCalls)
+	s.removedInitialised = true
 	return s.removed
 }
 
 func (s *Stack) RemovedComponent(addr stackaddrs.Component) []*RemovedComponent {
-	return s.Removed().localComponents[addr]
+	return s.Removed().components[addr]
 }
 
 // ApplyableComponents returns the combination of removed blocks and declared
@@ -344,7 +310,7 @@ func (s *Stack) ApplyableComponents(addr stackaddrs.Component) (*Component, []*R
 
 // KnownComponentInstances returns a set of the component instances that belong
 // to the given component from the current state or plan.
-func (s *Stack) KnownComponentInstances(component stackaddrs.Component, phase EvalPhase) collections.Set[stackaddrs.ComponentInstance] {
+func (s *Stack) KnownComponentInstances(component stackaddrs.Component, phase EvalPhase) iter.Seq[stackaddrs.ComponentInstance] {
 	switch phase {
 	case PlanPhase:
 		return s.main.PlanPrevState().ComponentInstances(stackaddrs.AbsComponent{
@@ -352,14 +318,14 @@ func (s *Stack) KnownComponentInstances(component stackaddrs.Component, phase Ev
 			Item:  component,
 		})
 	case ApplyPhase:
-		return s.main.PlanBeingApplied().ComponentInstances(stackaddrs.AbsComponent{
+		return s.main.PlanBeingApplied().ComponentInstanceAddresses(stackaddrs.AbsComponent{
 			Stack: s.addr,
 			Item:  component,
 		})
 	default:
 		// We're not executing with an existing state in the other phases, so
 		// we have no known instances.
-		return collections.NewSet[stackaddrs.ComponentInstance]()
+		return func(yield func(stackaddrs.ComponentInstance) bool) {}
 	}
 }
 
@@ -655,7 +621,7 @@ func (s *Stack) PlanChanges(ctx context.Context) ([]stackplan.PlannedChange, tfd
 
 	// We're going to validate that all the removed blocks in this stack resolve
 	// to unique instance addresses.
-	for _, blocks := range s.Removed().localComponents {
+	for _, blocks := range s.Removed().components {
 		seen := make(map[addrs.InstanceKey]*RemovedComponentInstance)
 		for _, block := range blocks {
 			insts, unknown := block.InstancesFor(ctx, s.addr, PlanPhase)
@@ -678,7 +644,7 @@ func (s *Stack) PlanChanges(ctx context.Context) ([]stackplan.PlannedChange, tfd
 		}
 	}
 
-	for _, blocks := range s.Removed().localStackCalls {
+	for _, blocks := range s.Removed().stackCalls {
 		seen := collections.NewMap[stackaddrs.StackInstance, *RemovedStackCallInstance]()
 		for _, block := range blocks {
 			insts, unknown := block.InstancesFor(ctx, s.addr, PlanPhase)
@@ -717,7 +683,7 @@ func (s *Stack) PlanChanges(ctx context.Context) ([]stackplan.PlannedChange, tfd
 
 	var changes []stackplan.PlannedChange
 Instance:
-	for inst := range s.main.PlanPrevState().AllComponentInstances().All() {
+	for inst := range s.main.PlanPrevState().AllComponentInstances() {
 
 		// We track here whether this component instance has any associated
 		// resources. If this component is empty, and not referenced in the

internal/stacks/stackruntime/internal/stackeval/stack_call.go

@@ -10,7 +10,6 @@ import (
 	"github.com/zclconf/go-cty/cty"
 
 	"github.com/hashicorp/terraform/internal/addrs"
-	"github.com/hashicorp/terraform/internal/collections"
 	"github.com/hashicorp/terraform/internal/instances"
 	"github.com/hashicorp/terraform/internal/promising"
 	"github.com/hashicorp/terraform/internal/stacks/stackaddrs"
@@ -47,8 +46,8 @@ func newStackCall(main *Main, addr stackaddrs.AbsStackCall, stack *Stack, config
 
 // GetExternalRemovedBlocks fetches the removed blocks that target the stack
 // instances being created by this stack call.
-func (c *StackCall) GetExternalRemovedBlocks() (collections.Map[stackaddrs.ConfigComponent, []*RemovedComponent], collections.Map[stackaddrs.Stack, []*RemovedStackCall]) {
-	return c.stack.Removed().ForStackCall(c.addr.Item)
+func (c *StackCall) GetExternalRemovedBlocks() *Removed {
+	return c.stack.Removed().Next(c.addr.Item.Name)
 }
 
 // ForEachValue returns the result of evaluating the "for_each" expression

internal/stacks/stackruntime/internal/stackeval/stack_call_instance.go

@@ -72,8 +72,7 @@ func (c *StackCallInstance) CalledStackAddr() stackaddrs.StackInstance {
 
 func (c *StackCallInstance) Stack(ctx context.Context, phase EvalPhase) *Stack {
 	stack, err := c.stack.For(phase).Do(ctx, c.tracingName(), func(ctx context.Context) (*Stack, error) {
-		components, embeddedStackCalls := c.call.GetExternalRemovedBlocks()
-		return newStack(c.main, c.CalledStackAddr(), c.call.stack, c.call.config.TargetConfig(), components, embeddedStackCalls, c.mode, c.deferred), nil
+		return newStack(c.main, c.CalledStackAddr(), c.call.stack, c.call.config.TargetConfig(), c.call.GetExternalRemovedBlocks(), c.mode, c.deferred), nil
 	})
 	if err != nil {
 		// we don't have cycles in here, and we don't return an error so this

internal/stacks/stackruntime/internal/stackeval/stack_config.go

@@ -43,7 +43,7 @@ type StackConfig struct {
 	localValues       map[stackaddrs.LocalValue]*LocalValueConfig
 	outputValues      map[stackaddrs.OutputValue]*OutputValueConfig
 	stackCalls        map[stackaddrs.StackCall]*StackCallConfig
-	removedStackCalls collections.Map[stackaddrs.Stack, []*RemovedStackCallConfig]
+	removedStackCalls collections.Map[stackaddrs.ConfigStackCall, []*RemovedStackCallConfig]
 	components        map[stackaddrs.Component]*ComponentConfig
 	removedComponents collections.Map[stackaddrs.ConfigComponent, []*RemovedComponentConfig]
 	providers         map[stackaddrs.ProviderConfig]*ProviderConfig
@@ -65,7 +65,7 @@ func newStackConfig(main *Main, addr stackaddrs.Stack, parent *StackConfig, conf
 		localValues:       make(map[stackaddrs.LocalValue]*LocalValueConfig, len(config.Stack.Declarations.LocalValues)),
 		outputValues:      make(map[stackaddrs.OutputValue]*OutputValueConfig, len(config.Stack.Declarations.OutputValues)),
 		stackCalls:        make(map[stackaddrs.StackCall]*StackCallConfig, len(config.Stack.Declarations.EmbeddedStacks)),
-		removedStackCalls: collections.NewMap[stackaddrs.Stack, []*RemovedStackCallConfig](),
+		removedStackCalls: collections.NewMap[stackaddrs.ConfigStackCall, []*RemovedStackCallConfig](),
 		components:        make(map[stackaddrs.Component]*ComponentConfig, len(config.Stack.Declarations.Components)),
 		removedComponents: collections.NewMap[stackaddrs.ConfigComponent, []*RemovedComponentConfig](),
 		providers:         make(map[stackaddrs.ProviderConfig]*ProviderConfig, len(config.Stack.Declarations.ProviderConfigs)),
@@ -361,14 +361,14 @@ func (s *StackConfig) StackCalls() map[stackaddrs.StackCall]*StackCallConfig {
 	return ret
 }
 
-func (s *StackConfig) RemovedStackCall(addr stackaddrs.Stack) []*RemovedStackCallConfig {
+func (s *StackConfig) RemovedStackCall(addr stackaddrs.ConfigStackCall) []*RemovedStackCallConfig {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
 	ret, ok := s.removedStackCalls.GetOk(addr)
 	if !ok {
 		for _, cfg := range s.config.Stack.RemovedEmbeddedStacks.Get(addr) {
-			removed := newRemovedStackCallConfig(s.main, append(s.addr, addr...), s, cfg)
+			removed := newRemovedStackCallConfig(s.main, addr, s, cfg)
 			ret = append(ret, removed)
 		}
 		s.removedStackCalls.Put(addr, ret)
@@ -376,8 +376,8 @@ func (s *StackConfig) RemovedStackCall(addr stackaddrs.Stack) []*RemovedStackCal
 	return ret
 }
 
-func (s *StackConfig) RemovedStackCalls() collections.Map[stackaddrs.Stack, []*RemovedStackCallConfig] {
-	ret := collections.NewMap[stackaddrs.Stack, []*RemovedStackCallConfig]()
+func (s *StackConfig) RemovedStackCalls() collections.Map[stackaddrs.ConfigStackCall, []*RemovedStackCallConfig] {
+	ret := collections.NewMap[stackaddrs.ConfigStackCall, []*RemovedStackCallConfig]()
 	for addr := range s.config.Stack.RemovedEmbeddedStacks.All() {
 		ret.Put(addr, s.RemovedStackCall(addr))
 	}

internal/stacks/stackruntime/internal/stackeval/walk_dynamic.go

@@ -60,7 +60,7 @@ func walkDynamicObjectsInStack[Output any](
 	for call := range stack.EmbeddedStackCalls() {
 		walkEmbeddedStack(ctx, walk, stack, call, phase, visit)
 	}
-	for call := range stack.Removed().localStackCalls {
+	for call := range stack.Removed().stackCalls {
 		if stack.EmbeddedStackCall(call) != nil {
 			continue
 		}
@@ -70,7 +70,7 @@ func walkDynamicObjectsInStack[Output any](
 	for component := range stack.Components() {
 		walkComponent(ctx, walk, stack, component, phase, visit)
 	}
-	for component := range stack.Removed().localComponents {
+	for component := range stack.Removed().components {
 		if stack.Component(component) != nil {
 			continue // then we processed this as part of the component stage
 		}
@@ -180,7 +180,7 @@ func walkComponent[Output any](
 		})
 	}
 
-	for _, block := range stack.Removed().localComponents[addr] {
+	for _, block := range stack.Removed().components[addr] {
 		visit(ctx, walk, block) // first, just visit the removed block directly
 
 		wg.Add(1)
@@ -233,7 +233,7 @@ func walkComponent[Output any](
 		unknownComponentBlockClaimedSomething := unknownComponentBlock == nil
 
 		knownInstances := stack.KnownComponentInstances(addr, phase)
-		for inst := range knownInstances.All() {
+		for inst := range knownInstances {
 			if claimedInstances.Has(inst) {
 				// don't need the mutex any more since this will be fully
 				// initialised when all the wait groups are finished.
@@ -348,7 +348,7 @@ func walkEmbeddedStack[Output any](
 		})
 	}
 
-	for _, block := range stack.Removed().localStackCalls[addr] {
+	for _, block := range stack.Removed().stackCalls[addr] {
 		visit(ctx, walk, block)
 
 		wg.Add(1)
@@ -386,7 +386,7 @@ func walkEmbeddedStack[Output any](
 		unknownStackCallClaimedSomething := unknownStackCall == nil
 
 		knownStacks := stack.KnownEmbeddedStacks(addr, phase)
-		for inst := range knownStacks.All() {
+		for inst := range knownStacks {
 			if claimedInstances.Has(inst) {
 				continue
 			}

internal/stacks/stackstate/state.go

@@ -4,6 +4,8 @@
 package stackstate
 
 import (
+	"iter"
+
 	"github.com/zclconf/go-cty/cty"
 	"google.golang.org/protobuf/types/known/anypb"
 
@@ -21,9 +23,9 @@ import (
 // not be modified after it's been constructed; results of planning or applying
 // changes are represented in other ways inside the stacks language runtime.
 type State struct {
-	componentInstances collections.Map[stackaddrs.AbsComponentInstance, *componentInstanceState]
-	outputs            map[stackaddrs.OutputValue]cty.Value
-	inputs             map[stackaddrs.InputVariable]cty.Value
+	root    *stackInstanceState
+	outputs map[stackaddrs.OutputValue]cty.Value
+	inputs  map[stackaddrs.InputVariable]cty.Value
 
 	// discardUnsupportedKeys is the set of state keys that we encountered
 	// during decoding which are of types that are not supported by this
@@ -39,7 +41,7 @@ type State struct {
 // NewState constructs a new, empty state.
 func NewState() *State {
 	return &State{
-		componentInstances:     collections.NewMap[stackaddrs.AbsComponentInstance, *componentInstanceState](),
+		root:                   newStackInstanceState(stackaddrs.RootStackInstance),
 		outputs:                make(map[stackaddrs.OutputValue]cty.Value),
 		inputs:                 make(map[stackaddrs.InputVariable]cty.Value),
 		discardUnsupportedKeys: statekeys.NewKeySet(),
@@ -74,7 +76,11 @@ func (s *State) RootOutputValue(addr stackaddrs.OutputValue) cty.Value {
 }
 
 func (s *State) HasComponentInstance(addr stackaddrs.AbsComponentInstance) bool {
-	return s.componentInstances.HasKey(addr)
+	stack := s.root.getDescendent(addr.Stack)
+	if stack == nil {
+		return false
+	}
+	return stack.getComponentInstance(addr.Item) != nil
 }
 
 // AllComponentInstances returns a set of addresses for all of the component
@@ -86,16 +92,10 @@ func (s *State) HasComponentInstance(addr stackaddrs.AbsComponentInstance) bool
 // instance record tracked in raw state, but it can potentially be absent in
 // exceptional cases such as if Terraform Core crashed partway through the
 // previous run.
-func (s *State) AllComponentInstances() collections.Set[stackaddrs.AbsComponentInstance] {
-	var ret collections.Set[stackaddrs.AbsComponentInstance]
-	if s.componentInstances.Len() == 0 {
-		return ret
-	}
-	ret = collections.NewSet[stackaddrs.AbsComponentInstance]()
-	for key := range s.componentInstances.All() {
-		ret.Add(key)
+func (s *State) AllComponentInstances() iter.Seq[stackaddrs.AbsComponentInstance] {
+	return func(yield func(stackaddrs.AbsComponentInstance) bool) {
+		s.root.iterate(yield)
 	}
-	return ret
 }
 
 // ComponentInstances returns the set of component instances that belong to the
@@ -103,46 +103,43 @@ func (s *State) AllComponentInstances() collections.Set[stackaddrs.AbsComponentI
 // state.
 //
 // This will always be a subset of AllComponentInstances.
-func (s *State) ComponentInstances(addr stackaddrs.AbsComponent) collections.Set[stackaddrs.ComponentInstance] {
-	ret := collections.NewSet[stackaddrs.ComponentInstance]()
-	for key := range s.componentInstances.All() {
-		if key.Stack.String() != addr.Stack.String() {
-			// Then
-			continue
+func (s *State) ComponentInstances(addr stackaddrs.AbsComponent) iter.Seq[stackaddrs.ComponentInstance] {
+	return func(yield func(stackaddrs.ComponentInstance) bool) {
+		target := s.root.getDescendent(addr.Stack)
+		if target == nil {
+			return
 		}
-		if key.Item.Component.Name != addr.Item.Name {
-			continue
+
+		for key := range target.components[addr.Item] {
+			yield(stackaddrs.ComponentInstance{
+				Component: addr.Item,
+				Key:       key,
+			})
 		}
-		ret.Add(key.Item)
 	}
-	return ret
 }
 
 // StackInstances returns the set of known stack instances for the given stack
 // call.
-func (s *State) StackInstances(call stackaddrs.AbsStackCall) collections.Set[stackaddrs.StackInstance] {
-	ret := collections.NewSet[stackaddrs.StackInstance]()
-	for key := range s.componentInstances.All() {
-		if len(key.Stack) == 0 {
-			continue
+func (s *State) StackInstances(call stackaddrs.AbsStackCall) iter.Seq[stackaddrs.StackInstance] {
+	return func(yield func(stackaddrs.StackInstance) bool) {
+		target := s.root.getDescendent(call.Stack)
+		if target == nil {
+			return
 		}
 
-		last := key.Stack[len(key.Stack)-1]
-		path := key.Stack[:len(key.Stack)-1]
-
-		if path.String() != call.Stack.String() {
-			continue
+		for _, stack := range target.children[call.Item.Name] {
+			yield(stack.address)
 		}
-		if last.Name != call.Item.Name {
-			continue
-		}
-		ret.Add(key.Stack)
 	}
-	return ret
 }
 
 func (s *State) componentInstanceState(addr stackaddrs.AbsComponentInstance) *componentInstanceState {
-	return s.componentInstances.Get(addr)
+	target := s.root.getDescendent(addr.Stack)
+	if target == nil {
+		return nil
+	}
+	return target.getComponentInstance(addr.Item)
 }
 
 // DependenciesForComponent returns the list of components that are required by
@@ -216,8 +213,8 @@ func (s *State) IdentitiesForComponent(addr stackaddrs.AbsComponentInstance) map
 // with the given address.
 func (s *State) ComponentInstanceResourceInstanceObjects(addr stackaddrs.AbsComponentInstance) collections.Set[stackaddrs.AbsResourceInstanceObject] {
 	var ret collections.Set[stackaddrs.AbsResourceInstanceObject]
-	cs, ok := s.componentInstances.GetOk(addr)
-	if !ok {
+	cs := s.componentInstanceState(addr)
+	if cs == nil {
 		return ret
 	}
 	ret = collections.NewSet[stackaddrs.AbsResourceInstanceObject]()
@@ -231,23 +228,6 @@ func (s *State) ComponentInstanceResourceInstanceObjects(addr stackaddrs.AbsComp
 	return ret
 }
 
-// AllResourceInstanceObjects returns a set of addresses for all of the resource
-// 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 key, elem := range s.componentInstances.All() {
-		componentAddr := key
-		for _, elem := range elem.resourceInstanceObjects.Elems {
-			objKey := stackaddrs.AbsResourceInstanceObject{
-				Component: componentAddr,
-				Item:      elem.Key,
-			}
-			ret.Add(objKey)
-		}
-	}
-	return ret
-}
-
 // ResourceInstanceObjectSrc returns the source (i.e. still encoded) version of
 // the resource instance object for the given address, or nil if no such
 // object is tracked in the state.
@@ -267,8 +247,8 @@ func (s *State) ResourceInstanceObjectSrc(addr stackaddrs.AbsResourceInstanceObj
 // function that operates on the configuration of a component instance rather
 // than the state of one.
 func (s *State) RequiredProviderInstances(component stackaddrs.AbsComponentInstance) addrs.Set[addrs.RootProviderConfig] {
-	state, ok := s.componentInstances.GetOk(component)
-	if !ok {
+	state := s.componentInstanceState(component)
+	if state == nil {
 		// Then we have no state for this component, which is fine.
 		return addrs.MakeSet[addrs.RootProviderConfig]()
 	}
@@ -284,8 +264,8 @@ func (s *State) RequiredProviderInstances(component stackaddrs.AbsComponentInsta
 }
 
 func (s *State) resourceInstanceObjectState(addr stackaddrs.AbsResourceInstanceObject) *resourceInstanceObjectState {
-	cs, ok := s.componentInstances.GetOk(addr.Component)
-	if !ok {
+	cs := s.componentInstanceState(addr.Component)
+	if cs == nil {
 		return nil
 	}
 	return cs.resourceInstanceObjects.Get(addr.Item)
@@ -346,17 +326,40 @@ func (s *State) addInputVariable(addr stackaddrs.InputVariable, value cty.Value)
 }
 
 func (s *State) ensureComponentInstanceState(addr stackaddrs.AbsComponentInstance) *componentInstanceState {
-	if existing, ok := s.componentInstances.GetOk(addr); ok {
-		return existing
+	current := s.root
+	for _, step := range addr.Stack {
+		next := current.getChild(step)
+		if next == nil {
+			next = newStackInstanceState(append(current.address, step))
+
+			children, ok := current.children[step.Name]
+			if !ok {
+				children = make(map[addrs.InstanceKey]*stackInstanceState)
+			}
+			children[step.Key] = next
+			current.children[step.Name] = children
+		}
+		current = next
 	}
-	s.componentInstances.Put(addr, &componentInstanceState{
-		dependencies:            collections.NewSet[stackaddrs.AbsComponent](),
-		dependents:              collections.NewSet[stackaddrs.AbsComponent](),
-		outputValues:            make(map[addrs.OutputValue]cty.Value),
-		inputVariables:          make(map[addrs.InputVariable]cty.Value),
-		resourceInstanceObjects: addrs.MakeMap[addrs.AbsResourceInstanceObject, *resourceInstanceObjectState](),
-	})
-	return s.componentInstances.Get(addr)
+
+	component := current.getComponentInstance(addr.Item)
+	if component == nil {
+		component = &componentInstanceState{
+			dependencies:            collections.NewSet[stackaddrs.AbsComponent](),
+			dependents:              collections.NewSet[stackaddrs.AbsComponent](),
+			outputValues:            make(map[addrs.OutputValue]cty.Value),
+			inputVariables:          make(map[addrs.InputVariable]cty.Value),
+			resourceInstanceObjects: addrs.MakeMap[addrs.AbsResourceInstanceObject, *resourceInstanceObjectState](),
+		}
+
+		components, ok := current.components[addr.Item.Component]
+		if !ok {
+			components = make(map[addrs.InstanceKey]*componentInstanceState)
+		}
+		components[addr.Item.Key] = component
+		current.components[addr.Item.Component] = components
+	}
+	return component
 }
 
 func (s *State) addResourceInstanceObject(addr stackaddrs.AbsResourceInstanceObject, src *states.ResourceInstanceObjectSrc, providerConfigAddr addrs.AbsProviderConfig) {
@@ -394,3 +397,72 @@ type resourceInstanceObjectState struct {
 	src                *states.ResourceInstanceObjectSrc
 	providerConfigAddr addrs.AbsProviderConfig
 }
+
+type stackInstanceState struct {
+	address    stackaddrs.StackInstance
+	components map[stackaddrs.Component]map[addrs.InstanceKey]*componentInstanceState
+	children   map[string]map[addrs.InstanceKey]*stackInstanceState
+}
+
+func newStackInstanceState(address stackaddrs.StackInstance) *stackInstanceState {
+	return &stackInstanceState{
+		address:    address,
+		components: make(map[stackaddrs.Component]map[addrs.InstanceKey]*componentInstanceState),
+		children:   make(map[string]map[addrs.InstanceKey]*stackInstanceState),
+	}
+}
+
+func (s *stackInstanceState) getDescendent(stack stackaddrs.StackInstance) *stackInstanceState {
+	if len(stack) == 0 {
+		return s
+	}
+
+	next := s.getChild(stack[0])
+	if next == nil {
+		return nil
+	}
+	return next.getDescendent(stack[1:])
+}
+
+func (s *stackInstanceState) getChild(step stackaddrs.StackInstanceStep) *stackInstanceState {
+	stacks, ok := s.children[step.Name]
+	if !ok {
+		return nil
+	}
+	return stacks[step.Key]
+}
+
+func (s *stackInstanceState) getComponentInstance(component stackaddrs.ComponentInstance) *componentInstanceState {
+	components, ok := s.components[component.Component]
+	if !ok {
+		return nil
+	}
+	return components[component.Key]
+}
+
+func (s *stackInstanceState) iterate(yield func(stackaddrs.AbsComponentInstance) bool) bool {
+	for component, components := range s.components {
+		for key := range components {
+			proceed := yield(stackaddrs.AbsComponentInstance{
+				Stack: s.address,
+				Item: stackaddrs.ComponentInstance{
+					Component: component,
+					Key:       key,
+				},
+			})
+			if !proceed {
+				return false
+			}
+		}
+	}
+
+	for _, children := range s.children {
+		for _, child := range children {
+			if !child.iterate(yield) {
+				return false
+			}
+		}
+	}
+
+	return true
+}