stackstate: Initial work on decoding and the state model

This is a sketch of the overall structure of the prior state decoder and
the model type it populates.

Before we can complete this we'll need to slightly rework how the apply
phase emits the raw events that this is consuming, and in particular to
change the raw state representation to be JSON-based to match with how
Terraform Core expects to receive it once reloaded. That will follow in
later commits.

internal/stacks/stackaddrs/in_component.go

@@ -79,6 +79,10 @@ var _ collections.UniqueKeyer[AbsResource] = AbsResource{}
 // particular component instance.
 type AbsResourceInstance = InAbsComponentInstance[addrs.AbsResourceInstance]
 
+// AbsResourceInstanceObject represents an object associated with an instance
+// of a resource from inside a particular component instance.
+type AbsResourceInstanceObject = InAbsComponentInstance[addrs.AbsResourceInstanceObject]
+
 // AbsModuleInstance represents an instance of a module from inside a
 // particular component instance.
 //

internal/stacks/stackstate/from_proto.go

@@ -0,0 +1,132 @@
+// Copyright (c) HashiCorp, Inc.
+// SPDX-License-Identifier: BUSL-1.1
+
+package stackstate
+
+import (
+	"fmt"
+
+	"github.com/hashicorp/terraform/internal/addrs"
+	"github.com/hashicorp/terraform/internal/stacks/stackaddrs"
+	"github.com/hashicorp/terraform/internal/stacks/stackstate/statekeys"
+	"github.com/hashicorp/terraform/internal/stacks/tfstackdata1"
+	"google.golang.org/protobuf/proto"
+	"google.golang.org/protobuf/reflect/protoreflect"
+	"google.golang.org/protobuf/types/known/anypb"
+)
+
+func LoadFromProto(msgs map[string]*anypb.Any) (*State, error) {
+	ret := NewState()
+	for rawKey, rawMsg := range msgs {
+		key, err := statekeys.Parse(rawKey)
+		if err != nil {
+			// "invalid" here means that it was either not syntactically
+			// valid at all or was a recognized type but with the wrong
+			// syntax for that type.
+			// An unrecognized key type is NOT invalid; we handle that below.
+			return nil, fmt.Errorf("invalid tracking key %q in state: %w", rawKey, err)
+		}
+
+		if !statekeys.RecognizedType(key) {
+			// There are three different strategies for dealing with
+			// unrecognized keys, which we recognize based on naming
+			// conventions of the key types.
+			switch handling := key.KeyType().UnrecognizedKeyHandling(); handling {
+
+			case statekeys.FailIfUnrecognized:
+				// This is for keys whose messages materially change the
+				// meaning of the state and so cannot be ignored. Keys
+				// with this treatment are forwards-incompatible (old versions
+				// of Terraform will fail to load a state containing them) so
+				// should be added only as a last resort.
+				return nil, fmt.Errorf("state was created by a newer version of Terraform Core (unrecognized tracking key %q)", rawKey)
+
+			case statekeys.PreserveIfUnrecognized:
+				// This is for keys whose messages can safely be left entirely
+				// unchanged if applying a plan with a version of Terraform
+				// that doesn't understand them. Keys in this category should
+				// typically be standalone and not refer to or depend on any
+				// other objects in the state, to ensure that removing or
+				// updating other objects will not cause the preserved message
+				// to become misleading or invalid.
+				// We don't need to do anything special with these ones because
+				// the caller should preserve any object we don't explicitly
+				// update or delete during the apply phase.
+
+			case statekeys.DiscardIfUnrecognized:
+				// This is for keys which can be discarded when planning or
+				// applying with an older version of Terraform that doesn't
+				// understand them. This category is for optional ancillary
+				// information -- not actually required for correct subsequent
+				// planning -- especially if it could be recomputed again and
+				// repopulated if later planning and applying with a newer
+				// version of Terraform Core.
+				// For these ones we need to remember their keys so that we
+				// can emit "delete" messages early in the apply phase to
+				// actually discard them from the caller's records.
+				ret.discardUnsupportedKeys.Add(key)
+
+			default:
+				// Should not get here. The above should be exhaustive.
+				panic(fmt.Sprintf("unsupported UnrecognizedKeyHandling value %s", handling))
+			}
+		}
+
+		msg, err := anypb.UnmarshalNew(rawMsg, proto.UnmarshalOptions{})
+		if err != nil {
+			return nil, fmt.Errorf("invalid raw value for raw state key %q: %w", rawKey, err)
+		}
+
+		switch key := key.(type) {
+
+		case statekeys.ComponentInstance:
+			err := handleComponentInstanceMsg(key, msg, ret)
+			if err != nil {
+				return nil, err
+			}
+
+		case statekeys.ResourceInstanceObject:
+			err := handleResourceInstanceObjectMsg(key, msg, ret)
+			if err != nil {
+				return nil, err
+			}
+
+		default:
+			// Should not get here: the above should be exhaustive for all
+			// possible key types.
+			panic(fmt.Sprintf("unsupported state key type %T", key))
+		}
+	}
+	return ret, nil
+}
+
+func handleComponentInstanceMsg(key statekeys.ComponentInstance, msg protoreflect.ProtoMessage, state *State) error {
+	// For this particular object type all of the information is in the key,
+	// for now at least.
+	_, ok := msg.(*tfstackdata1.StateComponentInstanceV1)
+	if !ok {
+		return fmt.Errorf("unsupported message type %T for %s state", msg, key.ComponentInstanceAddr)
+	}
+
+	state.ensureComponentInstanceState(key.ComponentInstanceAddr)
+	return nil
+}
+
+func handleResourceInstanceObjectMsg(key statekeys.ResourceInstanceObject, msg protoreflect.ProtoMessage, state *State) error {
+	fullAddr := stackaddrs.AbsResourceInstanceObject{
+		Component: key.ResourceInstance.Component,
+		Item: addrs.AbsResourceInstanceObject{
+			ResourceInstance: key.ResourceInstance.Item,
+			DeposedKey:       key.DeposedKey,
+		},
+	}
+
+	riMsg, ok := msg.(*tfstackdata1.StateResourceInstanceObjectV1)
+	if !ok {
+		return fmt.Errorf("unsupported message type %T for state of %s", msg, fullAddr.String())
+	}
+
+	// TODO: Implement this
+	_ = riMsg
+	return fmt.Errorf("resource instance object state decoding not yet implemented")
+}

internal/stacks/stackstate/state.go

@@ -0,0 +1,183 @@
+// Copyright (c) HashiCorp, Inc.
+// SPDX-License-Identifier: BUSL-1.1
+
+package stackstate
+
+import (
+	"github.com/hashicorp/terraform/internal/addrs"
+	"github.com/hashicorp/terraform/internal/collections"
+	"github.com/hashicorp/terraform/internal/stacks/stackaddrs"
+	"github.com/hashicorp/terraform/internal/stacks/stackstate/statekeys"
+	"github.com/hashicorp/terraform/internal/states"
+)
+
+// State represents a previous run's state snapshot.
+//
+// Unlike [states.State] and its associates, State is an immutable data
+// structure constructed to represent only the previous run state. It should
+// 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]
+
+	// discardUnsupportedKeys is the set of state keys that we encountered
+	// during decoding which are of types that are not supported by this
+	// version of Terraform, if and only if they are of a type which is
+	// specified as being discarded when unrecognized. We should emit
+	// events during the apply phase to delete the objects associated with
+	// these keys.
+	discardUnsupportedKeys collections.Set[statekeys.Key]
+}
+
+// NewState constructs a new, empty state.
+func NewState() *State {
+	return &State{
+		componentInstances:     collections.NewMap[stackaddrs.AbsComponentInstance, *componentInstanceState](),
+		discardUnsupportedKeys: statekeys.NewKeySet(),
+	}
+}
+
+// AllComponentInstances returns a set of addresses for all of the component
+// instances that are tracked in the state.
+//
+// This includes both instances that were explicitly represented in the source
+// raw state _and_ any that were missing but implied by a resource instance
+// existing inside them. There should typically be an explicit component
+// 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 _, elem := range s.componentInstances.Elems() {
+		ret.Add(elem.K)
+	}
+	return ret
+}
+
+func (s *State) componentInstanceState(addr stackaddrs.AbsComponentInstance) *componentInstanceState {
+	return s.componentInstances.Get(addr)
+}
+
+// ComponentInstanceResourceInstanceObjects returns a set of addresses for
+// all of the resource instance objects belonging to the component instance
+// 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 {
+		return ret
+	}
+	ret = collections.NewSet[stackaddrs.AbsResourceInstanceObject]()
+	for _, elem := range cs.resourceInstanceObjects.Elems {
+		objKey := stackaddrs.AbsResourceInstanceObject{
+			Component: addr,
+			Item:      elem.Key,
+		}
+		ret.Add(objKey)
+	}
+	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 _, elem := range s.componentInstances.Elems() {
+		componentAddr := elem.K
+		for _, elem := range elem.V.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.
+func (s *State) ResourceInstanceObjectSrc(addr stackaddrs.AbsResourceInstanceObject) *states.ResourceInstanceObjectSrc {
+	rios := s.resourceInstanceObjectState(addr)
+	if rios == nil {
+		return nil
+	}
+	return rios.src
+}
+
+func (s *State) resourceInstanceObjectState(addr stackaddrs.AbsResourceInstanceObject) *resourceInstanceObjectState {
+	cs, ok := s.componentInstances.GetOk(addr.Component)
+	if !ok {
+		return nil
+	}
+	return cs.resourceInstanceObjects.Get(addr.Item)
+}
+
+// ComponentInstanceStateForModulesRuntime returns a [states.State]
+// representation of the objects tracked for the given component instance.
+//
+// This produces only a very bare-bones [states.State] that should be
+// sufficient for use as a prior state for the modules runtime's plan function
+// to consider, but likely won't be of much other use.
+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() {
+			rios := s.resourceInstanceObjectState(objAddr)
+
+			if objAddr.Item.IsCurrent() {
+				ss.SetResourceInstanceCurrent(
+					objAddr.Item.ResourceInstance,
+					rios.src, rios.providerConfigAddr,
+				)
+			} else {
+				ss.SetResourceInstanceDeposed(
+					objAddr.Item.ResourceInstance, objAddr.Item.DeposedKey,
+					rios.src, rios.providerConfigAddr,
+				)
+			}
+		}
+	})
+}
+
+// RawKeysToDiscard returns a set of raw state keys that the apply phase should
+// emit "delete" events for to remove objects from the raw state map that
+// will no longer be relevant or meaningful after this plan is applied.
+//
+// Do not modify the returned set.
+func (s *State) RawKeysToDiscard() collections.Set[statekeys.Key] {
+	return s.discardUnsupportedKeys
+}
+
+func (s *State) ensureComponentInstanceState(addr stackaddrs.AbsComponentInstance) *componentInstanceState {
+	if existing, ok := s.componentInstances.GetOk(addr); ok {
+		return existing
+	}
+	s.componentInstances.Put(addr, &componentInstanceState{
+		resourceInstanceObjects: addrs.MakeMap[addrs.AbsResourceInstanceObject, *resourceInstanceObjectState](),
+	})
+	return s.componentInstances.Get(addr)
+}
+
+func (s *State) addResourceInstanceObject(addr stackaddrs.AbsResourceInstanceObject, src *states.ResourceInstanceObjectSrc, providerConfigAddr addrs.AbsProviderConfig) {
+	cs := s.ensureComponentInstanceState(addr.Component)
+
+	cs.resourceInstanceObjects.Put(addr.Item, &resourceInstanceObjectState{
+		src:                src,
+		providerConfigAddr: providerConfigAddr,
+	})
+}
+
+type componentInstanceState struct {
+	resourceInstanceObjects addrs.Map[addrs.AbsResourceInstanceObject, *resourceInstanceObjectState]
+}
+
+type resourceInstanceObjectState struct {
+	src                *states.ResourceInstanceObjectSrc
+	providerConfigAddr addrs.AbsProviderConfig
+}

internal/stacks/stackstate/statekeys/collections.go

@@ -0,0 +1,34 @@
+// Copyright (c) HashiCorp, Inc.
+// SPDX-License-Identifier: BUSL-1.1
+
+package statekeys
+
+import (
+	"github.com/hashicorp/terraform/internal/collections"
+)
+
+type KeySet collections.Set[Key]
+
+// NewKeySet returns an initialized set of [Key] that's ready to use and
+// treats two keys as unique if they have the same string representation.
+func NewKeySet() collections.Set[Key] {
+	return collections.NewSetFunc[Key](stateKeyUniqueKey)
+}
+
+// NewKeyMap returns an initialized map from [Key] to V that's ready to use and
+// treats two keys as unique if they have the same string representation.
+func NewKeyMap[V any]() collections.Map[Key, V] {
+	return collections.NewMapFunc[Key, V](stateKeyUniqueKey)
+}
+
+// stateKeyCollectionsKey is an internal adapter so that [statekeys.Key] values
+// can be used as [collections.Set] elements and [collections.Map] keys.
+type stateKeyCollectionsKey string
+
+// IsUniqueKey implements collections.UniqueKey.
+func (stateKeyCollectionsKey) IsUniqueKey(Key) {
+}
+
+func stateKeyUniqueKey(k Key) collections.UniqueKey[Key] {
+	return stateKeyCollectionsKey(String(k))
+}

internal/stacks/stackstate/statekeys/key_parse_test.go

@@ -336,6 +336,11 @@ func TestParse(t *testing.T) {
 					if gotAsStr != test.Input {
 						t.Errorf("valid key of type %T did not round-trip\ngot:  %s\nwant: %s", got, gotAsStr, test.Input)
 					}
+					if test.WantUnrecognizedHandling != UnrecognizedKeyHandling(0) {
+						if got, want := got.KeyType().UnrecognizedKeyHandling(), test.WantUnrecognizedHandling; got != want {
+							t.Errorf("unexpected UnrecognizedKeyHandling\ngot:  %s\nwant: %s", got, want)
+						}
+					}
 				} else if err == nil {
 					t.Error("Parse returned nil Key and nil error")
 				}