dynblock: Preserve marks from for_each expression into result

Previously if the for_each expression was marked then expansion would
fail because marked expressions are never directly iterable.

Now instead we'll allow marked for_each and preserve the marks into the
values produced by the resulting block as much as we can. This runs into
the classic problem that HCL blocks are not values themselves and so
cannot carry marks directly, but we can at least make sure that the values
of any leaf arguments end up marked.

5 files changed, 157 insertions, 20 deletions

diff --git a/ext/dynblock/expand_body.go b/ext/dynblock/expand_body.go
index 2734e93..e630f18 100644
--- a/ext/dynblock/expand_body.go
+++ b/ext/dynblock/expand_body.go
@@ -16,6 +16,7 @@ type expandBody struct {
 	original   hcl.Body
 	forEachCtx *hcl.EvalContext
 	iteration  *iteration // non-nil if we're nested inside another "dynamic" block
+	valueMarks cty.ValueMarks
 
 	checkForEach []func(cty.Value, hcl.Expression, *hcl.EvalContext) hcl.Diagnostics
 
@@ -125,7 +126,7 @@ func (b *expandBody) extendSchema(schema *hcl.BodySchema) *hcl.BodySchema {
 }
 
 func (b *expandBody) prepareAttributes(rawAttrs hcl.Attributes) hcl.Attributes {
-	if len(b.hiddenAttrs) == 0 && b.iteration == nil {
+	if len(b.hiddenAttrs) == 0 && b.iteration == nil && len(b.valueMarks) == 0 {
 		// Easy path: just pass through the attrs from the original body verbatim
 		return rawAttrs
 	}
@@ -142,13 +143,24 @@ func (b *expandBody) prepareAttributes(rawAttrs hcl.Attributes) hcl.Attributes {
 		if b.iteration != nil {
 			attr := *rawAttr // shallow copy so we can mutate it
 			attr.Expr = exprWrap{
-				Expression: attr.Expr,
-				i:          b.iteration,
+				Expression:  attr.Expr,
+				i:           b.iteration,
+				resultMarks: b.valueMarks,
 			}
 			attrs[name] = &attr
 		} else {
-			// If we have no active iteration then no wrapping is required.
-			attrs[name] = rawAttr
+			// If we have no active iteration then no wrapping is required
+			// unless we have marks to apply.
+			if len(b.valueMarks) != 0 {
+				attr := *rawAttr // shallow copy so we can mutate it
+				attr.Expr = exprWrap{
+					Expression:  attr.Expr,
+					resultMarks: b.valueMarks,
+				}
+				attrs[name] = &attr
+			} else {
+				attrs[name] = rawAttr
+			}
 		}
 	}
 	return attrs
@@ -192,8 +204,9 @@ func (b *expandBody) expandBlocks(schema *hcl.BodySchema, rawBlocks hcl.Blocks,
 				continue
 			}
 
-			if spec.forEachVal.IsKnown() {
-				for it := spec.forEachVal.ElementIterator(); it.Next(); {
+			forEachVal, marks := spec.forEachVal.Unmark()
+			if forEachVal.IsKnown() {
+				for it := forEachVal.ElementIterator(); it.Next(); {
 					key, value := it.Element()
 					i := b.iteration.MakeChild(spec.iteratorName, key, value)
 
@@ -202,7 +215,7 @@ func (b *expandBody) expandBlocks(schema *hcl.BodySchema, rawBlocks hcl.Blocks,
 					if block != nil {
 						// Attach our new iteration context so that attributes
 						// and other nested blocks can refer to our iterator.
-						block.Body = b.expandChild(block.Body, i)
+						block.Body = b.expandChild(block.Body, i, marks)
 						blocks = append(blocks, block)
 					}
 				}
@@ -214,7 +227,10 @@ func (b *expandBody) expandBlocks(schema *hcl.BodySchema, rawBlocks hcl.Blocks,
 				block, blockDiags := spec.newBlock(i, b.forEachCtx)
 				diags = append(diags, blockDiags...)
 				if block != nil {
-					block.Body = unknownBody{b.expandChild(block.Body, i)}
+					block.Body = unknownBody{
+						template:   b.expandChild(block.Body, i, marks),
+						valueMarks: marks,
+					}
 					blocks = append(blocks, block)
 				}
 			}
@@ -226,7 +242,7 @@ func (b *expandBody) expandBlocks(schema *hcl.BodySchema, rawBlocks hcl.Blocks,
 				// case it contains expressions that refer to our inherited
 				// iterators, or nested "dynamic" blocks.
 				expandedBlock := *rawBlock // shallow copy
-				expandedBlock.Body = b.expandChild(rawBlock.Body, b.iteration)
+				expandedBlock.Body = b.expandChild(rawBlock.Body, b.iteration, nil)
 				blocks = append(blocks, &expandedBlock)
 			}
 		}
@@ -235,11 +251,12 @@ func (b *expandBody) expandBlocks(schema *hcl.BodySchema, rawBlocks hcl.Blocks,
 	return blocks, diags
 }
 
-func (b *expandBody) expandChild(child hcl.Body, i *iteration) hcl.Body {
+func (b *expandBody) expandChild(child hcl.Body, i *iteration, valueMarks cty.ValueMarks) hcl.Body {
 	chiCtx := i.EvalContext(b.forEachCtx)
 	ret := Expand(child, chiCtx)
 	ret.(*expandBody).iteration = i
 	ret.(*expandBody).checkForEach = b.checkForEach
+	ret.(*expandBody).valueMarks = valueMarks
 	return ret
 }
 
diff --git a/ext/dynblock/expand_body_test.go b/ext/dynblock/expand_body_test.go
index bec6c21..3b245ee 100644
--- a/ext/dynblock/expand_body_test.go
+++ b/ext/dynblock/expand_body_test.go
@@ -8,6 +8,7 @@ import (
 	"testing"
 
 	"github.com/davecgh/go-spew/spew"
+	"github.com/google/go-cmp/cmp"
 	"github.com/hashicorp/hcl/v2"
 	"github.com/hashicorp/hcl/v2/hcldec"
 	"github.com/hashicorp/hcl/v2/hcltest"
@@ -710,6 +711,69 @@ func TestExpandUnknownBodies(t *testing.T) {
 
 }
 
+func TestExpandMarkedForEach(t *testing.T) {
+	srcBody := hcltest.MockBody(&hcl.BodyContent{
+		Blocks: hcl.Blocks{
+			{
+				Type:        "dynamic",
+				Labels:      []string{"b"},
+				LabelRanges: []hcl.Range{{}},
+				Body: hcltest.MockBody(&hcl.BodyContent{
+					Attributes: hcltest.MockAttrs(map[string]hcl.Expression{
+						"for_each": hcltest.MockExprLiteral(cty.TupleVal([]cty.Value{
+							cty.StringVal("hey"),
+						}).Mark("boop")),
+						"iterator": hcltest.MockExprTraversalSrc("dyn_b"),
+					}),
+					Blocks: hcl.Blocks{
+						{
+							Type: "content",
+							Body: hcltest.MockBody(&hcl.BodyContent{
+								Attributes: hcltest.MockAttrs(map[string]hcl.Expression{
+									"val0": hcltest.MockExprLiteral(cty.StringVal("static c 1")),
+									"val1": hcltest.MockExprTraversalSrc("dyn_b.value"),
+								}),
+							}),
+						},
+					},
+				}),
+			},
+		},
+	})
+
+	dynBody := Expand(srcBody, nil)
+
+	t.Run("Decode", func(t *testing.T) {
+		decSpec := &hcldec.BlockListSpec{
+			TypeName: "b",
+			Nested: &hcldec.ObjectSpec{
+				"val0": &hcldec.AttrSpec{
+					Name: "val0",
+					Type: cty.String,
+				},
+				"val1": &hcldec.AttrSpec{
+					Name: "val1",
+					Type: cty.String,
+				},
+			},
+		}
+
+		want := cty.ListVal([]cty.Value{
+			cty.ObjectVal(map[string]cty.Value{
+				"val0": cty.StringVal("static c 1").Mark("boop"),
+				"val1": cty.StringVal("hey").Mark("boop"),
+			}),
+		})
+		got, diags := hcldec.Decode(dynBody, decSpec, nil)
+		if diags.HasErrors() {
+			t.Fatalf("unexpected errors\n%s", diags.Error())
+		}
+		if diff := cmp.Diff(want, got, ctydebug.CmpOptions); diff != "" {
+			t.Errorf("wrong result\n%s", diff)
+		}
+	})
+}
+
 func TestExpandInvalidIteratorError(t *testing.T) {
 	srcBody := hcltest.MockBody(&hcl.BodyContent{
 		Blocks: hcl.Blocks{
diff --git a/ext/dynblock/expand_spec.go b/ext/dynblock/expand_spec.go
index 0231c4a..55a69ba 100644
--- a/ext/dynblock/expand_spec.go
+++ b/ext/dynblock/expand_spec.go
@@ -77,7 +77,8 @@ func (b *expandBody) decodeSpec(blockS *hcl.BlockHeaderSchema, rawSpec *hcl.Bloc
 		}
 	}
 
-	if !eachVal.CanIterateElements() && eachVal.Type() != cty.DynamicPseudoType {
+	unmarkedEachVal, _ := eachVal.Unmark()
+	if !unmarkedEachVal.CanIterateElements() && unmarkedEachVal.Type() != cty.DynamicPseudoType {
 		// We skip this error for DynamicPseudoType because that means we either
 		// have a null (which is checked immediately below) or an unknown
 		// (which is handled in the expandBody Content methods).
@@ -91,7 +92,7 @@ func (b *expandBody) decodeSpec(blockS *hcl.BlockHeaderSchema, rawSpec *hcl.Bloc
 		})
 		return nil, diags
 	}
-	if eachVal.IsNull() {
+	if unmarkedEachVal.IsNull() {
 		diags = append(diags, &hcl.Diagnostic{
 			Severity:    hcl.DiagError,
 			Summary:     "Invalid dynamic for_each value",
@@ -212,6 +213,28 @@ func (s *expandSpec) newBlock(i *iteration, ctx *hcl.EvalContext) (*hcl.Block, h
 			})
 			return nil, diags
 		}
+		if labelVal.IsMarked() {
+			// This situation is tricky because HCL just works generically
+			// with marks and so doesn't have any good language to talk about
+			// the meaning of specific mark types, but yet we cannot allow
+			// marked values here because the HCL API guarantees that a block's
+			// labels are always known static constant Go strings.
+			// Therefore this is a low-quality error message but at least
+			// better than panicking below when we call labelVal.AsString.
+			// If this becomes a problem then we could potentially add a new
+			// option for the public function [Expand] to allow calling
+			// applications to specify custom label validation functions that
+			// could then supersede this generic message.
+			diags = append(diags, &hcl.Diagnostic{
+				Severity:    hcl.DiagError,
+				Summary:     "Invalid dynamic block label",
+				Detail:      "This value has dynamic marks that make it unsuitable for use as a block label.",
+				Subject:     labelExpr.Range().Ptr(),
+				Expression:  labelExpr,
+				EvalContext: lCtx,
+			})
+			return nil, diags
+		}
 
 		labels = append(labels, labelVal.AsString())
 		labelRanges = append(labelRanges, labelExpr.Range())
diff --git a/ext/dynblock/expr_wrap.go b/ext/dynblock/expr_wrap.go
index 625bf9c..5763641 100644
--- a/ext/dynblock/expr_wrap.go
+++ b/ext/dynblock/expr_wrap.go
@@ -11,6 +11,19 @@ import (
 type exprWrap struct {
 	hcl.Expression
 	i *iteration
+
+	// resultMarks is a set of marks that must be applied to whatever
+	// value results from this expression. We do this whenever a
+	// dynamic block's for_each expression produced a marked result,
+	// since in that case any nested expressions inside are treated
+	// as being derived from that for_each expression.
+	//
+	// (calling applications might choose to reject marks by passing
+	// an [OptCheckForEach] to [Expand] and returning an error when
+	// marks are present, but this mechanism is here to help achieve
+	// reasonable behavior for situations where marks are permitted,
+	// which is the default.)
+	resultMarks cty.ValueMarks
 }
 
 func (e exprWrap) Variables() []hcl.Traversal {
@@ -34,8 +47,13 @@ func (e exprWrap) Variables() []hcl.Traversal {
 }
 
 func (e exprWrap) Value(ctx *hcl.EvalContext) (cty.Value, hcl.Diagnostics) {
+	if e.i == nil {
+		// If we don't have an active iteration then we can just use the
+		// given EvalContext directly.
+		return e.prepareValue(e.Expression.Value(ctx))
+	}
 	extCtx := e.i.EvalContext(ctx)
-	return e.Expression.Value(extCtx)
+	return e.prepareValue(e.Expression.Value(extCtx))
 }
 
 // UnwrapExpression returns the expression being wrapped by this instance.
@@ -43,3 +61,7 @@ func (e exprWrap) Value(ctx *hcl.EvalContext) (cty.Value, hcl.Diagnostics) {
 func (e exprWrap) UnwrapExpression() hcl.Expression {
 	return e.Expression
 }
+
+func (e exprWrap) prepareValue(val cty.Value, diags hcl.Diagnostics) (cty.Value, hcl.Diagnostics) {
+	return val.WithMarks(e.resultMarks), diags
+}
diff --git a/ext/dynblock/unknown_body.go b/ext/dynblock/unknown_body.go
index 6cd59c7..b1fdfc0 100644
--- a/ext/dynblock/unknown_body.go
+++ b/ext/dynblock/unknown_body.go
@@ -5,6 +5,7 @@ package dynblock
 
 import (
 	"github.com/hashicorp/hcl/v2"
+	"github.com/hashicorp/hcl/v2/hcldec"
 	"github.com/zclconf/go-cty/cty"
 )
 
@@ -18,16 +19,26 @@ import (
 // we instead arrange for everything _inside_ the block to be unknown instead,
 // to give the best possible approximation.
 type unknownBody struct {
-	template hcl.Body
+	template   hcl.Body
+	valueMarks cty.ValueMarks
 }
 
 var _ hcl.Body = unknownBody{}
 
-// hcldec.UnkownBody impl
+// hcldec.UnknownBody impl
 func (b unknownBody) Unknown() bool {
 	return true
 }
 
+// hcldec.MarkedBody impl
+func (b unknownBody) BodyValueMarks() cty.ValueMarks {
+	// We'll pass through to our template if it is a MarkedBody
+	if t, ok := b.template.(hcldec.MarkedBody); ok {
+		return t.BodyValueMarks()
+	}
+	return nil
+}
+
 func (b unknownBody) Content(schema *hcl.BodySchema) (*hcl.BodyContent, hcl.Diagnostics) {
 	content, diags := b.template.Content(schema)
 	content = b.fixupContent(content)
@@ -41,7 +52,7 @@ func (b unknownBody) Content(schema *hcl.BodySchema) (*hcl.BodyContent, hcl.Diag
 func (b unknownBody) PartialContent(schema *hcl.BodySchema) (*hcl.BodyContent, hcl.Body, hcl.Diagnostics) {
 	content, remain, diags := b.template.PartialContent(schema)
 	content = b.fixupContent(content)
-	remain = unknownBody{remain} // remaining content must also be wrapped
+	remain = unknownBody{template: remain, valueMarks: b.valueMarks} // remaining content must also be wrapped
 
 	// We're intentionally preserving the diagnostics reported from the
 	// inner body so that we can still report where the template body doesn't
@@ -69,8 +80,8 @@ func (b unknownBody) fixupContent(got *hcl.BodyContent) *hcl.BodyContent {
 	if len(got.Blocks) > 0 {
 		ret.Blocks = make(hcl.Blocks, 0, len(got.Blocks))
 		for _, gotBlock := range got.Blocks {
-			new := *gotBlock                      // shallow copy
-			new.Body = unknownBody{gotBlock.Body} // nested content must also be marked unknown
+			new := *gotBlock                                                          // shallow copy
+			new.Body = unknownBody{template: gotBlock.Body, valueMarks: b.valueMarks} // nested content must also be marked unknown
 			ret.Blocks = append(ret.Blocks, &new)
 		}
 	}
@@ -85,7 +96,7 @@ func (b unknownBody) fixupAttrs(got hcl.Attributes) hcl.Attributes {
 	ret := make(hcl.Attributes, len(got))
 	for name, gotAttr := range got {
 		new := *gotAttr // shallow copy
-		new.Expr = hcl.StaticExpr(cty.DynamicVal, gotAttr.Expr.Range())
+		new.Expr = hcl.StaticExpr(cty.DynamicVal.WithMarks(b.valueMarks), gotAttr.Expr.Range())
 		ret[name] = &new
 	}
 	return ret