microsoft · rbuckton · Mar 18, 2016 · Feb 4, 2016 · Feb 5, 2016 · Feb 5, 2016
diff --git a/src/compiler/printer.ts b/src/compiler/printer.ts
@@ -152,8 +152,7 @@ const _super = (function (geti, seti) {
             let isEmitNotificationEnabled: (node: Node) => boolean;
             let expressionSubstitution: (node: Expression) => Expression;
             let identifierSubstitution: (node: Identifier) => Identifier;
-            let onBeforeEmitNode: (node: Node) => void;
-            let onAfterEmitNode: (node: Node) => void;
+            let onEmitNode: (node: Node, emit: (node: Node) => void) => void;
             let nodeToGeneratedName: string[];
             let generatedNameSet: Map<string>;
             let tempFlags: TempFlags;
@@ -215,8 +214,7 @@ const _super = (function (geti, seti) {
                 isEmitNotificationEnabled = undefined;
                 expressionSubstitution = undefined;
                 identifierSubstitution = undefined;
-                onBeforeEmitNode = undefined;
-                onAfterEmitNode = undefined;
+                onEmitNode = undefined;
                 tempFlags = TempFlags.Auto;
                 currentSourceFile = undefined;
                 currentText = undefined;
@@ -237,8 +235,7 @@ const _super = (function (geti, seti) {
                 isEmitNotificationEnabled = context.isEmitNotificationEnabled;
                 expressionSubstitution = context.expressionSubstitution;
                 identifierSubstitution = context.identifierSubstitution;
-                onBeforeEmitNode = context.onBeforeEmitNode;
-                onAfterEmitNode = context.onAfterEmitNode;
+                onEmitNode = context.onEmitNode;
                 return printSourceFile;
             }
 
@@ -252,46 +249,62 @@ const _super = (function (geti, seti) {
                 return node;
             }
 
+            /**
+             * Emits a node.
+             */
             function emit(node: Node) {
-                emitWithWorker(node, emitWorker);
+                emitNodeWithNotificationOption(node, emitWithoutNotificationOption);
+            }
+
+            /**
+             * Emits a node without calling onEmitNode.
+             * NOTE: Do not call this method directly.
+             */
+            function emitWithoutNotificationOption(node: Node) {
+                emitNodeWithWorker(node, emitWorker);
             }
 
+            /**
+             * Emits an expression node.
+             */
             function emitExpression(node: Expression) {
-                emitWithWorker(node, emitExpressionWorker);
+                emitNodeWithNotificationOption(node, emitExpressionWithoutNotificationOption);
+            }
+
+            /**
+             * Emits an expression without calling onEmitNode.
+             * NOTE: Do not call this method directly.
+             */
+            function emitExpressionWithoutNotificationOption(node: Expression) {
+                emitNodeWithWorker(node, emitExpressionWorker);
             }
 
-            function emitWithWorker(node: Node, emitWorker: (node: Node) => void) {
+            /**
+             * Emits a node with emit notification if available.
+             */
+            function emitNodeWithNotificationOption(node: Node, emit: (node: Node) => void) {
                 if (node) {
-                    const adviseOnEmit = isEmitNotificationEnabled(node);
-                    if (adviseOnEmit && onBeforeEmitNode) {
-                        onBeforeEmitNode(node);
+                    if (isEmitNotificationEnabled(node)) {
+                        onEmitNode(node, emit);
+                    }
+                    else {
+                        emit(node);
                     }
+                }
+            }
 
+            function emitNodeWithWorker(node: Node, emitWorker: (node: Node) => void) {
+                if (node) {
                     const leadingComments = getLeadingComments(node, getNotEmittedParent);
                     const trailingComments = getTrailingComments(node, getNotEmittedParent);
                     emitLeadingComments(node, leadingComments);
                     emitStart(node);
                     emitWorker(node);
                     emitEnd(node);
                     emitTrailingComments(node, trailingComments);
-
-                    if (adviseOnEmit && onAfterEmitNode) {
-                        onAfterEmitNode(node);
-                    }
                 }
             }
 
-            function getNotEmittedParent(node: Node): Node {
-                if (getNodeEmitFlags(node) & NodeEmitFlags.EmitCommentsOfNotEmittedParent) {
-                    const parent = getOriginalNode(node).parent;
-                    if (getNodeEmitFlags(parent) & NodeEmitFlags.IsNotEmittedNode) {
-                        return parent;
-                    }
-                }
-
-                return undefined;
-            }
-
             function emitWorker(node: Node): void {
                 const kind = node.kind;
                 switch (kind) {
@@ -2367,6 +2380,17 @@ const _super = (function (geti, seti) {
                     && rangeEndIsOnSameLineAsRangeStart(block, block);
             }
 
+            function getNotEmittedParent(node: Node): Node {
+                if (getNodeEmitFlags(node) & NodeEmitFlags.EmitCommentsOfNotEmittedParent) {
+                    const parent = getOriginalNode(node).parent;
+                    if (getNodeEmitFlags(parent) & NodeEmitFlags.IsNotEmittedNode) {
+                        return parent;
+                    }
+                }
+
+                return undefined;
+            }
+
             function isUniqueName(name: string): boolean {
                 return !resolver.hasGlobalName(name) &&
                     !hasProperty(currentFileIdentifiers, name) &&

diff --git a/src/compiler/transformer.ts b/src/compiler/transformer.ts
@@ -78,8 +78,7 @@ namespace ts {
             expressionSubstitution: node => node,
             enableExpressionSubstitution,
             isExpressionSubstitutionEnabled,
-            onBeforeEmitNode: node => { },
-            onAfterEmitNode: node => { },
+            onEmitNode: (node, emit) => emit(node),
             enableEmitNotification,
             isEmitNotificationEnabled,
         };

diff --git a/src/compiler/transformers/ts.ts b/src/compiler/transformers/ts.ts
@@ -28,13 +28,11 @@ namespace ts {
         const languageVersion = getEmitScriptTarget(compilerOptions);
 
         // Save the previous transformation hooks.
-        const previousOnBeforeEmitNode = context.onBeforeEmitNode;
-        const previousOnAfterEmitNode = context.onAfterEmitNode;
+        const previousOnEmitNode = context.onEmitNode;
         const previousExpressionSubstitution = context.expressionSubstitution;
 
         // Set new transformation hooks.
-        context.onBeforeEmitNode = onBeforeEmitNode;
-        context.onAfterEmitNode = onAfterEmitNode;
+        context.onEmitNode = onEmitNode;
         context.expressionSubstitution = substituteExpression;
 
         // These variables contain state that changes as we descend into the tree.
@@ -64,19 +62,16 @@ namespace ts {
         let currentDecoratedClassAliases: Map<Identifier>;
 
         /**
-         * Keeps track of how deeply nested we are within any containing namespaces
-         * when performing just-in-time substitution while printing an expression identifier.
-         * If the nest level is greater than zero, then we are performing a substitution
-         * inside of a namespace and we should perform the more costly checks to determine
-         * whether the identifier points to an exported declaration.
+         * Keeps track of whether  we are within any containing namespaces when performing
+         * just-in-time substitution while printing an expression identifier.
          */
-        let namespaceNestLevel: number;
+        let isEnclosedInNamespace: boolean;
 
         /**
-         * This array keeps track of containers where `super` is valid, for use with
+         * This keeps track of containers where `super` is valid, for use with
          * just-in-time substitution for `super` expressions inside of async methods.
          */
-        let superContainerStack: SuperContainer[];
+        let currentSuperContainer: SuperContainer;
 
         return transformSourceFile;
 
@@ -2423,21 +2418,24 @@ namespace ts {
             //      x_1.y = ...;
             //  })(x || (x = {}));
             statements.push(
-                setOriginalNode(
-                    createStatement(
-                        createCall(
-                            createParen(
-                                createFunctionExpression(
-                                    /*asteriskToken*/ undefined,
-                                    /*name*/ undefined,
-                                    [createParameter(currentNamespaceLocalName)],
-                                    transformModuleBody(node)
-                                )
-                            ),
-                            [moduleParam]
-                        )
+                setNodeEmitFlags(
+                    setOriginalNode(
+                        createStatement(
+                            createCall(
+                                createParen(
+                                    createFunctionExpression(
+                                        /*asteriskToken*/ undefined,
+                                        /*name*/ undefined,
+                                        [createParameter(currentNamespaceLocalName)],
+                                        transformModuleBody(node)
+                                    )
+                                ),
+                                [moduleParam]
+                            )
+                        ),
+                        node
                     ),
-                    node
+                    NodeEmitFlags.AdviseOnEmitNode
                 )
             );
 
@@ -2636,62 +2634,51 @@ namespace ts {
                 : getClassPrototype(node);
         }
 
-        function onBeforeEmitNode(node: Node): void {
-            previousOnBeforeEmitNode(node);
+        function isClassWithDecorators(node: Node): node is ClassDeclaration {
+            return node.kind === SyntaxKind.ClassDeclaration && node.decorators !== undefined;
+        }
 
+        function isSuperContainer(node: Node): node is SuperContainer {
             const kind = node.kind;
-            if (enabledSubstitutions & TypeScriptSubstitutionFlags.DecoratedClasses
-                && kind === SyntaxKind.ClassDeclaration
-                && node.decorators) {
-                currentDecoratedClassAliases[getOriginalNodeId(node)] = decoratedClassAliases[getOriginalNodeId(node)];
-            }
+            return kind === SyntaxKind.ClassDeclaration
+                || kind === SyntaxKind.Constructor
+                || kind === SyntaxKind.MethodDeclaration
+                || kind === SyntaxKind.GetAccessor
+                || kind === SyntaxKind.SetAccessor;
+        }
 
-            if (enabledSubstitutions & TypeScriptSubstitutionFlags.NamespaceExports
-                && (kind === SyntaxKind.ClassDeclaration
-                    || kind === SyntaxKind.Constructor
-                    || kind === SyntaxKind.MethodDeclaration
-                    || kind === SyntaxKind.GetAccessor
-                    || kind === SyntaxKind.SetAccessor)) {
+        function isTransformedModuleDeclaration(node: Node): boolean {
+            return getOriginalNode(node).kind === SyntaxKind.ModuleDeclaration;
+        }
 
-                if (!superContainerStack) {
-                    superContainerStack = [];
-                }
+        function onEmitNode(node: Node, emit: (node: Node) => void): void {
+            const savedIsEnclosedInNamespace = isEnclosedInNamespace;
+            const savedCurrentSuperContainer = currentSuperContainer;
 
-                superContainerStack.push(<SuperContainer>node);
+            // If we need support substitutions for aliases for decorated classes,
+            // we should enable it here.
+            if (enabledSubstitutions & TypeScriptSubstitutionFlags.DecoratedClasses && isClassWithDecorators(node)) {
+                currentDecoratedClassAliases[getOriginalNodeId(node)] = decoratedClassAliases[getOriginalNodeId(node)];
             }
 
-            if (enabledSubstitutions & TypeScriptSubstitutionFlags.NamespaceExports
-                && kind === SyntaxKind.ModuleDeclaration) {
-                namespaceNestLevel++;
+            // If we need to support substitutions for `super` in an async method,
+            // we should track it here.
+            if (enabledSubstitutions & TypeScriptSubstitutionFlags.AsyncMethodsWithSuper && isSuperContainer(node)) {
+                currentSuperContainer = node;
             }
-        }
-
-        function onAfterEmitNode(node: Node): void {
-            previousOnAfterEmitNode(node);
 
-            const kind = node.kind;
-            if (enabledSubstitutions & TypeScriptSubstitutionFlags.DecoratedClasses
-                && kind === SyntaxKind.ClassDeclaration
-                && node.decorators) {
-                currentDecoratedClassAliases[getOriginalNodeId(node)] = undefined;
+            if (enabledSubstitutions & TypeScriptSubstitutionFlags.NamespaceExports && isTransformedModuleDeclaration(node)) {
+                isEnclosedInNamespace = true;
             }
 
-            if (enabledSubstitutions & TypeScriptSubstitutionFlags.NamespaceExports
-                && (kind === SyntaxKind.ClassDeclaration
-                    || kind === SyntaxKind.Constructor
-                    || kind === SyntaxKind.MethodDeclaration
-                    || kind === SyntaxKind.GetAccessor
-                    || kind === SyntaxKind.SetAccessor)) {
+            previousOnEmitNode(node, emit);
 
-                if (superContainerStack) {
-                    superContainerStack.pop();
-                }
+            if (enabledSubstitutions & TypeScriptSubstitutionFlags.DecoratedClasses && isClassWithDecorators(node)) {
+                currentDecoratedClassAliases[getOriginalNodeId(node)] = undefined;
             }
 
-            if (enabledSubstitutions & TypeScriptSubstitutionFlags.NamespaceExports
-                && kind === SyntaxKind.ModuleDeclaration) {
-                namespaceNestLevel--;
-            }
+            isEnclosedInNamespace = savedIsEnclosedInNamespace;
+            currentSuperContainer = savedCurrentSuperContainer;
         }
 
         function substituteExpression(node: Expression): Expression {
@@ -2702,7 +2689,7 @@ namespace ts {
                     return substituteExpressionIdentifier(<Identifier>node);
             }
 
-            if (enabledSubstitutions & TypeScriptSubstitutionFlags.NamespaceExports) {
+            if (enabledSubstitutions & TypeScriptSubstitutionFlags.AsyncMethodsWithSuper) {
                 switch (node.kind) {
                     case SyntaxKind.CallExpression:
                         return substituteCallExpression(<CallExpression>node);
@@ -2735,8 +2722,7 @@ namespace ts {
                 }
             }
 
-            if (enabledSubstitutions & TypeScriptSubstitutionFlags.NamespaceExports
-                && namespaceNestLevel > 0) {
+            if (enabledSubstitutions & TypeScriptSubstitutionFlags.NamespaceExports && isEnclosedInNamespace) {
                 // If we are nested within a namespace declaration, we may need to qualifiy
                 // an identifier that is exported from a merged namespace.
                 const original = getOriginalNode(node);
@@ -2802,8 +2788,8 @@ namespace ts {
         }
 
         function enableExpressionSubstitutionForAsyncMethodsWithSuper() {
-            if ((enabledSubstitutions & TypeScriptSubstitutionFlags.NamespaceExports) === 0) {
-                enabledSubstitutions |= TypeScriptSubstitutionFlags.NamespaceExports;
+            if ((enabledSubstitutions & TypeScriptSubstitutionFlags.AsyncMethodsWithSuper) === 0) {
+                enabledSubstitutions |= TypeScriptSubstitutionFlags.AsyncMethodsWithSuper;
 
                 // We need to enable substitutions for call, property access, and element access
                 // if we need to rewrite super calls.
@@ -2844,9 +2830,6 @@ namespace ts {
 
                 // We need to be notified when entering and exiting namespaces.
                 context.enableEmitNotification(SyntaxKind.ModuleDeclaration);
-
-                // Keep track of namespace nesting depth
-                namespaceNestLevel = 0;
             }
         }
 
@@ -2871,9 +2854,8 @@ namespace ts {
         }
 
         function getSuperContainerAsyncMethodFlags() {
-            const container = lastOrUndefined(superContainerStack);
-            return container !== undefined
-                && resolver.getNodeCheckFlags(getOriginalNode(container)) & (NodeCheckFlags.AsyncMethodWithSuper | NodeCheckFlags.AsyncMethodWithSuperBinding);
+            return currentSuperContainer !== undefined
+                && resolver.getNodeCheckFlags(getOriginalNode(currentSuperContainer)) & (NodeCheckFlags.AsyncMethodWithSuper | NodeCheckFlags.AsyncMethodWithSuperBinding);
         }
     }
 }
diff --git a/src/compiler/types.ts b/src/compiler/types.ts
@@ -2867,16 +2867,10 @@ namespace ts {
         isEmitNotificationEnabled(node: Node): boolean;
 
         /**
-         * Hook used to notify transformers immediately before the pretty printer
-         * emits a node.
+         * Hook used to allow transformers to capture state before or after
+         * the printer emits a node.
          */
-        onBeforeEmitNode?: (node: Node) => void;
-
-        /**
-         * Hook used to notify transformers immediately after the pretty printer
-         * emits a node.
-         */
-        onAfterEmitNode?: (node: Node) => void;
+        onEmitNode?: (node: Node, emit: (node: Node) => void) => void;
     }
 
     /* @internal */