cue: adjust Environment for comprehensions

Not doing so could result in bad dereferences.

CL https://cuelang.org/cl/529517 introduced a way to
access the expression associated with a comprehension.

In cases where this expression was evaluated within the
original Environment of the Comprehension, this resulted
in a panic due to an unaligned Environment. The problem
was that this value expression needed to be evaluated
using an adjusted Environment that accounts for the
interstitial scopes introduced by the for and let clauses.
In general, these scopes cannot be added ahead of time,
as the associated values are not known before evaluation.

Like is done in other parts of the code, the new EnvExpr
method and function adjusts the Environment by inserting
dummy scopes that allow references in the expression that
reference values outside of the comprehension to be
resolved properly.

The old Expr method and function are retained for use cases
that only need to refer to the value expression without the
need to evaluate or resolve it.

Closes #2911

Signed-off-by: Marcel van Lohuizen <mpvl@gmail.com>
Change-Id: Ibecd4738115456d408f6c7eb505fd00c66e9d3f0
Reviewed-on: https://review.gerrithub.io/c/cue-lang/cue/+/1186144
Reviewed-by: Daniel Martí <mvdan@mvdan.cc>
TryBot-Result: CUEcueckoo <cueckoo@gmail.com>
Unity-Result: CUE porcuepine <cue.porcuepine@gmail.com>
(cherry picked from commit 3e6a8be)
Reviewed-on: https://review.gerrithub.io/c/cue-lang/cue/+/1190979
Reviewed-by: Paul Jolly <paul@myitcv.io>
TryBot-Result: CUEcueckoo <cueckoo@cuelang.org>

cue/types.go

@@ -645,12 +645,17 @@ func Dereference(v Value) Value {
 		return v
 	}
 
-	c := n.Conjuncts[0]
-	r, _ := c.Expr().(adt.Resolver)
+	env, expr := n.Conjuncts[0].EnvExpr()
+
+	// TODO: consider supporting unwrapping of structs or comprehensions around
+	// a single embedded reference.
+	r, _ := expr.(adt.Resolver)
 	if r == nil {
 		return v
 	}
 
+	c := adt.MakeRootConjunct(env, expr)
+
 	ctx := v.ctx()
 	n, b := ctx.Resolve(c, r)
 	if b != nil {
@@ -1067,7 +1072,8 @@ func (v hiddenValue) Split() []Value {
 	}
 	a := []Value{}
 	for _, x := range v.v.Conjuncts {
-		a = append(a, remakeValue(v, x.Env, x.Expr()))
+		env, expr := x.EnvExpr()
+		a = append(a, remakeValue(v, env, expr))
 	}
 	return a
 }
@@ -1983,7 +1989,9 @@ func (v Value) ReferencePath() (root Value, p Path) {
 	ctx := v.ctx()
 	c := v.v.Conjuncts[0]
 
-	x, path := reference(v.idx, ctx, c.Env, c.Expr())
+	env, expr := c.EnvExpr()
+
+	x, path := reference(v.idx, ctx, env, expr)
 	if x == nil {
 		return Value{}, Path{}
 	}
@@ -2319,8 +2327,7 @@ func (v Value) Expr() (Op, []Value) {
 		case 1:
 			// the default case, processed below.
 			c := v.v.Conjuncts[0]
-			env = c.Env
-			expr = c.Expr()
+			env, expr = c.EnvExpr()
 			if w, ok := expr.(*adt.Vertex); ok {
 				return Value{v.idx, w, v.parent_}.Expr()
 			}

cue/types_test.go

@@ -3084,6 +3084,9 @@ func TestReferencePath(t *testing.T) {
 	}, {
 		input: "v: w: x: a.b.c, a: b: c: 1",
 		want:  "a.b.c",
+	}, {
+		input: "if true { v: w: x: a, a: 1 }",
+		want:  "a",
 	}, {
 		input: "v: w: x: w.a.b.c, v: w: a: b: c: 1",
 		want:  "v.w.a.b.c",
@@ -3447,6 +3450,38 @@ func TestPathCorrection(t *testing.T) {
 			v = v.Lookup("t")
 			return v
 		},
+	}, {
+		input: `
+		x: { if true { v: a } }
+		a: b
+		b: 2
+		`,
+		want: "b",
+		lookup: func(inst *Instance) Value {
+			v := inst.Value().LookupPath(ParsePath("x.v"))
+			v = Dereference(v)
+			return v
+		},
+	}, {
+		input: `
+		package foo
+
+		#A:{ if true { #B: #T } }
+
+		#T: {
+			a: #S.#U
+			#S: #U: {}
+		}
+		`,
+		want: "#T.#S.#U",
+		lookup: func(inst *Instance) Value {
+			f, _ := inst.Value().LookupField("#A")
+			f, _ = f.Value.LookupField("#B")
+			v := f.Value
+			v = Dereference(v)
+			v = v.Lookup("a")
+			return v
+		},
 	}}
 	for _, tc := range testCases {
 		if tc.skip {
@@ -3723,6 +3758,9 @@ func TestExpr(t *testing.T) {
 	}, {
 		input: `v: {>30, <40}`,
 		want:  `&(>(30) <(40))`,
+	}, {
+		input: `a: string, if true { v: a }`,
+		want:  `.(〈〉 "a")`,
 	}}
 	for _, tc := range testCases {
 		t.Run(tc.input, func(t *testing.T) {

internal/core/adt/composite.go

@@ -1164,12 +1164,39 @@ func (c *Conjunct) Elem() Elem {
 	}
 }
 
-// Expr retrieves the expression form of the contained conjunct.
-// If it is a field or comprehension, it will return its associated value.
+// Expr retrieves the expression form of the contained conjunct. If it is a
+// field or comprehension, it will return its associated value. This is only to
+// be used for syntactic operations where evaluation of the expression is not
+// required. To get an expression paired with the correct environment, use
+// EnvExpr.
+//
+// TODO: rename to RawExpr.
 func (c *Conjunct) Expr() Expr {
 	return ToExpr(c.x)
 }
 
+// EnvExpr returns the expression form of the contained conjunct alongside an
+// Environment in which this expression should be evaluated.
+func (c Conjunct) EnvExpr() (*Environment, Expr) {
+	return EnvExpr(c.Env, c.Elem())
+}
+
+// EnvExpr returns the expression represented by Elem alongside an Environment
+// with the necessary adjustments in which the resulting expression can be
+// evaluated.
+func EnvExpr(env *Environment, elem Elem) (*Environment, Expr) {
+	for {
+		if c, ok := elem.(*Comprehension); ok {
+			env = linkChildren(env, c)
+			c := MakeConjunct(env, c.Value, CloseInfo{})
+			elem = c.Elem()
+			continue
+		}
+		break
+	}
+	return env, ToExpr(elem)
+}
+
 // ToExpr extracts the underlying expression for a Node. If something is already
 // an Expr, it will return it as is, if it is a field, it will return its value,
 // and for comprehensions it returns the yielded struct.

internal/core/dep/dep.go

@@ -621,6 +621,7 @@ func (c *visitor) markComprehension(env *adt.Environment, y *adt.Comprehension)
 	for i := y.Nest(); i > 0; i-- {
 		env = &adt.Environment{Up: env, Vertex: empty}
 	}
+	// TODO: consider using adt.EnvExpr and remove the above loop.
 	c.markExpr(env, adt.ToExpr(y.Value))
 }
 

internal/core/export/adt.go

@@ -263,7 +263,7 @@ func (e *exporter) adt(env *adt.Environment, expr adt.Elem) ast.Expr {
 	case *adt.ConjunctGroup:
 		a := []ast.Expr{}
 		for _, c := range *x {
-			v := e.expr(env, c.Expr())
+			v := e.expr(c.EnvExpr())
 			a = append(a, v)
 		}
 		return ast.NewBinExpr(token.AND, a...)
@@ -293,6 +293,7 @@ func (e *exporter) adt(env *adt.Environment, expr adt.Elem) ast.Expr {
 			env = &adt.Environment{Up: env, Vertex: empty}
 		}
 
+		// TODO: consider using adt.EnvExpr.
 		return e.adt(env, adt.ToExpr(x.Value))
 
 	default:
@@ -704,6 +705,7 @@ func (e *exporter) comprehension(env *adt.Environment, comp *adt.Comprehension)
 		env = &adt.Environment{Up: env, Vertex: empty}
 	}
 
+	// TODO: consider using adt.EnvExpr.
 	v := e.expr(env, adt.ToExpr(comp.Value))
 	if _, ok := v.(*ast.StructLit); !ok {
 		v = ast.NewStruct(ast.Embed(v))

internal/core/export/export.go

@@ -571,7 +571,7 @@ func (e *exporter) resolveLet(env *adt.Environment, x *adt.LetReference) ast.Exp
 
 			return e.expr(env, x.X)
 		}
-		return e.expr(env, ref.Conjuncts[0].Expr())
+		return e.expr(ref.Conjuncts[0].EnvExpr())
 
 	case let.Expr == nil:
 		label := e.uniqueLetIdent(x.Label, x.X)