function/stdlib/keys: accept marked arguments

cty/function/stdlib/collection.go

@@ -539,6 +539,7 @@ var KeysFunc = function.New(&function.Spec{
 			Name:         "inputMap",
 			Type:         cty.DynamicPseudoType,
 			AllowUnknown: true,
+			AllowMarked:  true,
 		},
 	},
 	Type: func(args []cty.Value) (cty.Type, error) {
@@ -563,7 +564,11 @@ var KeysFunc = function.New(&function.Spec{
 		}
 	},
 	Impl: func(args []cty.Value, retType cty.Type) (cty.Value, error) {
-		m := args[0]
+		// We must unmark the value before we can use ElementIterator on it, and
+		// then re-apply the same marks (possibly none) when we return. Since we
+		// don't mark map keys, we can throw away any nested marks, which would
+		// only apply to values.
+		m, marks := args[0].Unmark()
 		var keys []cty.Value
 
 		switch {
@@ -576,28 +581,28 @@ var KeysFunc = function.New(&function.Spec{
 			}
 			sort.Strings(names) // same ordering guaranteed by cty's ElementIterator
 			if len(names) == 0 {
-				return cty.EmptyTupleVal, nil
+				return cty.EmptyTupleVal.WithMarks(marks), nil
 			}
 			keys = make([]cty.Value, len(names))
 			for i, name := range names {
 				keys[i] = cty.StringVal(name)
 			}
-			return cty.TupleVal(keys), nil
+			return cty.TupleVal(keys).WithMarks(marks), nil
 		default:
 			if !m.IsKnown() {
-				return cty.UnknownVal(retType), nil
+				return cty.UnknownVal(retType).WithMarks(marks), nil
 			}
 
 			// cty guarantees that ElementIterator will iterate in lexicographical
 			// order by key.
-			for it := args[0].ElementIterator(); it.Next(); {
+			for it := m.ElementIterator(); it.Next(); {
 				k, _ := it.Element()
 				keys = append(keys, k)
 			}
 			if len(keys) == 0 {
-				return cty.ListValEmpty(cty.String), nil
+				return cty.ListValEmpty(cty.String).WithMarks(marks), nil
 			}
-			return cty.ListVal(keys), nil
+			return cty.ListVal(keys).WithMarks(marks), nil
 		}
 	},
 })

cty/function/stdlib/collection_test.go

@@ -1327,3 +1327,113 @@ func TestZipMap(t *testing.T) {
 		})
 	}
 }
+
+func TestKeys(t *testing.T) {
+	tests := []struct {
+		Collection cty.Value
+		Want       cty.Value
+		Err        string
+	}{
+		{
+			cty.MapValEmpty(cty.String),
+			cty.ListValEmpty(cty.String),
+			``,
+		},
+		{
+			cty.MapValEmpty(cty.String).Mark("a"),
+			cty.ListValEmpty(cty.String).Mark("a"),
+			``,
+		},
+		{
+			cty.NullVal(cty.Map(cty.String)),
+			cty.NilVal,
+			`argument must not be null`,
+		},
+		{
+			cty.MapVal(map[string]cty.Value{"hello": cty.StringVal("world")}),
+			cty.ListVal([]cty.Value{cty.StringVal("hello")}),
+			``,
+		},
+		{ // The map itself is not marked, just an inner element.
+			cty.MapVal(map[string]cty.Value{"hello": cty.StringVal("world").Mark("a")}),
+			cty.ListVal([]cty.Value{cty.StringVal("hello")}),
+			``,
+		},
+		{ // The entire map is marked, so the resulting list is also marked.
+			cty.MapVal(map[string]cty.Value{"hello": cty.StringVal("world")}).Mark("a"),
+			cty.ListVal([]cty.Value{cty.StringVal("hello")}).Mark("a"),
+			``,
+		},
+		{ // Marked both inside and outside.
+			cty.MapVal(map[string]cty.Value{"hello": cty.StringVal("world").Mark("a")}).Mark("a"),
+			cty.ListVal([]cty.Value{cty.StringVal("hello")}).Mark("a"),
+			``,
+		},
+		{
+			cty.ObjectVal(map[string]cty.Value{"hello": cty.StringVal("world")}),
+			cty.TupleVal([]cty.Value{cty.StringVal("hello")}),
+			``,
+		},
+		{
+			cty.EmptyObjectVal,
+			cty.EmptyTupleVal,
+			``,
+		},
+		{
+			cty.EmptyObjectVal.Mark("a"),
+			cty.EmptyTupleVal.Mark("a"),
+			``,
+		},
+		{
+			cty.NullVal(cty.EmptyObject),
+			cty.NilVal,
+			`argument must not be null`,
+		},
+		{
+			cty.UnknownVal(cty.EmptyObject),
+			cty.EmptyTupleVal,
+			``,
+		},
+		{
+			cty.UnknownVal(cty.Object(map[string]cty.Type{"a": cty.String})),
+			cty.TupleVal([]cty.Value{cty.StringVal("a")}),
+			``,
+		},
+		{ // The object itself is not marked, just an inner attribute value.
+			cty.ObjectVal(map[string]cty.Value{"hello": cty.StringVal("world").Mark("a")}),
+			cty.TupleVal([]cty.Value{cty.StringVal("hello")}),
+			``,
+		},
+		{ // The entire object is marked, so the resulting tuple is also marked.
+			cty.ObjectVal(map[string]cty.Value{"hello": cty.StringVal("world")}).Mark("a"),
+			cty.TupleVal([]cty.Value{cty.StringVal("hello")}).Mark("a"),
+			``,
+		},
+		{ // Marked both inside and outside.
+			cty.ObjectVal(map[string]cty.Value{"hello": cty.StringVal("world").Mark("a")}).Mark("a"),
+			cty.TupleVal([]cty.Value{cty.StringVal("hello")}).Mark("a"),
+			``,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(fmt.Sprintf("Keys(%#v)", test.Collection), func(t *testing.T) {
+			got, err := Keys(test.Collection)
+			if test.Err != "" {
+				if err == nil {
+					t.Fatal("succeeded; want error")
+				}
+				if got, want := err.Error(), test.Err; got != want {
+					t.Fatalf("wrong error\ngot:  %s\nwant: %s", got, want)
+				}
+				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)
+			}
+		})
+	}
+}