Change this capturing algorithm for converted loops

rbuckton · rbuckton · commit 619e116dedae · 2016-09-24T15:07:54.000-07:00
diff --git a/src/compiler/transformers/es6.ts b/src/compiler/transformers/es6.ts
@@ -165,10 +165,9 @@ namespace ts {
         let currentNode: Node;
         let enclosingBlockScopeContainer: Node;
         let enclosingBlockScopeContainerParent: Node;
-        let containingNonArrowFunction: FunctionLikeDeclaration | ClassElement;
-
-        /** Tracks the container that determines whether `super.x` is a static. */
-        let superScopeContainer: FunctionLikeDeclaration | ClassElement;
+        let enclosingFunction: FunctionLikeDeclaration;
+        let enclosingNonArrowFunction: FunctionLikeDeclaration;
+        let enclosingNonAsyncFunctionBody: FunctionLikeDeclaration | ClassElement;
 
         /**
          * Used to track if we are emitting body of the converted loop
@@ -182,11 +181,6 @@ namespace ts {
          */
         let enabledSubstitutions: ES6SubstitutionFlags;
 
-        /**
-         * This is used to determine whether we need to emit `_this` instead of `this`.
-         */
-        let useCapturedThis: boolean;
-
         return transformSourceFile;
 
         function transformSourceFile(node: SourceFile) {
@@ -206,13 +200,13 @@ namespace ts {
         }
 
         function saveStateAndInvoke<T>(node: Node, f: (node: Node) => T): T {
-            const savedContainingNonArrowFunction = containingNonArrowFunction;
-            const savedSuperScopeContainer = superScopeContainer;
-            const savedCurrentParent = currentParent;
-            const savedCurrentNode = currentNode;
+            const savedEnclosingFunction = enclosingFunction;
+            const savedEnclosingNonArrowFunction = enclosingNonArrowFunction;
+            const savedEnclosingNonAsyncFunctionBody = enclosingNonAsyncFunctionBody;
             const savedEnclosingBlockScopeContainer = enclosingBlockScopeContainer;
             const savedEnclosingBlockScopeContainerParent = enclosingBlockScopeContainerParent;
-
+            const savedCurrentParent = currentParent;
+            const savedCurrentNode = currentNode;
             const savedConvertedLoopState = convertedLoopState;
             if (nodeStartsNewLexicalEnvironment(node)) {
                 // don't treat content of nodes that start new lexical environment as part of converted loop copy
@@ -223,12 +217,13 @@ namespace ts {
             const visited = f(node);
 
             convertedLoopState = savedConvertedLoopState;
-            containingNonArrowFunction = savedContainingNonArrowFunction;
-            superScopeContainer = savedSuperScopeContainer;
-            currentParent = savedCurrentParent;
-            currentNode = savedCurrentNode;
+            enclosingFunction = savedEnclosingFunction;
+            enclosingNonArrowFunction = savedEnclosingNonArrowFunction;
+            enclosingNonAsyncFunctionBody = savedEnclosingNonAsyncFunctionBody;
             enclosingBlockScopeContainer = savedEnclosingBlockScopeContainer;
             enclosingBlockScopeContainerParent = savedEnclosingBlockScopeContainerParent;
+            currentParent = savedCurrentParent;
+            currentNode = savedCurrentNode;
             return visited;
         }
 
@@ -251,22 +246,13 @@ namespace ts {
         }
 
         function visitorForConvertedLoopWorker(node: Node): VisitResult<Node> {
-            const savedUseCapturedThis = useCapturedThis;
-
-            if (nodeStartsNewLexicalEnvironment(node)) {
-                useCapturedThis = false;
-            }
-
             let result: VisitResult<Node>;
-
             if (shouldCheckNode(node)) {
                 result = visitJavaScript(node);
             }
             else {
                 result = visitNodesInConvertedLoop(node);
             }
-
-            useCapturedThis = savedUseCapturedThis;
             return result;
         }
 
@@ -409,30 +395,25 @@ namespace ts {
         }
 
         function onBeforeVisitNode(node: Node) {
-            const currentGrandparent = currentParent;
-            currentParent = currentNode;
-            currentNode = node;
-
-            if (currentParent) {
-                if (isBlockScope(currentParent, currentGrandparent)) {
-                    enclosingBlockScopeContainer = currentParent;
-                    enclosingBlockScopeContainerParent = currentGrandparent;
+            if (currentNode) {
+                if (isBlockScope(currentNode, currentParent)) {
+                    enclosingBlockScopeContainer = currentNode;
+                    enclosingBlockScopeContainerParent = currentParent;
                 }
 
-                switch (currentParent.kind) {
-                    case SyntaxKind.FunctionExpression:
-                    case SyntaxKind.Constructor:
-                    case SyntaxKind.MethodDeclaration:
-                    case SyntaxKind.GetAccessor:
-                    case SyntaxKind.SetAccessor:
-                    case SyntaxKind.FunctionDeclaration:
-                        containingNonArrowFunction = <FunctionLikeDeclaration>currentParent;
-                        if (!(containingNonArrowFunction.emitFlags & NodeEmitFlags.AsyncFunctionBody)) {
-                            superScopeContainer = containingNonArrowFunction;
+                if (isFunctionLike(currentNode)) {
+                    enclosingFunction = currentNode;
+                    if (currentNode.kind !== SyntaxKind.ArrowFunction) {
+                        enclosingNonArrowFunction = currentNode;
+                        if (!(currentNode.emitFlags & NodeEmitFlags.AsyncFunctionBody)) {
+                            enclosingNonAsyncFunctionBody = currentNode;
                         }
-                        break;
+                    }
                 }
             }
+
+            currentParent = currentNode;
+            currentNode = node;
         }
 
         function visitSwitchStatement(node: SwitchStatement): SwitchStatement {
@@ -468,9 +449,8 @@ namespace ts {
 
         function visitThisKeyword(node: Node): Node {
             Debug.assert(convertedLoopState !== undefined);
-
-            if (useCapturedThis) {
-                // if useCapturedThis is true then 'this' keyword is contained inside an arrow function.
+            if (enclosingFunction && enclosingFunction.kind === SyntaxKind.ArrowFunction) {
+                // if the enclosing function is an ArrowFunction is then we use the captured 'this' keyword.
                 convertedLoopState.containsLexicalThis = true;
                 return node;
             }
@@ -1306,12 +1286,7 @@ namespace ts {
                 enableSubstitutionsForCapturedThis();
             }
 
-            const savedUseCapturedThis = useCapturedThis;
-            useCapturedThis = true;
-
             const func = transformFunctionLikeToExpression(node, /*location*/ node, /*name*/ undefined);
-
-            useCapturedThis = savedUseCapturedThis;
             setNodeEmitFlags(func, NodeEmitFlags.CapturesThis);
             return func;
         }
@@ -1354,9 +1329,9 @@ namespace ts {
          * @param name The name of the new FunctionExpression.
          */
         function transformFunctionLikeToExpression(node: FunctionLikeDeclaration, location: TextRange, name: Identifier): FunctionExpression {
-            const savedContainingNonArrowFunction = containingNonArrowFunction;
+            const savedContainingNonArrowFunction = enclosingNonArrowFunction;
             if (node.kind !== SyntaxKind.ArrowFunction) {
-                containingNonArrowFunction = node;
+                enclosingNonArrowFunction = node;
             }
 
             const expression = setOriginalNode(
@@ -1372,7 +1347,7 @@ namespace ts {
                 /*original*/ node
             );
 
-            containingNonArrowFunction = savedContainingNonArrowFunction;
+            enclosingNonArrowFunction = savedContainingNonArrowFunction;
             return expression;
         }
 
@@ -2066,8 +2041,8 @@ namespace ts {
             }
 
             const isAsyncBlockContainingAwait =
-                containingNonArrowFunction
-                && (containingNonArrowFunction.emitFlags & NodeEmitFlags.AsyncFunctionBody) !== 0
+                enclosingNonArrowFunction
+                && (enclosingNonArrowFunction.emitFlags & NodeEmitFlags.AsyncFunctionBody) !== 0
                 && (node.statement.transformFlags & TransformFlags.ContainsYield) !== 0;
 
             let loopBodyFlags: NodeEmitFlags = 0;
@@ -2829,9 +2804,9 @@ namespace ts {
          * Visits the `super` keyword
          */
         function visitSuperKeyword(node: PrimaryExpression): LeftHandSideExpression {
-            return superScopeContainer
-                && isClassElement(superScopeContainer)
-                && !hasModifier(superScopeContainer, ModifierFlags.Static)
+            return enclosingNonAsyncFunctionBody
+                && isClassElement(enclosingNonAsyncFunctionBody)
+                && !hasModifier(enclosingNonAsyncFunctionBody, ModifierFlags.Static)
                 && currentParent.kind !== SyntaxKind.CallExpression
                     ? createPropertyAccess(createIdentifier("_super"), "prototype")
                     : createIdentifier("_super");
@@ -2856,17 +2831,16 @@ namespace ts {
          * @param node The node to be printed.
          */
         function onEmitNode(node: Node, emit: (node: Node) => void) {
-            const savedUseCapturedThis = useCapturedThis;
+            const savedEnclosingFunction = enclosingFunction;
 
             if (enabledSubstitutions & ES6SubstitutionFlags.CapturedThis && isFunctionLike(node)) {
-                // If we are tracking a captured `this`, push a bit that indicates whether the
-                // containing function is an arrow function.
-                useCapturedThis = (getNodeEmitFlags(node) & NodeEmitFlags.CapturesThis) !== 0;
+                // If we are tracking a captured `this`, keep track of the enclosing function.
+                enclosingFunction = node;
             }
 
             previousOnEmitNode(node, emit);
 
-            useCapturedThis = savedUseCapturedThis;
+            enclosingFunction = savedEnclosingFunction;
         }
 
         /**
@@ -2994,7 +2968,9 @@ namespace ts {
          * @param node The ThisKeyword node.
          */
         function substituteThisKeyword(node: PrimaryExpression): PrimaryExpression {
-            if (enabledSubstitutions & ES6SubstitutionFlags.CapturedThis && useCapturedThis) {
+            if (enabledSubstitutions & ES6SubstitutionFlags.CapturedThis
+                && enclosingFunction
+                && enclosingFunction.emitFlags & NodeEmitFlags.CapturesThis) {
                 return createIdentifier("_this", /*location*/ node);
             }
 
diff --git a/tests/baselines/reference/blockScopedBindingCaptureThisInFunction.js b/tests/baselines/reference/blockScopedBindingCaptureThisInFunction.js
@@ -0,0 +1,21 @@
+//// [blockScopedBindingCaptureThisInFunction.ts]
+// https://github.com/Microsoft/TypeScript/issues/11038
+() => function () {
+    for (let someKey in {}) {
+        this.helloWorld();
+        () => someKey;
+    }
+};
+
+//// [blockScopedBindingCaptureThisInFunction.js]
+// https://github.com/Microsoft/TypeScript/issues/11038
+(function () { return function () {
+    var _loop_1 = function (someKey) {
+        this_1.helloWorld();
+        (function () { return someKey; });
+    };
+    var this_1 = this;
+    for (var someKey in {}) {
+        _loop_1(someKey);
+    }
+}; });
diff --git a/tests/baselines/reference/blockScopedBindingCaptureThisInFunction.symbols b/tests/baselines/reference/blockScopedBindingCaptureThisInFunction.symbols
@@ -0,0 +1,11 @@
+=== tests/cases/compiler/blockScopedBindingCaptureThisInFunction.ts ===
+// https://github.com/Microsoft/TypeScript/issues/11038
+() => function () {
+    for (let someKey in {}) {
+>someKey : Symbol(someKey, Decl(blockScopedBindingCaptureThisInFunction.ts, 2, 12))
+
+        this.helloWorld();
+        () => someKey;
+>someKey : Symbol(someKey, Decl(blockScopedBindingCaptureThisInFunction.ts, 2, 12))
+    }
+};
diff --git a/tests/baselines/reference/blockScopedBindingCaptureThisInFunction.types b/tests/baselines/reference/blockScopedBindingCaptureThisInFunction.types
@@ -0,0 +1,21 @@
+=== tests/cases/compiler/blockScopedBindingCaptureThisInFunction.ts ===
+// https://github.com/Microsoft/TypeScript/issues/11038
+() => function () {
+>() => function () {    for (let someKey in {}) {        this.helloWorld();        () => someKey;    }} : () => () => void
+>function () {    for (let someKey in {}) {        this.helloWorld();        () => someKey;    }} : () => void
+
+    for (let someKey in {}) {
+>someKey : string
+>{} : {}
+
+        this.helloWorld();
+>this.helloWorld() : any
+>this.helloWorld : any
+>this : any
+>helloWorld : any
+
+        () => someKey;
+>() => someKey : () => string
+>someKey : string
+    }
+};
diff --git a/tests/cases/compiler/blockScopedBindingCaptureThisInFunction.ts b/tests/cases/compiler/blockScopedBindingCaptureThisInFunction.ts
@@ -0,0 +1,7 @@
+// https://github.com/Microsoft/TypeScript/issues/11038
+() => function () {
+    for (let someKey in {}) {
+        this.helloWorld();
+        () => someKey;
+    }
+};
