From db968733da419cf905116d6a16c55dde22d5b794 Mon Sep 17 00:00:00 2001
From: James Bardin <j.bardin@gmail.com>
Date: Fri, 16 Nov 2018 15:26:16 -0500
Subject: [PATCH] re-count the flatmapped containers

When applying a legacy diff, recount the flatmapped containers. We can't
trust helper/schema to return the correct value, if it even exists.
---
 terraform/diff.go      | 47 ++++++++++++++++++++++++++++++++----------
 terraform/diff_test.go | 37 +++++++++++++++++++++++++++++++++
 2 files changed, 73 insertions(+), 11 deletions(-)

diff --git a/terraform/diff.go b/terraform/diff.go
index a28d7e482..ae2618316 100644
--- a/terraform/diff.go
+++ b/terraform/diff.go
@@ -7,6 +7,7 @@ import (
 	"reflect"
 	"regexp"
 	"sort"
+	"strconv"
 	"strings"
 	"sync"
 
@@ -617,17 +618,9 @@ func (d *InstanceDiff) applyCollectionDiff(attrName string, oldAttrs map[string]
 		}
 	}
 
-	// Verify we have the index count.
-	// If it wasn't added from a diff, check it from the previous value.
-	// Make sure we keep the count if it existed before, so we can tell if it
-	// existed, or was null.
-	if !setIndex {
-		old := oldAttrs[idx]
-		if old != "" {
-			result[idx] = old
-		}
-	}
-
+	// Don't trust helper/schema to return a valid count, or even have one at
+	// all.
+	result[idx] = countFlatmapContainerValues(idx, result)
 	return result, nil
 }
 
@@ -686,9 +679,41 @@ func (d *InstanceDiff) applySetDiff(attrName string, oldAttrs map[string]string,
 		}
 	}
 
+	result[idx] = countFlatmapContainerValues(idx, result)
+
 	return result, nil
 }
 
+// countFlatmapContainerValues returns the number of values in the flatmapped container
+// (set, map, list) indexed by key. The key argument is expected to include the
+// trailing ".#", or ".%".
+func countFlatmapContainerValues(key string, attrs map[string]string) string {
+	if len(key) < 3 || !(strings.HasSuffix(key, ".#") || strings.HasSuffix(key, ".%")) {
+		panic(fmt.Sprintf("invalid index value %q", key))
+	}
+
+	prefix := key[:len(key)-1]
+	items := map[string]int{}
+
+	for k := range attrs {
+		if k == key {
+			continue
+		}
+		if !strings.HasPrefix(k, prefix) {
+			continue
+		}
+
+		suffix := k[len(prefix):]
+		dot := strings.Index(suffix, ".")
+		if dot > 0 {
+			suffix = suffix[:dot]
+		}
+
+		items[suffix]++
+	}
+	return strconv.Itoa(len(items))
+}
+
 // ResourceAttrDiff is the diff of a single attribute of a resource.
 type ResourceAttrDiff struct {
 	Old         string      // Old Value
diff --git a/terraform/diff_test.go b/terraform/diff_test.go
index b2ace77f4..e7ee0d818 100644
--- a/terraform/diff_test.go
+++ b/terraform/diff_test.go
@@ -3,6 +3,7 @@ package terraform
 import (
 	"fmt"
 	"reflect"
+	"strconv"
 	"strings"
 	"testing"
 
@@ -1213,3 +1214,39 @@ CREATE: nodeA
   longfoo:   "foo" => "bar" (forces new resource)
   secretfoo: "<sensitive>" => "<sensitive>" (attribute changed)
 `
+
+func TestCountFlatmapContainerValues(t *testing.T) {
+	for i, tc := range []struct {
+		attrs map[string]string
+		key   string
+		count string
+	}{
+		{
+			attrs: map[string]string{"set.2.list.#": "9999", "set.2.list.0": "x", "set.2.list.0.z": "y", "set.2.attr": "bar", "set.#": "9999"},
+			key:   "set.2.list.#",
+			count: "1",
+		},
+		{
+			attrs: map[string]string{"set.2.list.#": "9999", "set.2.list.0": "x", "set.2.list.0.z": "y", "set.2.attr": "bar", "set.#": "9999"},
+			key:   "set.#",
+			count: "1",
+		},
+		{
+			attrs: map[string]string{"set.2.list.0": "x", "set.2.list.0.z": "y", "set.2.attr": "bar", "set.#": "9999"},
+			key:   "set.#",
+			count: "1",
+		},
+		{
+			attrs: map[string]string{"map.#": "3", "map.a": "b", "map.a.#": "0", "map.b": "4"},
+			key:   "map.#",
+			count: "2",
+		},
+	} {
+		t.Run(strconv.Itoa(i), func(t *testing.T) {
+			count := countFlatmapContainerValues(tc.key, tc.attrs)
+			if count != tc.count {
+				t.Fatalf("expected %q, got %q", tc.count, count)
+			}
+		})
+	}
+}