cue: adjust Environment for comprehensions

Not doing so could result in bad dereferences.

Closes #2911

Signed-off-by: Marcel van Lohuizen <[email protected]>
Change-Id: Iff82628f3475c75b35e030569eb87e0363d66ac8
Dispatch-Trailer: {"type":"trybot","CL":1185796,"patchset":1,"ref":"refs/changes/96/1185796/1","targetBranch":"master"}

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 that where evaluation of the expression is
+// not required. To get an expression paired with the correct environment, use
+// ToExpr.
+//
+// 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)