From 93ef015336631f784b1f06c7a9d8372ffbd9f635 Mon Sep 17 00:00:00 2001
From: James Bardin <j.bardin@gmail.com>
Date: Wed, 1 May 2019 14:22:50 -0400
Subject: [PATCH] more precise type handling in flatten

FlattenFunc can return lists and tuples when individual elements are
unknown. Only return an unknown tuple if the number of elements cannot
be determined because it contains an unknown series.

Make sure flatten can handle non-series elements, which were previously
lost due to passing a slice value as the argument.
---
 lang/funcs/collection.go      | 40 +++++++++++++++-------
 lang/funcs/collection_test.go | 64 ++++++++++++++++++++++++++++++++---
 2 files changed, 87 insertions(+), 17 deletions(-)

diff --git a/lang/funcs/collection.go b/lang/funcs/collection.go
index d72d7261b..3ec505b00 100644
--- a/lang/funcs/collection.go
+++ b/lang/funcs/collection.go
@@ -441,8 +441,11 @@ var FlattenFunc = function.New(&function.Spec{
 			return cty.NilType, fmt.Errorf("can only flatten lists, sets and tuples")
 		}
 
-		outputList := make([]cty.Value, 0)
-		retVal := flattener(outputList, args[0])
+		retVal, known := flattener(args[0])
+		if !known {
+			return cty.DynamicPseudoType, nil
+		}
+
 		tys := make([]cty.Type, len(retVal))
 		for i, ty := range retVal {
 			tys[i] = ty.Type()
@@ -451,30 +454,41 @@ var FlattenFunc = function.New(&function.Spec{
 	},
 	Impl: func(args []cty.Value, retType cty.Type) (ret cty.Value, err error) {
 		inputList := args[0]
-		if !inputList.IsWhollyKnown() {
-			return cty.UnknownVal(retType), nil
-		}
 		if inputList.LengthInt() == 0 {
 			return cty.EmptyTupleVal, nil
 		}
-		outputList := make([]cty.Value, 0)
 
-		return cty.TupleVal(flattener(outputList, inputList)), nil
+		out, known := flattener(inputList)
+		if !known {
+			return cty.UnknownVal(retType), nil
+		}
+
+		return cty.TupleVal(out), nil
 	},
 })
 
-// Flatten until it's not a cty.List
-func flattener(finalList []cty.Value, flattenList cty.Value) []cty.Value {
+// Flatten until it's not a cty.List, and return whether the value is known.
+// We can flatten lists with unknown values, as long as they are not
+// lists themselves.
+func flattener(flattenList cty.Value) ([]cty.Value, bool) {
+	out := make([]cty.Value, 0)
 	for it := flattenList.ElementIterator(); it.Next(); {
 		_, val := it.Element()
-
 		if val.Type().IsListType() || val.Type().IsSetType() || val.Type().IsTupleType() {
-			finalList = flattener(finalList, val)
+			if !val.IsKnown() {
+				return out, false
+			}
+
+			res, known := flattener(val)
+			if !known {
+				return res, known
+			}
+			out = append(out, res...)
 		} else {
-			finalList = append(finalList, val)
+			out = append(out, val)
 		}
 	}
-	return finalList
+	return out, true
 }
 
 // KeysFunc contructs a function that takes a map and returns a sorted list of the map keys.
diff --git a/lang/funcs/collection_test.go b/lang/funcs/collection_test.go
index eadc3cddb..e6dbb04fa 100644
--- a/lang/funcs/collection_test.go
+++ b/lang/funcs/collection_test.go
@@ -1043,6 +1043,44 @@ func TestFlatten(t *testing.T) {
 			}),
 			false,
 		},
+		// handle single elements as arguments
+		{
+			cty.TupleVal([]cty.Value{
+				cty.ListVal([]cty.Value{
+					cty.StringVal("a"),
+					cty.StringVal("b"),
+				}),
+				cty.StringVal("c"),
+			}),
+			cty.TupleVal([]cty.Value{
+				cty.StringVal("a"),
+				cty.StringVal("b"),
+				cty.StringVal("c"),
+			}), false,
+		},
+		// handle single elements and mixed primitive types as arguments
+		{
+			cty.TupleVal([]cty.Value{
+				cty.ListVal([]cty.Value{
+					cty.StringVal("a"),
+					cty.StringVal("b"),
+				}),
+				cty.StringVal("c"),
+				cty.TupleVal([]cty.Value{
+					cty.StringVal("x"),
+					cty.NumberIntVal(1),
+				}),
+			}),
+			cty.TupleVal([]cty.Value{
+				cty.StringVal("a"),
+				cty.StringVal("b"),
+				cty.StringVal("c"),
+				cty.StringVal("x"),
+				cty.NumberIntVal(1),
+			}),
+			false,
+		},
+		// Primitive unknowns should still be flattened to a tuple
 		{
 			cty.ListVal([]cty.Value{
 				cty.ListVal([]cty.Value{
@@ -1054,8 +1092,26 @@ func TestFlatten(t *testing.T) {
 					cty.StringVal("d"),
 				}),
 			}),
-			cty.DynamicVal,
-			false,
+			cty.TupleVal([]cty.Value{
+				cty.StringVal("a"),
+				cty.StringVal("b"),
+				cty.UnknownVal(cty.String),
+				cty.StringVal("d"),
+			}), false,
+		},
+		// An unknown series should return an unknown dynamic value
+		{
+			cty.TupleVal([]cty.Value{
+				cty.ListVal([]cty.Value{
+					cty.StringVal("a"),
+					cty.StringVal("b"),
+				}),
+				cty.TupleVal([]cty.Value{
+					cty.UnknownVal(cty.List(cty.String)),
+					cty.StringVal("d"),
+				}),
+			}),
+			cty.UnknownVal(cty.DynamicPseudoType), false,
 		},
 		{
 			cty.ListValEmpty(cty.String),
@@ -1102,8 +1158,8 @@ func TestFlatten(t *testing.T) {
 		},
 	}
 
-	for _, test := range tests {
-		t.Run(fmt.Sprintf("flatten(%#v)", test.List), func(t *testing.T) {
+	for i, test := range tests {
+		t.Run(fmt.Sprintf("%d-flatten(%#v)", i, test.List), func(t *testing.T) {
 			got, err := Flatten(test.List)
 
 			if test.Err {