functions: FlattenFunc

lang/funcs/collection.go

@@ -334,6 +334,43 @@ var ChunklistFunc = function.New(&function.Spec{
 	},
 })
 
+// FlattenFunc contructs a function that takes a list and replaces any elements
+// that are lists with a flattened sequence of the list contents.
+var FlattenFunc = function.New(&function.Spec{
+	Params: []function.Parameter{
+		{
+			Name: "list",
+			Type: cty.List(cty.DynamicPseudoType),
+		},
+	},
+	Type: function.StaticReturnType(cty.List(cty.DynamicPseudoType)),
+	Impl: func(args []cty.Value, retType cty.Type) (ret cty.Value, err error) {
+		inputList := args[0]
+
+		if inputList.LengthInt() == 0 {
+			return cty.ListValEmpty(cty.DynamicPseudoType), nil
+		}
+		outputList := make([]cty.Value, 0)
+
+		return cty.ListVal(flattener(outputList, inputList)), nil
+	},
+})
+
+// Flatten until it's not a cty.List
+func flattener(finalList []cty.Value, flattenList cty.Value) []cty.Value {
+
+	for it := flattenList.ElementIterator(); it.Next(); {
+		_, val := it.Element()
+
+		if val.Type().IsListType() {
+			finalList = flattener(finalList, val)
+		} else {
+			finalList = append(finalList, val)
+		}
+	}
+	return finalList
+}
+
 // MatchkeysFunc contructs a function that constructs a new list by taking a
 // subset of elements from one list whose indexes match the corresponding
 // indexes of values in another list.
@@ -353,16 +390,18 @@ var MatchkeysFunc = function.New(&function.Spec{
 		},
 	},
 	Type: func(args []cty.Value) (cty.Type, error) {
-		if args[1].Type() != args[2].Type() {
+		if !args[1].Type().Equals(args[2].Type()) {
 			return cty.NilType, fmt.Errorf("lists must be of the same type")
 		}
 
-		if args[0].LengthInt() != args[1].LengthInt() {
-			return cty.NilType, fmt.Errorf("length of keys and values should be equal")
-		}
 		return args[0].Type(), nil
 	},
 	Impl: func(args []cty.Value, retType cty.Type) (ret cty.Value, err error) {
+
+		if args[0].LengthInt() != args[1].LengthInt() {
+			return cty.ListValEmpty(retType.ElementType()), fmt.Errorf("length of keys and values should be equal")
+		}
+
 		output := make([]cty.Value, 0)
 
 		values := args[0]
@@ -371,7 +410,7 @@ var MatchkeysFunc = function.New(&function.Spec{
 
 		// if searchset is empty, return an empty list.
 		if searchset.LengthInt() == 0 {
-			return cty.ListValEmpty(*retType.ListElementType()), nil
+			return cty.ListValEmpty(retType.ElementType()), nil
 		}
 
 		i := 0
@@ -383,6 +422,9 @@ var MatchkeysFunc = function.New(&function.Spec{
 				if err != nil {
 					return cty.NilVal, err
 				}
+				if !eq.IsKnown() {
+					return cty.ListValEmpty(retType.ElementType()), nil
+				}
 				if eq.True() {
 					v := values.Index(cty.NumberIntVal(int64(i)))
 					output = append(output, v)
@@ -394,7 +436,7 @@ var MatchkeysFunc = function.New(&function.Spec{
 
 		// if we haven't matched any key, then output is an empty list.
 		if len(output) == 0 {
-			return cty.ListValEmpty(*retType.ListElementType()), nil
+			return cty.ListValEmpty(retType.ElementType()), nil
 		}
 		return cty.ListVal(output), nil
 	},
@@ -459,6 +501,12 @@ func Chunklist(list, size cty.Value) (cty.Value, error) {
 	return ChunklistFunc.Call([]cty.Value{list, size})
 }
 
+// Flatten takes a list and replaces any elements that are lists with a flattened
+// sequence of the list contents.
+func Flatten(list cty.Value) (cty.Value, error) {
+	return FlattenFunc.Call([]cty.Value{list})
+}
+
 // Matchkeys constructs a new list by taking a subset of elements from one list
 // whose indexes match the corresponding indexes of values in another list.
 func Matchkeys(values, keys, searchset cty.Value) (cty.Value, error) {

lang/funcs/collection_test.go

@@ -787,6 +787,58 @@ func TestChunklist(t *testing.T) {
 	}
 }
 
+func TestFlatten(t *testing.T) {
+	tests := []struct {
+		List cty.Value
+		Want cty.Value
+		Err  bool
+	}{
+		{
+			cty.ListVal([]cty.Value{
+				cty.ListVal([]cty.Value{
+					cty.StringVal("a"),
+					cty.StringVal("b"),
+				}),
+				cty.ListVal([]cty.Value{
+					cty.StringVal("c"),
+					cty.StringVal("d"),
+				}),
+			}),
+			cty.ListVal([]cty.Value{
+				cty.StringVal("a"),
+				cty.StringVal("b"),
+				cty.StringVal("c"),
+				cty.StringVal("d"),
+			}),
+			false,
+		},
+		{
+			cty.ListValEmpty(cty.String),
+			cty.ListValEmpty(cty.DynamicPseudoType),
+			false,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(fmt.Sprintf("flatten(%#v)", test.List), func(t *testing.T) {
+			got, err := Flatten(test.List)
+
+			if test.Err {
+				if err == nil {
+					t.Fatal("succeeded; want error")
+				}
+				return
+			} else if err != nil {
+				t.Fatalf("unexpected error: %s", err)
+			}
+
+			if !got.RawEquals(test.Want) {
+				t.Errorf("wrong result\ngot:  %#v\nwant: %#v", got, test.Want)
+			}
+		})
+	}
+}
+
 func TestMatchkeys(t *testing.T) {
 	tests := []struct {
 		Keys      cty.Value

lang/functions.go

@@ -54,8 +54,7 @@ func (s *Scope) Functions() map[string]function.Function {
 			"chunklist":    funcs.ChunklistFunc,
 			"file":         funcs.MakeFileFunc(s.BaseDir, false),
 			"filebase64":   funcs.MakeFileFunc(s.BaseDir, true),
-			"matchkeys":    funcs.MatchkeysFunc,
-			"flatten":      unimplFunc, // TODO
+			"flatten":      funcs.FlattenFunc,
 			"floor":        funcs.FloorFunc,
 			"format":       stdlib.FormatFunc,
 			"formatlist":   stdlib.FormatListFunc,
@@ -71,6 +70,7 @@ func (s *Scope) Functions() map[string]function.Function {
 			"lookup":       unimplFunc, // TODO
 			"lower":        stdlib.LowerFunc,
 			"map":          unimplFunc, // TODO
+			"matchkeys":    funcs.MatchkeysFunc,
 			"max":          stdlib.MaxFunc,
 			"md5":          funcs.Md5Func,
 			"merge":        unimplFunc, // TODO