From 3729e6a705b16943a4a2c252304601f3b1caf894 Mon Sep 17 00:00:00 2001
From: James Bardin <j.bardin@gmail.com>
Date: Fri, 13 Mar 2020 13:05:37 -0400
Subject: [PATCH] update MaybeFixUpResourceInstanceAddressForCount

We need to fixup count by config path, not by absolute path, so update
MaybeFixUpResourceInstanceAddressForCount to work with
addrs.ConfigResource.
---
 states/state.go | 23 +++++++++++++++++++
 states/sync.go  | 60 ++++++++++++++++++++++++++++---------------------
 2 files changed, 57 insertions(+), 26 deletions(-)

diff --git a/states/state.go b/states/state.go
index 1f842359e..7777b9144 100644
--- a/states/state.go
+++ b/states/state.go
@@ -70,6 +70,17 @@ func (s *State) Module(addr addrs.ModuleInstance) *Module {
 	return s.Modules[addr.String()]
 }
 
+// ModuleInstances returns the set of Module states that matches the given path.
+func (s *State) ModuleInstances(addr addrs.Module) []*Module {
+	var ms []*Module
+	for _, m := range s.Modules {
+		if m.Addr.Module().Equal(addr) {
+			ms = append(ms, m)
+		}
+	}
+	return ms
+}
+
 // RemoveModule removes the module with the given address from the state,
 // unless it is the root module. The root module cannot be deleted, and so
 // this method will panic if that is attempted.
@@ -133,6 +144,18 @@ func (s *State) Resource(addr addrs.AbsResource) *Resource {
 	return ms.Resource(addr.Resource)
 }
 
+// Resources returns the set of resources that match the given configuration path.
+func (s *State) Resources(addr addrs.ConfigResource) []*Resource {
+	var ret []*Resource
+	for _, m := range s.ModuleInstances(addr.Module) {
+		r := m.Resource(addr.Resource)
+		if r != nil {
+			ret = append(ret, r)
+		}
+	}
+	return ret
+}
+
 // ResourceInstance returns the state for the resource instance with the given
 // address, or nil if no such resource is tracked in the state.
 func (s *State) ResourceInstance(addr addrs.AbsResourceInstance) *ResourceInstance {
diff --git a/states/sync.go b/states/sync.go
index 47fb16d6e..3a70a5c40 100644
--- a/states/sync.go
+++ b/states/sync.go
@@ -246,40 +246,48 @@ func (s *SyncState) RemoveResourceIfEmpty(addr addrs.AbsResource) bool {
 // The state is modified in-place if necessary, moving a resource instance
 // between the two addresses. The return value is true if a change was made,
 // and false otherwise.
-func (s *SyncState) MaybeFixUpResourceInstanceAddressForCount(addr addrs.AbsResource, countEnabled bool) bool {
+func (s *SyncState) MaybeFixUpResourceInstanceAddressForCount(addr addrs.ConfigResource, countEnabled bool) bool {
 	s.lock.Lock()
 	defer s.lock.Unlock()
 
-	ms := s.state.Module(addr.Module)
-	if ms == nil {
+	// get all modules instances that may match this state
+	modules := s.state.ModuleInstances(addr.Module)
+	if len(modules) == 0 {
 		return false
 	}
 
-	relAddr := addr.Resource
-	rs := ms.Resource(relAddr)
-	if rs == nil {
-		return false
-	}
-	huntKey := addrs.NoKey
-	replaceKey := addrs.InstanceKey(addrs.IntKey(0))
-	if !countEnabled {
-		huntKey, replaceKey = replaceKey, huntKey
+	changed := false
+
+	for _, ms := range modules {
+		relAddr := addr.Resource
+		rs := ms.Resource(relAddr)
+		if rs == nil {
+			continue
+		}
+
+		huntKey := addrs.NoKey
+		replaceKey := addrs.InstanceKey(addrs.IntKey(0))
+		if !countEnabled {
+			huntKey, replaceKey = replaceKey, huntKey
+		}
+
+		is, exists := rs.Instances[huntKey]
+		if !exists {
+			continue
+		}
+
+		if _, exists := rs.Instances[replaceKey]; exists {
+			// If the replacement key also exists then we'll do nothing and keep both.
+			continue
+		}
+
+		// If we get here then we need to "rename" from hunt to replace
+		rs.Instances[replaceKey] = is
+		delete(rs.Instances, huntKey)
+		changed = true
 	}
 
-	is, exists := rs.Instances[huntKey]
-	if !exists {
-		return false
-	}
-
-	if _, exists := rs.Instances[replaceKey]; exists {
-		// If the replacement key also exists then we'll do nothing and keep both.
-		return false
-	}
-
-	// If we get here then we need to "rename" from hunt to replace
-	rs.Instances[replaceKey] = is
-	delete(rs.Instances, huntKey)
-	return true
+	return changed
 }
 
 // SetResourceInstanceCurrent saves the given instance object as the current