From ebc7d209a7acbd7bf21bace51a2c651be97f49c1 Mon Sep 17 00:00:00 2001
From: Mitchell Hashimoto <mitchell.hashimoto@gmail.com>
Date: Mon, 19 Sep 2016 09:28:24 -0700
Subject: [PATCH] terraform: new graph fixes ".0" and "" boundaries on counts

---
 terraform/eval_count_boundary.go      | 78 +++++++++++++++++++++++++++
 terraform/graph_builder_apply.go      |  7 +++
 terraform/node_count_boundary.go      | 14 +++++
 terraform/transform_count_boundary.go | 28 ++++++++++
 4 files changed, 127 insertions(+)
 create mode 100644 terraform/eval_count_boundary.go
 create mode 100644 terraform/node_count_boundary.go
 create mode 100644 terraform/transform_count_boundary.go

diff --git a/terraform/eval_count_boundary.go b/terraform/eval_count_boundary.go
new file mode 100644
index 000000000..91e2b904e
--- /dev/null
+++ b/terraform/eval_count_boundary.go
@@ -0,0 +1,78 @@
+package terraform
+
+import (
+	"log"
+)
+
+// EvalCountFixZeroOneBoundaryGlobal is an EvalNode that fixes up the state
+// when there is a resource count with zero/one boundary, i.e. fixing
+// a resource named "aws_instance.foo" to "aws_instance.foo.0" and vice-versa.
+//
+// This works on the global state.
+type EvalCountFixZeroOneBoundaryGlobal struct{}
+
+// TODO: test
+func (n *EvalCountFixZeroOneBoundaryGlobal) Eval(ctx EvalContext) (interface{}, error) {
+	// Get the state and lock it since we'll potentially modify it
+	state, lock := ctx.State()
+	lock.Lock()
+	defer lock.Unlock()
+
+	// Prune the state since we require a clean state to work
+	state.prune()
+
+	// Go through each modules since the boundaries are restricted to a
+	// module scope.
+	for _, m := range state.Modules {
+		if err := n.fixModule(m); err != nil {
+			return nil, err
+		}
+	}
+
+	return nil, nil
+}
+
+func (n *EvalCountFixZeroOneBoundaryGlobal) fixModule(m *ModuleState) error {
+	// Counts keeps track of keys and their counts
+	counts := make(map[string]int)
+	for k, _ := range m.Resources {
+		// Parse the key
+		key, err := ParseResourceStateKey(k)
+		if err != nil {
+			return err
+		}
+
+		// Set the index to -1 so that we can keep count
+		key.Index = -1
+
+		// Increment
+		counts[key.String()]++
+	}
+
+	// Go through the counts and do the fixup for each resource
+	for raw, count := range counts {
+		// Search and replace this resource
+		search := raw
+		replace := raw + ".0"
+		if count < 2 {
+			search, replace = replace, search
+		}
+		log.Printf("[TRACE] EvalCountFixZeroOneBoundaryGlobal: count %d, search %q, replace %q", count, search, replace)
+
+		// Look for the resource state. If we don't have one, then it is okay.
+		rs, ok := m.Resources[search]
+		if !ok {
+			continue
+		}
+
+		// If the replacement key exists, we just keep both
+		if _, ok := m.Resources[replace]; ok {
+			continue
+		}
+
+		m.Resources[replace] = rs
+		delete(m.Resources, search)
+	}
+
+	return nil
+}
diff --git a/terraform/graph_builder_apply.go b/terraform/graph_builder_apply.go
index 7a6e735ca..4889bda55 100644
--- a/terraform/graph_builder_apply.go
+++ b/terraform/graph_builder_apply.go
@@ -82,8 +82,15 @@ func (b *ApplyGraphBuilder) Steps() []GraphTransformer {
 		// Connect references so ordering is correct
 		&ReferenceTransformer{},
 
+		// Add the node to fix the state count boundaries
+		&CountBoundaryTransformer{},
+
 		// Single root
 		&RootTransformer{},
+
+		// Perform the transitive reduction to make our graph a bit
+		// more sane if possible (it usually is possible).
+		&TransitiveReductionTransformer{},
 	}
 
 	return steps
diff --git a/terraform/node_count_boundary.go b/terraform/node_count_boundary.go
new file mode 100644
index 000000000..bd32c79f3
--- /dev/null
+++ b/terraform/node_count_boundary.go
@@ -0,0 +1,14 @@
+package terraform
+
+// NodeCountBoundary fixes any "count boundarie" in the state: resources
+// that are named "foo.0" when they should be named "foo"
+type NodeCountBoundary struct{}
+
+func (n *NodeCountBoundary) Name() string {
+	return "meta.count-boundary (count boundary fixup)"
+}
+
+// GraphNodeEvalable
+func (n *NodeCountBoundary) EvalTree() EvalNode {
+	return &EvalCountFixZeroOneBoundaryGlobal{}
+}
diff --git a/terraform/transform_count_boundary.go b/terraform/transform_count_boundary.go
new file mode 100644
index 000000000..83415f352
--- /dev/null
+++ b/terraform/transform_count_boundary.go
@@ -0,0 +1,28 @@
+package terraform
+
+import (
+	"github.com/hashicorp/terraform/dag"
+)
+
+// CountBoundaryTransformer adds a node that depends on everything else
+// so that it runs last in order to clean up the state for nodes that
+// are on the "count boundary": "foo.0" when only one exists becomes "foo"
+type CountBoundaryTransformer struct{}
+
+func (t *CountBoundaryTransformer) Transform(g *Graph) error {
+	node := &NodeCountBoundary{}
+	g.Add(node)
+
+	// Depends on everything
+	for _, v := range g.Vertices() {
+		// Don't connect to ourselves
+		if v == node {
+			continue
+		}
+
+		// Connect!
+		g.Connect(dag.BasicEdge(node, v))
+	}
+
+	return nil
+}