Trampolines for large binary expressions (microsoft#36248)

* WIP

* Test no longer crashes, but emit trampoline is incomplete and skips pipeline phases

* Fix lints, use non-generator trampoline in emit (still skips pipeline)

* Final version with emitBinaryExprssion work stack that is only used if possible

* Fix lints

* retarget to es2015 for testing

* Use bespoke state machine trampolines in binder and checker

* Remove now extraneous code in parser

* Adjust fixupParentReferences to use a depth first preorder traversal rather than breadth first

* Reintroduce incremental fast bail in fixupParentReferences

* Revert target to es5

* PR feedback

* Small edit for devops rebuild with updated definition

* Fix comment nits, add internally extraneous check back into transformer

Author: weswigham

src/compiler/binder.ts

@@ -1445,28 +1445,156 @@ namespace ts {
             }
         }
 
+        const enum BindBinaryExpressionFlowState {
+            BindThenBindChildren,
+            MaybeBindLeft,
+            BindToken,
+            BindRight,
+            FinishBind
+        }
+
         function bindBinaryExpressionFlow(node: BinaryExpression) {
-            const operator = node.operatorToken.kind;
-            if (operator === SyntaxKind.AmpersandAmpersandToken || operator === SyntaxKind.BarBarToken || operator === SyntaxKind.QuestionQuestionToken) {
-                if (isTopLevelLogicalExpression(node)) {
-                    const postExpressionLabel = createBranchLabel();
-                    bindLogicalExpression(node, postExpressionLabel, postExpressionLabel);
-                    currentFlow = finishFlowLabel(postExpressionLabel);
+            const workStacks: {
+                expr: BinaryExpression[],
+                state: BindBinaryExpressionFlowState[],
+                inStrictMode: (boolean | undefined)[],
+                parent: (Node | undefined)[],
+                subtreeFlags: (number | undefined)[]
+            } = {
+                expr: [node],
+                state: [BindBinaryExpressionFlowState.MaybeBindLeft],
+                inStrictMode: [undefined],
+                parent: [undefined],
+                subtreeFlags: [undefined]
+            };
+            let stackIndex = 0;
+            while (stackIndex >= 0) {
+                node = workStacks.expr[stackIndex];
+                switch (workStacks.state[stackIndex]) {
+                    case BindBinaryExpressionFlowState.BindThenBindChildren: {
+                        // This state is used only when recuring, to emulate the work that `bind` does before
+                        // reaching `bindChildren`. A normal call to `bindBinaryExpressionFlow` will already have done this work.
+                        node.parent = parent;
+                        const saveInStrictMode = inStrictMode;
+                        bindWorker(node);
+                        const saveParent = parent;
+                        parent = node;
+
+                        let subtreeFlagsState: number | undefined;
+                        // While this next part does the work of `bindChildren` before it descends into `bindChildrenWorker`
+                        // and uses `subtreeFlagsState` to queue up the work that needs to be done once the node is bound.
+                        if (skipTransformFlagAggregation) {
+                            // do nothing extra
+                        }
+                        else if (node.transformFlags & TransformFlags.HasComputedFlags) {
+                            skipTransformFlagAggregation = true;
+                            subtreeFlagsState = -1;
+                        }
+                        else {
+                            const savedSubtreeTransformFlags = subtreeTransformFlags;
+                            subtreeTransformFlags = 0;
+                            subtreeFlagsState = savedSubtreeTransformFlags;
+                        }
+
+                        advanceState(BindBinaryExpressionFlowState.MaybeBindLeft, saveInStrictMode, saveParent, subtreeFlagsState);
+                        break;
+                    }
+                    case BindBinaryExpressionFlowState.MaybeBindLeft: {
+                        const operator = node.operatorToken.kind;
+                        // TODO: bindLogicalExpression is recursive - if we want to handle deeply nested `&&` expressions
+                        // we'll need to handle the `bindLogicalExpression` scenarios in this state machine, too
+                        // For now, though, since the common cases are chained `+`, leaving it recursive is fine
+                        if (operator === SyntaxKind.AmpersandAmpersandToken || operator === SyntaxKind.BarBarToken || operator === SyntaxKind.QuestionQuestionToken) {
+                            if (isTopLevelLogicalExpression(node)) {
+                                const postExpressionLabel = createBranchLabel();
+                                bindLogicalExpression(node, postExpressionLabel, postExpressionLabel);
+                                currentFlow = finishFlowLabel(postExpressionLabel);
+                            }
+                            else {
+                                bindLogicalExpression(node, currentTrueTarget!, currentFalseTarget!);
+                            }
+                            completeNode();
+                        }
+                        else {
+                            advanceState(BindBinaryExpressionFlowState.BindToken);
+                            maybeBind(node.left);
+                        }
+                        break;
+                    }
+                    case BindBinaryExpressionFlowState.BindToken: {
+                        advanceState(BindBinaryExpressionFlowState.BindRight);
+                        maybeBind(node.operatorToken);
+                        break;
+                    }
+                    case BindBinaryExpressionFlowState.BindRight: {
+                        advanceState(BindBinaryExpressionFlowState.FinishBind);
+                        maybeBind(node.right);
+                        break;
+                    }
+                    case BindBinaryExpressionFlowState.FinishBind: {
+                        const operator = node.operatorToken.kind;
+                        if (isAssignmentOperator(operator) && !isAssignmentTarget(node)) {
+                            bindAssignmentTargetFlow(node.left);
+                            if (operator === SyntaxKind.EqualsToken && node.left.kind === SyntaxKind.ElementAccessExpression) {
+                                const elementAccess = <ElementAccessExpression>node.left;
+                                if (isNarrowableOperand(elementAccess.expression)) {
+                                    currentFlow = createFlowMutation(FlowFlags.ArrayMutation, currentFlow, node);
+                                }
+                            }
+                        }
+                        completeNode();
+                        break;
+                    }
+                    default: return Debug.fail(`Invalid state ${workStacks.state[stackIndex]} for bindBinaryExpressionFlow`);
                 }
-                else {
-                    bindLogicalExpression(node, currentTrueTarget!, currentFalseTarget!);
+            }
+
+            /**
+             * Note that `advanceState` sets the _current_ head state, and that `maybeBind` potentially pushes on a new
+             * head state; so `advanceState` must be called before any `maybeBind` during a state's execution.
+             */
+            function advanceState(state: BindBinaryExpressionFlowState, isInStrictMode?: boolean, parent?: Node, subtreeFlags?: number) {
+                workStacks.state[stackIndex] = state;
+                if (isInStrictMode !== undefined) {
+                    workStacks.inStrictMode[stackIndex] = isInStrictMode;
+                }
+                if (parent !== undefined) {
+                    workStacks.parent[stackIndex] = parent;
+                }
+                if (subtreeFlags !== undefined) {
+                    workStacks.subtreeFlags[stackIndex] = subtreeFlags;
                 }
             }
-            else {
-                bindEachChild(node);
-                if (isAssignmentOperator(operator) && !isAssignmentTarget(node)) {
-                    bindAssignmentTargetFlow(node.left);
-                    if (operator === SyntaxKind.EqualsToken && node.left.kind === SyntaxKind.ElementAccessExpression) {
-                        const elementAccess = <ElementAccessExpression>node.left;
-                        if (isNarrowableOperand(elementAccess.expression)) {
-                            currentFlow = createFlowMutation(FlowFlags.ArrayMutation, currentFlow, node);
-                        }
+
+            function completeNode() {
+                if (workStacks.inStrictMode[stackIndex] !== undefined) {
+                    if (workStacks.subtreeFlags[stackIndex] === -1) {
+                        skipTransformFlagAggregation = false;
+                        subtreeTransformFlags |= node.transformFlags & ~getTransformFlagsSubtreeExclusions(node.kind);
+                    }
+                    else if (workStacks.subtreeFlags[stackIndex] !== undefined) {
+                        subtreeTransformFlags = workStacks.subtreeFlags[stackIndex]! | computeTransformFlagsForNode(node, subtreeTransformFlags);
                     }
+                    inStrictMode = workStacks.inStrictMode[stackIndex]!;
+                    parent = workStacks.parent[stackIndex]!;
+                }
+                stackIndex--;
+            }
+
+            /**
+             * If `node` is a BinaryExpression, adds it to the local work stack, otherwise recursively binds it
+             */
+            function maybeBind(node: Node) {
+                if (node && isBinaryExpression(node)) {
+                    stackIndex++;
+                    workStacks.expr[stackIndex] = node;
+                    workStacks.state[stackIndex] = BindBinaryExpressionFlowState.BindThenBindChildren;
+                    workStacks.inStrictMode[stackIndex] = undefined;
+                    workStacks.parent[stackIndex] = undefined;
+                    workStacks.subtreeFlags[stackIndex] = undefined;
+                }
+                else {
+                    bind(node);
                 }
             }
         }

src/compiler/checker.ts

@@ -27136,12 +27136,89 @@ namespace ts {
             return (target.flags & TypeFlags.Nullable) !== 0 || isTypeComparableTo(source, target);
         }
 
+        const enum CheckBinaryExpressionState {
+            MaybeCheckLeft,
+            CheckRight,
+            FinishCheck
+        }
+
         function checkBinaryExpression(node: BinaryExpression, checkMode?: CheckMode) {
-            if (isInJSFile(node) && getAssignedExpandoInitializer(node)) {
-                return checkExpression(node.right, checkMode);
+            const workStacks: {
+                expr: BinaryExpression[],
+                state: CheckBinaryExpressionState[],
+                leftType: (Type | undefined)[]
+            } = {
+                expr: [node],
+                state: [CheckBinaryExpressionState.MaybeCheckLeft],
+                leftType: [undefined]
+            };
+            let stackIndex = 0;
+            let lastResult: Type | undefined;
+            while (stackIndex >= 0) {
+                node = workStacks.expr[stackIndex];
+                switch (workStacks.state[stackIndex]) {
+                    case CheckBinaryExpressionState.MaybeCheckLeft: {
+                        if (isInJSFile(node) && getAssignedExpandoInitializer(node)) {
+                            finishInvocation(checkExpression(node.right, checkMode));
+                            break;
+                        }
+                        checkGrammarNullishCoalesceWithLogicalExpression(node);
+                        const operator = node.operatorToken.kind;
+                        if (operator === SyntaxKind.EqualsToken && (node.left.kind === SyntaxKind.ObjectLiteralExpression || node.left.kind === SyntaxKind.ArrayLiteralExpression)) {
+                            finishInvocation(checkDestructuringAssignment(node.left, checkExpression(node.right, checkMode), checkMode, node.right.kind === SyntaxKind.ThisKeyword));
+                            break;
+                        }
+                        advanceState(CheckBinaryExpressionState.CheckRight);
+                        maybeCheckExpression(node.left);
+                        break;
+                    }
+                    case CheckBinaryExpressionState.CheckRight: {
+                        const leftType = lastResult!;
+                        workStacks.leftType[stackIndex] = leftType;
+                        const operator = node.operatorToken.kind;
+                        if (operator === SyntaxKind.AmpersandAmpersandToken || operator === SyntaxKind.BarBarToken || operator === SyntaxKind.QuestionQuestionToken) {
+                            checkTruthinessOfType(leftType, node.left);
+                        }
+                        advanceState(CheckBinaryExpressionState.FinishCheck);
+                        maybeCheckExpression(node.right);
+                        break;
+                    }
+                    case CheckBinaryExpressionState.FinishCheck: {
+                        const leftType = workStacks.leftType[stackIndex]!;
+                        const rightType = lastResult!;
+                        finishInvocation(checkBinaryLikeExpressionWorker(node.left, node.operatorToken, node.right, leftType, rightType, node));
+                        break;
+                    }
+                    default: return Debug.fail(`Invalid state ${workStacks.state[stackIndex]} for checkBinaryExpression`);
+                }
+            }
+
+            return lastResult!;
+
+            function finishInvocation(result: Type) {
+                lastResult = result;
+                stackIndex--;
+            }
+
+            /**
+             * Note that `advanceState` sets the _current_ head state, and that `maybeCheckExpression` potentially pushes on a new
+             * head state; so `advanceState` must be called before any `maybeCheckExpression` during a state's execution.
+             */
+            function advanceState(nextState: CheckBinaryExpressionState) {
+                workStacks.state[stackIndex] = nextState;
+            }
+
+            function maybeCheckExpression(node: Expression) {
+                if (isBinaryExpression(node)) {
+                    stackIndex++;
+                    workStacks.expr[stackIndex] = node;
+                    workStacks.state[stackIndex] = CheckBinaryExpressionState.MaybeCheckLeft;
+                    workStacks.leftType[stackIndex] = undefined;
+                }
+                else {
+                    lastResult = checkExpression(node, checkMode);
+                }
             }
-            checkGrammarNullishCoalesceWithLogicalExpression(node);
-            return checkBinaryLikeExpression(node.left, node.operatorToken, node.right, checkMode, node);
         }
 
         function checkGrammarNullishCoalesceWithLogicalExpression(node: BinaryExpression) {
@@ -27156,6 +27233,8 @@ namespace ts {
             }
         }
 
+        // Note that this and `checkBinaryExpression` above should behave mostly the same, except this elides some
+        // expression-wide checks and does not use a work stack to fold nested binary expressions into the same callstack frame
         function checkBinaryLikeExpression(left: Expression, operatorToken: Node, right: Expression, checkMode?: CheckMode, errorNode?: Node): Type {
             const operator = operatorToken.kind;
             if (operator === SyntaxKind.EqualsToken && (left.kind === SyntaxKind.ObjectLiteralExpression || left.kind === SyntaxKind.ArrayLiteralExpression)) {
@@ -27169,7 +27248,19 @@ namespace ts {
                 leftType = checkExpression(left, checkMode);
             }
 
-            let rightType = checkExpression(right, checkMode);
+            const rightType = checkExpression(right, checkMode);
+            return checkBinaryLikeExpressionWorker(left, operatorToken, right, leftType, rightType, errorNode);
+        }
+
+        function checkBinaryLikeExpressionWorker(
+            left: Expression,
+            operatorToken: Node,
+            right: Expression,
+            leftType: Type,
+            rightType: Type,
+            errorNode?: Node
+        ): Type {
+            const operator = operatorToken.kind;
             switch (operator) {
                 case SyntaxKind.AsteriskToken:
                 case SyntaxKind.AsteriskAsteriskToken:
@@ -30774,14 +30865,17 @@ namespace ts {
             checkSourceElement(node.statement);
         }
 
-        function checkTruthinessExpression(node: Expression, checkMode?: CheckMode) {
-            const type = checkExpression(node, checkMode);
+        function checkTruthinessOfType(type: Type, node: Node) {
             if (type.flags & TypeFlags.Void) {
                 error(node, Diagnostics.An_expression_of_type_void_cannot_be_tested_for_truthiness);
             }
             return type;
         }
 
+        function checkTruthinessExpression(node: Expression, checkMode?: CheckMode) {
+            return checkTruthinessOfType(checkExpression(node, checkMode), node);
+        }
+
         function checkForStatement(node: ForStatement) {
             // Grammar checking
             if (!checkGrammarStatementInAmbientContext(node)) {