From 2b14a6b332ae399ddf895fc09697e350e6316fda Mon Sep 17 00:00:00 2001
From: Kristin Laemmert <mildwonkey@users.noreply.github.com>
Date: Thu, 1 Aug 2019 15:52:52 -0400
Subject: [PATCH] lang/funcs: lookup() can work with maps of lists, maps and
 objects (#22269)

* lang/funcs: lookup() can work with maps of lists, maps and objects

lookup() can already handle aribtrary objects of (whatever) and should
handle maps of (whatever) similarly.
---
 lang/funcs/collection.go      | 20 ++++-------
 lang/funcs/collection_test.go | 66 +++++++++++++++++++++++++++++++++++
 2 files changed, 73 insertions(+), 13 deletions(-)

diff --git a/lang/funcs/collection.go b/lang/funcs/collection.go
index bcccc1fd2..e6898457b 100644
--- a/lang/funcs/collection.go
+++ b/lang/funcs/collection.go
@@ -660,6 +660,12 @@ var LookupFunc = function.New(&function.Spec{
 			}
 			return cty.DynamicPseudoType, function.NewArgErrorf(0, "the given object has no attribute %q", key)
 		case ty.IsMapType():
+			if len(args) == 3 {
+				_, err = convert.Convert(args[2], ty.ElementType())
+				if err != nil {
+					return cty.NilType, function.NewArgErrorf(2, "the default value must have the same type as the map elements")
+				}
+			}
 			return ty.ElementType(), nil
 		default:
 			return cty.NilType, function.NewArgErrorf(0, "lookup() requires a map as the first argument")
@@ -686,19 +692,7 @@ var LookupFunc = function.New(&function.Spec{
 				return mapVar.GetAttr(lookupKey), nil
 			}
 		} else if mapVar.HasIndex(cty.StringVal(lookupKey)) == cty.True {
-			v := mapVar.Index(cty.StringVal(lookupKey))
-			if ty := v.Type(); !ty.Equals(cty.NilType) {
-				switch {
-				case ty.Equals(cty.String):
-					return cty.StringVal(v.AsString()), nil
-				case ty.Equals(cty.Number):
-					return cty.NumberVal(v.AsBigFloat()), nil
-				case ty.Equals(cty.Bool):
-					return cty.BoolVal(v.True()), nil
-				default:
-					return cty.NilVal, errors.New("lookup() can only be used with maps of primitive types")
-				}
-			}
+			return mapVar.Index(cty.StringVal(lookupKey)), nil
 		}
 
 		if defaultValueSet {
diff --git a/lang/funcs/collection_test.go b/lang/funcs/collection_test.go
index 71a805e7b..cec18c049 100644
--- a/lang/funcs/collection_test.go
+++ b/lang/funcs/collection_test.go
@@ -1469,6 +1469,26 @@ func TestLookup(t *testing.T) {
 			cty.StringVal("baz"),
 		}),
 	})
+	mapOfMaps := cty.MapVal(map[string]cty.Value{
+		"foo": cty.MapVal(map[string]cty.Value{
+			"a": cty.StringVal("bar"),
+		}),
+		"baz": cty.MapVal(map[string]cty.Value{
+			"b": cty.StringVal("bat"),
+		}),
+	})
+	mapOfTuples := cty.MapVal(map[string]cty.Value{
+		"foo": cty.TupleVal([]cty.Value{cty.StringVal("bar")}),
+		"baz": cty.TupleVal([]cty.Value{cty.StringVal("bat")}),
+	})
+	objectOfMaps := cty.ObjectVal(map[string]cty.Value{
+		"foo": cty.MapVal(map[string]cty.Value{
+			"a": cty.StringVal("bar"),
+		}),
+		"baz": cty.MapVal(map[string]cty.Value{
+			"b": cty.StringVal("bat"),
+		}),
+	})
 	mapWithUnknowns := cty.MapVal(map[string]cty.Value{
 		"foo": cty.StringVal("bar"),
 		"baz": cty.UnknownVal(cty.String),
@@ -1507,6 +1527,34 @@ func TestLookup(t *testing.T) {
 			cty.NumberIntVal(42),
 			false,
 		},
+		{
+			[]cty.Value{
+				mapOfMaps,
+				cty.StringVal("foo"),
+			},
+			cty.MapVal(map[string]cty.Value{
+				"a": cty.StringVal("bar"),
+			}),
+			false,
+		},
+		{
+			[]cty.Value{
+				objectOfMaps,
+				cty.StringVal("foo"),
+			},
+			cty.MapVal(map[string]cty.Value{
+				"a": cty.StringVal("bar"),
+			}),
+			false,
+		},
+		{
+			[]cty.Value{
+				mapOfTuples,
+				cty.StringVal("foo"),
+			},
+			cty.TupleVal([]cty.Value{cty.StringVal("bar")}),
+			false,
+		},
 		{ // Invalid key
 			[]cty.Value{
 				simpleMap,
@@ -1541,6 +1589,15 @@ func TestLookup(t *testing.T) {
 			cty.StringVal("bar"),
 			false,
 		},
+		{ // Supplied default with valid (int) key
+			[]cty.Value{
+				simpleMap,
+				cty.StringVal("foobar"),
+				cty.NumberIntVal(-1),
+			},
+			cty.StringVal("-1"),
+			false,
+		},
 		{ // Supplied default with valid key
 			[]cty.Value{
 				mapWithObjects,
@@ -1559,6 +1616,15 @@ func TestLookup(t *testing.T) {
 			cty.StringVal(""),
 			false,
 		},
+		{ // Supplied default with type mismatch: expects a map return
+			[]cty.Value{
+				mapOfMaps,
+				cty.StringVal("foo"),
+				cty.StringVal(""),
+			},
+			cty.NilVal,
+			true,
+		},
 		{ // Supplied non-empty default with invalid key
 			[]cty.Value{
 				simpleMap,