resolver: resolve operand of first 'for' clause in enclosing block (g…

…oogle#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.

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
 

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)

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,