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

Author: liamcervante

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"),