From 66dd2d205d3e9882b40a008691ac0e1883c6fedd Mon Sep 17 00:00:00 2001
From: Anton Tayanovskyy <anton@pulumi.com>
Date: Wed, 12 Jun 2024 16:22:20 -0400
Subject: [PATCH] Support IndexStep.Apply for element values

---
 cty/path.go      | 13 +++++++++++++
 cty/path_test.go | 24 ++++++++++++++++++++++++
 cty/value_ops.go | 29 ++++++++++++++++++++++++-----
 3 files changed, 61 insertions(+), 5 deletions(-)

diff --git a/cty/path.go b/cty/path.go
index 4995a8c7..7600a71f 100644
--- a/cty/path.go
+++ b/cty/path.go
@@ -212,6 +212,19 @@ func (s IndexStep) Apply(val Value) (Value, error) {
 		return NilVal, errors.New("cannot index a null value")
 	}
 
+	// For sets, IndexStep.Key is the value itself.
+	if val.Type().IsSetType() {
+		findElementResult := val.findElement(s.Key)
+		has, _ := findElementResult.hasElement.Unmark()
+		if !has.IsKnown() {
+			return UnknownVal(val.Type().ElementType()), nil
+		}
+		if !has.True() {
+			return NilVal, errors.New("value does not have given element")
+		}
+		return findElementResult.matchingElement, nil
+	}
+
 	switch s.Key.Type() {
 	case Number:
 		if !(val.Type().IsListType() || val.Type().IsTupleType()) {
diff --git a/cty/path_test.go b/cty/path_test.go
index 93eaafc3..64ce7af6 100644
--- a/cty/path_test.go
+++ b/cty/path_test.go
@@ -160,6 +160,30 @@ func TestPathApply(t *testing.T) {
 			cty.StringVal("there").Mark(1),
 			``,
 		},
+		{
+			cty.SetVal([]cty.Value{
+				cty.ObjectVal(map[string]cty.Value{"x": cty.StringVal("X1")}),
+				cty.ObjectVal(map[string]cty.Value{"x": cty.StringVal("X2")}),
+			}),
+			cty.IndexPath(cty.ObjectVal(map[string]cty.Value{"x": cty.StringVal("X1")})).GetAttr("x"),
+			cty.StringVal("X1"),
+			``,
+		},
+		{
+			cty.SetVal([]cty.Value{
+				cty.ObjectVal(map[string]cty.Value{"x": cty.StringVal("X1")}),
+				cty.ObjectVal(map[string]cty.Value{"x": cty.StringVal("X2")}),
+			}),
+			cty.IndexPath(cty.ObjectVal(map[string]cty.Value{"x": cty.StringVal("X3")})).GetAttr("x"),
+			cty.NilVal,
+			`at step 0: value does not have given element`,
+		},
+		{
+			cty.UnknownVal(cty.Set(cty.Object(map[string]cty.Type{"x": cty.String}))),
+			cty.IndexPath(cty.ObjectVal(map[string]cty.Value{"x": cty.StringVal("X3")})).GetAttr("x"),
+			cty.UnknownVal(cty.String),
+			``,
+		},
 	}
 
 	for _, test := range tests {
diff --git a/cty/value_ops.go b/cty/value_ops.go
index c4584bd9..c33b4138 100644
--- a/cty/value_ops.go
+++ b/cty/value_ops.go
@@ -1021,10 +1021,24 @@ func (val Value) HasIndex(key Value) Value {
 //
 // This method will panic if the receiver is not a set, or if it is a null set.
 func (val Value) HasElement(elem Value) Value {
+	return val.findElement(elem).hasElement
+}
+
+type findElementResult struct {
+	hasElement      Value
+	matchingElement Value
+}
+
+// Like HasElement but also returns the matching set element found.
+func (val Value) findElement(elem Value) findElementResult {
 	if val.IsMarked() || elem.IsMarked() {
 		val, valMarks := val.Unmark()
 		elem, elemMarks := elem.Unmark()
-		return val.HasElement(elem).WithMarks(valMarks, elemMarks)
+		r := val.findElement(elem)
+		return findElementResult{
+			hasElement:      r.hasElement.WithMarks(valMarks, elemMarks),
+			matchingElement: r.matchingElement,
+		}
 	}
 
 	ty := val.Type()
@@ -1033,17 +1047,22 @@ func (val Value) HasElement(elem Value) Value {
 		panic("not a set type")
 	}
 	if !val.IsKnown() || !elem.IsKnown() {
-		return UnknownVal(Bool).RefineNotNull()
+		return findElementResult{hasElement: UnknownVal(Bool).RefineNotNull()}
 	}
 	if val.IsNull() {
 		panic("can't call HasElement on a null value")
 	}
 	if !ty.ElementType().Equals(elem.Type()) {
-		return False
+		return findElementResult{hasElement: False}
 	}
-
 	s := val.v.(set.Set[interface{}])
-	return BoolVal(s.Has(elem.v))
+	if !s.Has(elem.v) {
+		return findElementResult{hasElement: False}
+	}
+	return findElementResult{
+		hasElement:      True,
+		matchingElement: elem,
+	}
 }
 
 // Length returns the length of the receiver, which must be a collection type