resolver: resolve operand of first 'for' clause in enclosing block (#96)

In a comprehension such as [e for v in x], the expression x should
be resolved in the outer block, so [x for x in x] is legal.
This applies only to the first 'in' operand.

This behavior matches Python2 and Python3 and, apparently, Bazel,
though I'm pretty sure the existing test modified in this CL was
based on Bazel behavior, which must have changed in the meantime.

Updated spec and tests.

3 files changed, 41 insertions, 8 deletions

diff --git a/doc/spec.md b/doc/spec.md
index 8dec3fb..3b5a3b4 100644
--- a/doc/spec.md
+++ b/doc/spec.md
@@ -2086,6 +2086,32 @@ _ = [x for x in [2]]            # new variable x is local to the comprehension
 print(x)                        # 1
 ```
 
+The operand of a comprehension's first clause (always a `for`) is
+resolved in the lexical block enclosing the comprehension.
+In the examples below, identifiers referring to the outer variable
+named `x` have been distinguished by subscript.
+
+```python
+x₀ = (1, 2, 3)
+[x*x for x in x₀]               # [1, 4, 9]
+[x*x for x in x₀ if x%2 == 0]   # [4]
+```
+
+All subsequent `for` and `if` expressions are resolved within the
+comprehension's lexical block, as in this rather obscure example:
+
+```python
+x₀ = ([1, 2], [3, 4], [5, 6])
+[x*x for x in x₀ for x in x if x%2 == 0]     # [4, 16, 36]
+```
+
+which would be more clearly rewritten as:
+
+```python
+x = ([1, 2], [3, 4], [5, 6])
+[z*z for y in x for z in y if z%2 == 0]     # [4, 16, 36]
+```
+
 
 ### Function and method calls
 
diff --git a/resolve/resolve.go b/resolve/resolve.go
index ef3ddd6..17c922b 100644
--- a/resolve/resolve.go
+++ b/resolve/resolve.go
@@ -554,13 +554,21 @@ func (r *resolver) expr(e syntax.Expr) {
 		}
 
 	case *syntax.Comprehension:
+		// The 'in' operand of the first clause (always a ForClause)
+		// is resolved in the outer block; consider: [x for x in x].
+		clause := e.Clauses[0].(*syntax.ForClause)
+		r.expr(clause.X)
+
 		// A list/dict comprehension defines a new lexical block.
 		// Locals defined within the block will be allotted
 		// distinct slots in the locals array of the innermost
 		// enclosing container (function/module) block.
 		r.push(&block{comp: e})
+
 		const allowRebind = false
-		for _, clause := range e.Clauses {
+		r.assign(clause.Vars, allowRebind)
+
+		for _, clause := range e.Clauses[1:] {
 			switch clause := clause.(type) {
 			case *syntax.IfClause:
 				r.expr(clause.Cond)
diff --git a/testdata/assign.sky b/testdata/assign.sky
index 556041f..b209d26 100644
--- a/testdata/assign.sky
+++ b/testdata/assign.sky
@@ -210,14 +210,13 @@ assert.eq(type(set), "builtin_function_or_method")
 set = [1, 2, 3]
 assert.eq(type(set), "list")
 
----
-# All 'in x' expressions in a comprehension are evaluated
-# in the comprehension's lexical block.
-#
-# By contrast, Python yields [[1, 2], [1, 2]] because it evaluates
-# the first 'in x' in the environment enclosing the comprehension.
+# As in Python 2 and Python 3,
+# all 'in x' expressions in a comprehension are evaluated
+# in the comprehension's lexical block, except the first,
+# which is resolved in the outer block.
 x = [[1, 2]]
-_ = [x for x in x for y in x] ### "local variable x referenced before assignment"
+assert.eq([x for x in x for y in x],
+          [[1, 2], [1, 2]])
 
 ---
 # A comprehension establishes a single new lexical block,