function/stdlib: LookupFunc precise handling of marks

Several changes to Lookup to handle marks in a way that I think makes
more sense:

- If the entire collection is marked, preserve the marks on any result
  (whether successful or fallback)
- If a returned value from the collection is marked, preserve the marks
  from only that value, combined with any overall collection marks
- Retain marks on the fallback value when it is returned, combined with
  any overall collection marks
- Disregard marks on the key value, as map keys and object attributes
  cannot be marked
- Retain collection marks when returning an unknown value for a not
  wholly-known collection

cty/function/stdlib/collection.go

@@ -633,16 +633,19 @@ var KeysFunc = function.New(&function.Spec{
 var LookupFunc = function.New(&function.Spec{
 	Params: []function.Parameter{
 		{
-			Name: "inputMap",
-			Type: cty.DynamicPseudoType,
+			Name:        "inputMap",
+			Type:        cty.DynamicPseudoType,
+			AllowMarked: true,
 		},
 		{
-			Name: "key",
-			Type: cty.String,
+			Name:        "key",
+			Type:        cty.String,
+			AllowMarked: true,
 		},
 		{
-			Name: "default",
-			Type: cty.DynamicPseudoType,
+			Name:        "default",
+			Type:        cty.DynamicPseudoType,
+			AllowMarked: true,
 		},
 	},
 	Type: func(args []cty.Value) (ret cty.Type, err error) {
@@ -654,7 +657,8 @@ var LookupFunc = function.New(&function.Spec{
 				return cty.DynamicPseudoType, nil
 			}
 
-			key := args[1].AsString()
+			keyVal, _ := args[1].Unmark()
+			key := keyVal.AsString()
 			if ty.HasAttribute(key) {
 				return args[0].GetAttr(key).Type(), nil
 			} else if len(args) == 3 {
@@ -676,28 +680,39 @@ var LookupFunc = function.New(&function.Spec{
 		}
 	},
 	Impl: func(args []cty.Value, retType cty.Type) (ret cty.Value, err error) {
+		// leave default value marked
 		defaultVal := args[2]
 
-		mapVar := args[0]
-		lookupKey := args[1].AsString()
+		var markses []cty.ValueMarks
+
+		// unmark collection, retain marks to reapply later
+		mapVar, mapMarks := args[0].Unmark()
+		markses = append(markses, mapMarks)
+
+		// include marks on the key in the result
+		keyVal, keyMarks := args[1].Unmark()
+		if len(keyMarks) > 0 {
+			markses = append(markses, keyMarks)
+		}
+		lookupKey := keyVal.AsString()
 
 		if !mapVar.IsWhollyKnown() {
-			return cty.UnknownVal(retType), nil
+			return cty.UnknownVal(retType).WithMarks(markses...), nil
 		}
 
 		if mapVar.Type().IsObjectType() {
 			if mapVar.Type().HasAttribute(lookupKey) {
-				return mapVar.GetAttr(lookupKey), nil
+				return mapVar.GetAttr(lookupKey).WithMarks(markses...), nil
 			}
 		} else if mapVar.HasIndex(cty.StringVal(lookupKey)) == cty.True {
-			return mapVar.Index(cty.StringVal(lookupKey)), nil
+			return mapVar.Index(cty.StringVal(lookupKey)).WithMarks(markses...), nil
 		}
 
 		defaultVal, err = convert.Convert(defaultVal, retType)
 		if err != nil {
 			return cty.NilVal, err
 		}
-		return defaultVal, nil
+		return defaultVal.WithMarks(markses...), nil
 	},
 })
 

cty/function/stdlib/collection_test.go

@@ -1028,6 +1028,67 @@ func TestLookup(t *testing.T) {
 			cty.StringVal("nope"),
 			cty.StringVal("bar"),
 		},
+		{ // successful marked collection lookup returns marked value
+			cty.MapVal(map[string]cty.Value{
+				"boop": cty.StringVal("beep"),
+			}).Mark("a"),
+			cty.StringVal("boop"),
+			cty.StringVal("nope"),
+			cty.StringVal("beep").Mark("a"),
+		},
+		{ // apply collection marks to unknown return vaue
+			cty.MapVal(map[string]cty.Value{
+				"boop": cty.StringVal("beep"),
+				"frob": cty.UnknownVal(cty.String),
+			}).Mark("a"),
+			cty.StringVal("boop"),
+			cty.StringVal("nope"),
+			cty.UnknownVal(cty.String).Mark("a"),
+		},
+		{ // propagate collection marks to default when returning
+			cty.MapVal(map[string]cty.Value{
+				"boop": cty.StringVal("beep"),
+			}).Mark("a"),
+			cty.StringVal("frob"),
+			cty.StringVal("nope").Mark("b"),
+			cty.StringVal("nope").WithMarks(cty.NewValueMarks("a", "b")),
+		},
+		{ // on unmarked collection, return only marks from found value
+			cty.MapVal(map[string]cty.Value{
+				"boop": cty.StringVal("beep").Mark("a"),
+				"frob": cty.StringVal("honk").Mark("b"),
+			}),
+			cty.StringVal("frob"),
+			cty.StringVal("nope").Mark("c"),
+			cty.StringVal("honk").Mark("b"),
+		},
+		{ // on unmarked collection, return default exactly on missing
+			cty.MapVal(map[string]cty.Value{
+				"boop": cty.StringVal("beep").Mark("a"),
+				"frob": cty.StringVal("honk").Mark("b"),
+			}),
+			cty.StringVal("squish"),
+			cty.StringVal("nope").Mark("c"),
+			cty.StringVal("nope").Mark("c"),
+		},
+		{ // retain marks on default if converted
+			cty.MapVal(map[string]cty.Value{
+				"boop": cty.StringVal("beep").Mark("a"),
+				"frob": cty.StringVal("honk").Mark("b"),
+			}),
+			cty.StringVal("squish"),
+			cty.NumberIntVal(5).Mark("c"),
+			cty.StringVal("5").Mark("c"),
+		},
+		{ // propagate marks from key
+			cty.MapVal(map[string]cty.Value{
+				"boop": cty.StringVal("beep"),
+				"frob": cty.StringVal("honk"),
+			}),
+			cty.StringVal("boop").Mark("a"),
+			cty.StringVal("nope"),
+			cty.StringVal("beep").Mark("a"),
+		},
 	}
 
 	for _, test := range tests {