terraform: Move diff handling during dynamic expansion

terraform/graph.go

@@ -539,6 +539,8 @@ func graphAddDiff(g *depgraph.Graph, d *ModuleDiff) error {
 			rn.Diff = d
 		}
 
+		// If we are not expanding, then we assign the
+		// instance diff to the resource.
 		var rd *InstanceDiff
 		if rn.ExpandMode == ResourceExpandNone {
 			rd = diffs[0]
@@ -1648,16 +1650,9 @@ func (n *GraphNodeResource) Expand() (*depgraph.Graph, error) {
 		ModulePath: n.Resource.Info.ModulePath,
 	}
 
-	// Determine the nodes to create. If we're just looking for the
+	// Do the initial expansion of the nodes, attaching diffs if
-	// nodes to create, return that.
+	// applicable
-	n.expand(g, count)
+	n.expand(g, count, n.Diff)
-
-	// Add in the diff if we have it
-	if n.Diff != nil {
-		if err := graphAddDiff(g, n.Diff); err != nil {
-			return nil, err
-		}
-	}
 
 	// Add all the variable dependencies
 	graphAddVariableDeps(g)
@@ -1678,7 +1673,7 @@ func (n *GraphNodeResource) Expand() (*depgraph.Graph, error) {
 
 // expand expands this resource and adds the resources to the graph. It
 // adds both create and destroy resources.
-func (n *GraphNodeResource) expand(g *depgraph.Graph, count int) {
+func (n *GraphNodeResource) expand(g *depgraph.Graph, count int, diff *ModuleDiff) {
 	// Create the list of nouns
 	result := make([]*depgraph.Noun, 0, count)
 
@@ -1732,13 +1727,52 @@ func (n *GraphNodeResource) expand(g *depgraph.Graph, count int) {
 			}
 		}
 
+		// Add in the diff if we have it
+		var inDiff *InstanceDiff
+		if diff != nil {
+			// Looup the instance diff
+			if d, ok := diff.Resources[name]; ok {
+				inDiff = d
+			}
+
+			if inDiff == nil {
+				if count == 1 {
+					// If the count is one, check the state for ".0"
+					// appended, which might exist if we go from
+					// count > 1 to count == 1.
+					k := r.Id() + ".0"
+					inDiff = diff.Resources[k]
+				} else if i == 0 {
+					// If count is greater than one, check for state
+					// with just the ID, which might exist if we go
+					// from count == 1 to count > 1
+					inDiff = diff.Resources[r.Id()]
+				}
+			}
+		}
+
 		if state == nil {
 			state = &ResourceState{
 				Type: r.Type,
 			}
 		}
 
-		flags := FlagPrimary
+		if inDiff != nil && n.ExpandMode != ResourceExpandDestroy {
+			// Disable Destroy if we aren't doing a destroy expansion.
+			// There is a seperate expansion for the destruction action.
+			d := new(InstanceDiff)
+			*d = *inDiff
+			inDiff = d
+			inDiff.Destroy = false
+		} else if inDiff != nil && inDiff.Destroy && n.ExpandMode == ResourceExpandDestroy {
+			// If we are doing a destroy, make sure it is exclusively
+			// a destroy, since there is a seperate expansion for the apply
+			inDiff = new(InstanceDiff)
+			inDiff.Destroy = true
+		}
+
+		// Inherit the existing flags!
+		flags := n.Resource.Flags
 		if len(state.Tainted) > 0 {
 			flags |= FlagHasTainted
 		}
@@ -1748,6 +1782,7 @@ func (n *GraphNodeResource) expand(g *depgraph.Graph, count int) {
 		resource.CountIndex = i
 		resource.State = state.Primary
 		resource.Flags = flags
+		resource.Diff = inDiff
 
 		// Add the result
 		result = append(result, &depgraph.Noun{
@@ -1768,6 +1803,7 @@ func (n *GraphNodeResource) expand(g *depgraph.Graph, count int) {
 		resource.Config = NewResourceConfig(nil)
 		resource.State = rs.Primary
 		resource.Flags = FlagOrphan
+		resource.Diff = &InstanceDiff{Destroy: true}
 
 		noun := &depgraph.Noun{
 			Name: k,
@@ -1814,7 +1850,8 @@ func (n *GraphNodeResource) filterResources(g *depgraph.Graph, destroy bool) {
 			continue
 		}
 
-		if rn.Resource.Diff == nil || !rn.Resource.Diff.Destroy {
+		if rn.Resource.Flags&FlagOrphan != 0 ||
+			rn.Resource.Diff == nil || !rn.Resource.Diff.Destroy {
 			result = append(result, n)
 		}
 	}