Support rewriting of comprehension scopes (#1083)

* Support rewriting of comprehension scopes

* Revert "Workaround for #817 (#818)"

* Cleanup

Author: BCSharp

Src/IronPython/Compiler/Ast/Comprehension.cs

@@ -60,6 +60,15 @@ public override Ast Reduce() {
                 res
             );
         }
+
+        internal ComprehensionScope Scope { get; private protected set; }
+
+        internal Comprehension CopyForRewrite(ComprehensionScope scope) {
+            var newComprehension = (Comprehension)MemberwiseClone();
+            newComprehension.Scope = scope;
+            newComprehension.Parent = scope.Parent;
+            return newComprehension;
+        }
     }
 
     public sealed class ListComprehension : Comprehension {
@@ -106,10 +115,8 @@ public override void Walk(PythonWalker walker) {
             }
             walker.PostWalk(this);
         }
-
-        internal ComprehensionScope Scope { get; }
     }
-        
+
     public sealed class SetComprehension : Comprehension {
         private readonly ComprehensionIterator[] _iterators;
 
@@ -154,8 +161,6 @@ public override void Walk(PythonWalker walker) {
             }
             walker.PostWalk(this);
         }
-
-        internal ComprehensionScope Scope { get; }
     }
 
     public sealed class DictionaryComprehension : Comprehension {
@@ -207,19 +212,17 @@ public override void Walk(PythonWalker walker) {
             }
             walker.PostWalk(this);
         }
-
-        internal ComprehensionScope Scope { get; }
     }
 
     /// <summary>
     /// Scope for the comprehension.  Because scopes are usually statements and comprehensions are expressions
     /// this doesn't actually show up in the AST hierarchy and instead hangs off the comprehension expression.
     /// </summary>
     internal class ComprehensionScope : ScopeStatement {
-        private readonly Expression _comprehension;
+        private readonly Comprehension _comprehension;
         private static readonly MSAst.ParameterExpression _compContext = Ast.Parameter(typeof(CodeContext), "$compContext");
 
-        public ComprehensionScope(Expression comprehension) {
+        public ComprehensionScope(Comprehension comprehension) {
             _comprehension = comprehension;
         }
 
@@ -229,12 +232,12 @@ internal override bool ExposesLocalVariable(PythonVariable variable) {
             } else if (variable.Scope == this) {
                 return false;
             }
-            return _comprehension.Parent.ExposesLocalVariable(variable);
+            return Parent.ExposesLocalVariable(variable);
         }
 
         internal override MSAst.Expression/*!*/ GetParentClosureTuple() {
             Debug.Assert(NeedsLocalContext);
-            return MSAst.Expression.Call(null, typeof(PythonOps).GetMethod(nameof(PythonOps.GetClosureTupleFromContext)), _comprehension.Parent.LocalContext);
+            return MSAst.Expression.Call(null, typeof(PythonOps).GetMethod(nameof(PythonOps.GetClosureTupleFromContext)), Parent.LocalContext);
         }
 
         internal override bool TryBindOuter(ScopeStatement from, PythonReference reference, out PythonVariable variable) {
@@ -269,7 +272,7 @@ internal override PythonVariable BindReference(PythonNameBinder binder, PythonRe
             }
 
             // then bind in our parent scope
-            return _comprehension.Parent.BindReference(binder, reference);
+            return Parent.BindReference(binder, reference);
         }
 
         internal override Ast GetVariableExpression(PythonVariable variable) {
@@ -281,21 +284,23 @@ internal override Ast GetVariableExpression(PythonVariable variable) {
                 return expr;
             }
 
-            return _comprehension.Parent.GetVariableExpression(variable);
+            return Parent.GetVariableExpression(variable);
         }
 
         internal override Microsoft.Scripting.Ast.LightLambdaExpression GetLambda()
             => throw new NotImplementedException();
 
-        public override void Walk(PythonWalker walker) => _comprehension.Walk(walker);
+        public override void Walk(PythonWalker walker) {
+            _comprehension.Walk(walker);
+        }
 
         internal override Ast LocalContext {
             get {
                 if (NeedsLocalContext) {
                     return _compContext;
                 }
 
-                return _comprehension.Parent.LocalContext;
+                return Parent.LocalContext;
             }
         }
 
@@ -311,7 +316,7 @@ internal Ast AddVariables(Ast expression) {
             CreateVariables(locals, body);
 
             if (localContext != null) {
-                var createLocal = CreateLocalContext(_comprehension.Parent.LocalContext);
+                var createLocal = CreateLocalContext(Parent.LocalContext);
                 body.Add(Ast.Assign(_compContext, createLocal));
             }
 

Src/IronPython/Compiler/Ast/PythonAst.cs

@@ -699,7 +699,7 @@ internal void PrepareScope(ReadOnlyCollectionBuilder<MSAst.ParameterExpression>
 
         /// <summary>
         /// Rewrites the tree for performing lookups against globals instead of being bound
-        /// against the optimized scope.  This is used if the user compiles optimied code and then
+        /// against the optimized scope.  This is used if the user compiles optimized code and then
         /// runs it against a different scope.
         /// </summary>
         internal PythonAst MakeLookupCode() {
@@ -739,7 +739,7 @@ protected override MSAst.Expression VisitExtension(MSAst.Expression node) {
                     return PythonAst._globalContext;
                 }
 
-                // we need to re-write nested scoeps
+                // we need to re-write nested scopes
                 if (node is ScopeStatement scope) {
                     return base.VisitExtension(VisitScope(scope));
                 }
@@ -752,6 +752,10 @@ protected override MSAst.Expression VisitExtension(MSAst.Expression node) {
                     return base.VisitExtension(new GeneratorExpression((FunctionDefinition)VisitScope(generator.Function), generator.Iterable));
                 }
 
+                if (node is Comprehension comprehension) {
+                    return VisitComprehension(comprehension);
+                }
+
                 // update the global get/set/raw gets variables
                 if (node is PythonGlobalVariableExpression global) {
                     return new LookupGlobalVariable(
@@ -803,6 +807,22 @@ private ScopeStatement VisitScope(ScopeStatement scope) {
                 }
                 return newScope;
             }
+
+            private MSAst.Expression VisitComprehension(Comprehension comprehension) {
+                var newScope = (ComprehensionScope)comprehension.Scope.CopyForRewrite();
+                newScope.Parent = _curScope;
+                var newComprehension = comprehension.CopyForRewrite(newScope);
+
+                ScopeStatement prevScope = _curScope;
+                try {
+                    // rewrite the comprehension in a new scope
+                    _curScope = newScope;
+
+                    return base.VisitExtension(newComprehension);
+                } finally {
+                    _curScope = prevScope;
+                }
+            }
         }
 
         #endregion

Src/IronPythonTest/EngineTest.cs

@@ -618,6 +618,26 @@ public static void ScenarioTryGetMember() {
             Assert.AreEqual(result.ToString(), "IronPython.Runtime.Types.BuiltinFunction");
         }
 
+        [Test]
+        public static void ScenarioComprehensionScope() {
+            var engine = Python.CreateEngine();
+            var scope = engine.CreateScope();
+            var source = engine.CreateScriptSourceFromString("assert [lambda: i for i in [1,2]][0]() == 2");
+            var compiledCode = source.Compile();
+            compiledCode.Execute(scope);
+        }
+
+        [Test]
+        public static void ScenarioComprehensionScopeGlobal() {
+            var engine = Python.CreateEngine();
+            var source = engine.CreateScriptSourceFromString("assert [lambda: i for i in global_list][0]() == res");
+            var compiledCode = source.Compile();
+            var scope = engine.CreateScope();
+            scope.SetVariable("global_list", new[] { 1, 2 });
+            scope.SetVariable("res", 2);
+            compiledCode.Execute(scope);
+        }
+
         [Test]
         public static void ScenarioInterfaceExtensions() {
             var engine = Python.CreateEngine();

Src/StdLib/Lib/sre_compile.py

@@ -64,7 +64,7 @@
 )
 
 # Maps the lowercase code to lowercase codes which have the same uppercase.
-_ignorecase_fixes = {i: tuple([j for j in t if i != j]) # https://github.com/IronLanguages/ironpython3/issues/817
+_ignorecase_fixes = {i: tuple(j for j in t if i != j)
                      for t in _equivalences for i in t}
 
 def _compile(code, pattern, flags):

Tests/test_closure.py

@@ -2,6 +2,7 @@
 # The .NET Foundation licenses this file to you under the Apache 2.0 License.
 # See the LICENSE file in the project root for more information.
 
+import types
 import unittest
 
 from iptest import run_test
@@ -322,4 +323,38 @@ class C:
             exec("z = 7", globals())
             self.assertEqual(z, 7)
 
+    def test_gh817(self):
+        # https://github.com/IronLanguages/ironpython3/issues/817
+
+        for x in [lambda: i for i in [1,2]]:
+            self.assertEqual(x(), 2)
+
+        self.assertEqual([tuple(i for j in [1]) for i in [1]], [(1,)])
+        self.assertEqual({tuple(i for j in [1]) for i in [1]}, {(1,)})
+
+        self.assertEqual({i: tuple(j for j in t if i != j) for t in ((1,2),) for i in t}, {1: (2,), 2: (1,)})
+
+        self.assertEqual([lambda: i for i in [1]][0](), 1)
+
+        self.assertEqual([x * a for a in range(3) if a == 2 for x in range(5,7)], [10, 12])
+        self.assertRaises(UnboundLocalError, lambda: [x * a for a in range(3) if x == 2 for x in range(5,7)])
+
+        def foo1(z):
+            return [(x for x in range(y)) for y in range(z)]
+        self.assertEqual([list(x) for x in foo1(3)], [[], [0], [0, 1]])
+
+        def foo2(z):
+            return [(x + z for x in range(y)) for y in range(z)]
+        self.assertEqual([list(x) for x in foo2(3)], [[], [3], [3, 4]])
+
+        dl = [{ y: (z**2 for z in range(x)) for y in range(3)} for x in range(4)]
+        self.assertIsInstance(dl, list)
+        self.assertEqual(len(dl), 4)
+        for p, d in enumerate(dl):
+            self.assertIsInstance(d, dict)
+            for k in range(3):
+                g = d[k]
+                self.assertIsInstance(g, types.GeneratorType)
+                self.assertEqual(list(g), [0, 1, 4][:p])
+
 run_test(__name__)