cue/parser: allow let clause in comprehension

Closes #284

Change-Id: Ie051eb90f98f8e3cdf85d2aa3177ca6355619c9b
Reviewed-on: https://cue-review.googlesource.com/c/cue/+/7303
Reviewed-by: CUE cueckoo <[email protected]>
Reviewed-by: Marcel van Lohuizen <[email protected]>

Author: mpvl

cue/ast/astutil/resolve.go

@@ -413,20 +413,30 @@ func resolveIdent(s *scope, x *ast.Ident) bool {
 
 func scopeClauses(s *scope, clauses []ast.Clause) *scope {
 	for _, c := range clauses {
-		if f, ok := c.(*ast.ForClause); ok { // TODO(let): support let clause
-			walk(s, f.Source)
-			s = newScope(s.file, s, f, nil)
-			if f.Key != nil {
-				name, err := ast.ParseIdent(f.Key)
+		switch x := c.(type) {
+		case *ast.ForClause:
+			walk(s, x.Source)
+			s = newScope(s.file, s, x, nil)
+			if x.Key != nil {
+				name, err := ast.ParseIdent(x.Key)
 				if err == nil {
-					s.insert(name, f.Key, f)
+					s.insert(name, x.Key, x)
 				}
 			}
-			name, err := ast.ParseIdent(f.Value)
+			name, err := ast.ParseIdent(x.Value)
 			if err == nil {
-				s.insert(name, f.Value, f)
+				s.insert(name, x.Value, x)
 			}
-		} else {
+
+		case *ast.LetClause:
+			walk(s, x.Expr)
+			s = newScope(s.file, s, x, nil)
+			name, err := ast.ParseIdent(x.Ident)
+			if err == nil {
+				s.insert(name, x.Ident, x)
+			}
+
+		default:
 			walk(s, c)
 		}
 	}

cue/parser/parser.go

@@ -1123,11 +1123,20 @@ func (p *parser) parseComprehensionClauses(first bool) (clauses []ast.Clause, c
 				Condition: p.parseRHS(),
 			}))
 
-		// TODO:
-		// case token.LET:
-		// 	c := p.openComments()
-		// 	p.expect(token.LET)
-		// 	return nil, c
+		case token.LET:
+			c := p.openComments()
+			letPos := p.expect(token.LET)
+
+			ident := p.parseIdent()
+			assign := p.expect(token.BIND)
+			expr := p.parseRHS()
+
+			clauses = append(clauses, c.closeClause(p, &ast.LetClause{
+				Let:   letPos,
+				Ident: ident,
+				Equal: assign,
+				Expr:  expr,
+			}))
 
 		default:
 			return clauses, nil

cue/parser/parser_test.go

@@ -279,6 +279,17 @@ func TestParse(t *testing.T) {
 				}
 			 }`,
 		`{y: {a: 1, b: 2}, a: {for k: v in y if v>2 {"\(k)": v}}}`,
+	}, {
+		"nested comprehensions",
+		`{
+			y: { a: 1, b: 2}
+			a: {
+				for k, v in y let x = v+2 if x > 2 {
+					"\(k)": v
+				}
+			}
+		}`,
+		`{y: {a: 1, b: 2}, a: {for k: v in y let x=v+2 if x>2 {"\(k)": v}}}`,
 	}, {
 		"let declaration",
 		`{

cue/testdata/eval/comprehensions.txtar

@@ -13,6 +13,12 @@ b: {
 		l: 40
 	}
 }
+
+c: {
+  for k, v in a let y = v+10 if y > 50 {
+    "\(k)": y
+  }
+}
 -- out/eval --
 (struct){
   a: (struct){
@@ -26,6 +32,10 @@ b: {
     z: (int){ 50 }
     l: (int){ 40 }
   }
+  c: (struct){
+    y: (int){ 110 }
+    z: (int){ 60 }
+  }
 }
 -- out/compile --
 --- in.cue
@@ -48,4 +58,9 @@ b: {
       l: 40
     }
   }
+  c: {
+    for k, v in 〈1;a〉 let y = (〈0;v〉 + 10) if (〈0;y〉 > 50) {
+      "\(〈2;k〉)": 〈1;y〉
+    }
+  }
 }

internal/core/subsume/structural.go

@@ -275,6 +275,8 @@ func (c *collatedDecls) collate(env *adt.Environment, s *adt.StructLit) {
 			c.yielders = append(c.yielders, x)
 
 		case *adt.LetClause:
+			c.yielders = append(c.yielders, x)
+
 		case *adt.ValueClause:
 		}
 	}