From e0c6b3fcda9cc1766951bd935442215c91b098a2 Mon Sep 17 00:00:00 2001
From: Alisdair McDiarmid <alisdair@users.noreply.github.com>
Date: Fri, 7 May 2021 11:57:37 -0400
Subject: [PATCH] functions: Improve marks support for length

Similar to cty's implementation, we only need to preserve marks from the
value itself, not any nested values it may contain. This means that
taking the length of an umarked list with marked elements results in an
unmarked number.
---
 lang/funcs/collection.go      |  8 +++---
 lang/funcs/collection_test.go | 48 +++++++++++++++++++++++++++++++++++
 2 files changed, 53 insertions(+), 3 deletions(-)

diff --git a/lang/funcs/collection.go b/lang/funcs/collection.go
index ce6c8dc8f..0536dcfdf 100644
--- a/lang/funcs/collection.go
+++ b/lang/funcs/collection.go
@@ -20,6 +20,7 @@ var LengthFunc = function.New(&function.Spec{
 			Type:             cty.DynamicPseudoType,
 			AllowDynamicType: true,
 			AllowUnknown:     true,
+			AllowMarked:      true,
 		},
 	},
 	Type: func(args []cty.Value) (cty.Type, error) {
@@ -34,15 +35,16 @@ var LengthFunc = function.New(&function.Spec{
 	Impl: func(args []cty.Value, retType cty.Type) (cty.Value, error) {
 		coll := args[0]
 		collTy := args[0].Type()
+		marks := coll.Marks()
 		switch {
 		case collTy == cty.DynamicPseudoType:
-			return cty.UnknownVal(cty.Number), nil
+			return cty.UnknownVal(cty.Number).WithMarks(marks), nil
 		case collTy.IsTupleType():
 			l := len(collTy.TupleElementTypes())
-			return cty.NumberIntVal(int64(l)), nil
+			return cty.NumberIntVal(int64(l)).WithMarks(marks), nil
 		case collTy.IsObjectType():
 			l := len(collTy.AttributeTypes())
-			return cty.NumberIntVal(int64(l)), nil
+			return cty.NumberIntVal(int64(l)).WithMarks(marks), nil
 		case collTy == cty.String:
 			// We'll delegate to the cty stdlib strlen function here, because
 			// it deals with all of the complexities of tokenizing unicode
diff --git a/lang/funcs/collection_test.go b/lang/funcs/collection_test.go
index 9cc428a91..974a8437d 100644
--- a/lang/funcs/collection_test.go
+++ b/lang/funcs/collection_test.go
@@ -122,6 +122,54 @@ func TestLength(t *testing.T) {
 			cty.DynamicVal,
 			cty.UnknownVal(cty.Number),
 		},
+		{ // Marked collections return a marked length
+			cty.ListVal([]cty.Value{
+				cty.StringVal("hello"),
+				cty.StringVal("world"),
+			}).Mark("secret"),
+			cty.NumberIntVal(2).Mark("secret"),
+		},
+		{ // Marks on values in unmarked collections do not propagate
+			cty.ListVal([]cty.Value{
+				cty.StringVal("hello").Mark("a"),
+				cty.StringVal("world").Mark("b"),
+			}),
+			cty.NumberIntVal(2),
+		},
+		{ // Marked strings return a marked length
+			cty.StringVal("hello world").Mark("secret"),
+			cty.NumberIntVal(11).Mark("secret"),
+		},
+		{ // Marked tuples return a marked length
+			cty.TupleVal([]cty.Value{
+				cty.StringVal("hello"),
+				cty.StringVal("world"),
+			}).Mark("secret"),
+			cty.NumberIntVal(2).Mark("secret"),
+		},
+		{ // Marks on values in unmarked tuples do not propagate
+			cty.TupleVal([]cty.Value{
+				cty.StringVal("hello").Mark("a"),
+				cty.StringVal("world").Mark("b"),
+			}),
+			cty.NumberIntVal(2),
+		},
+		{ // Marked objects return a marked length
+			cty.ObjectVal(map[string]cty.Value{
+				"a": cty.StringVal("hello"),
+				"b": cty.StringVal("world"),
+				"c": cty.StringVal("nice to meet you"),
+			}).Mark("secret"),
+			cty.NumberIntVal(3).Mark("secret"),
+		},
+		{ // Marks on object attribute values do not propagate
+			cty.ObjectVal(map[string]cty.Value{
+				"a": cty.StringVal("hello").Mark("a"),
+				"b": cty.StringVal("world").Mark("b"),
+				"c": cty.StringVal("nice to meet you").Mark("c"),
+			}),
+			cty.NumberIntVal(3),
+		},
 	}
 
 	for _, test := range tests {