Improve __spreadArray perf, fix array packing and trailing omitted expressions

Author: rbuckton

src/compiler/factory/emitHelpers.ts

@@ -19,7 +19,7 @@ namespace ts {
         // ES2015 Helpers
         createExtendsHelper(name: Identifier): Expression;
         createTemplateObjectHelper(cooked: ArrayLiteralExpression, raw: ArrayLiteralExpression): Expression;
-        createSpreadArrayHelper(to: Expression, from: Expression): Expression;
+        createSpreadArrayHelper(to: Expression, from: Expression, packFrom: boolean): Expression;
         // ES2015 Destructuring Helpers
         createValuesHelper(expression: Expression): Expression;
         createReadHelper(iteratorRecord: Expression, count: number | undefined): Expression;
@@ -282,12 +282,12 @@ namespace ts {
             );
         }
 
-        function createSpreadArrayHelper(to: Expression, from: Expression) {
+        function createSpreadArrayHelper(to: Expression, from: Expression, packFrom: boolean) {
             context.requestEmitHelper(spreadArrayHelper);
             return factory.createCallExpression(
                 getUnscopedHelperName("__spreadArray"),
                 /*typeArguments*/ undefined,
-                [to, from]
+                packFrom ? [to, from, factory.createTrue()] : [to, from]
             );
         }
 
@@ -637,10 +637,14 @@ namespace ts {
         importName: "__spreadArray",
         scoped: false,
         text: `
-            var __spreadArray = (this && this.__spreadArray) || function (to, from) {
-                for (var i = 0, il = from.length, j = to.length; i < il; i++, j++)
-                    to[j] = from[i];
-                return to;
+            var __spreadArray = (this && this.__spreadArray) || function (to, from, pack) {
+                if (pack) for (var i = 0, l = from.length, ar; i < l; i++) {
+                    if (ar || !(i in from)) {
+                        if (!ar) ar = Array.prototype.slice.call(from, 0, i);
+                        ar[i] = from[i];
+                    }
+                }
+                return to.concat(ar || from);
             };`
     };
 
@@ -1001,10 +1005,10 @@ namespace ts {
             })(name => super[name], (name, value) => super[name] = value);`
     };
 
-    export function isCallToHelper(firstSegment: Expression, helperName: __String) {
+    export function isCallToHelper(firstSegment: Expression, helperName: __String): boolean {
         return isCallExpression(firstSegment)
             && isIdentifier(firstSegment.expression)
-            && (getEmitFlags(firstSegment.expression) & EmitFlags.HelperName)
+            && (getEmitFlags(firstSegment.expression) & EmitFlags.HelperName) !== 0
             && firstSegment.expression.escapedText === helperName;
     }
 }

src/compiler/factory/nodeFactory.ts

@@ -525,12 +525,25 @@ namespace ts {
                 elements = [];
             }
             else if (isNodeArray(elements)) {
-                // Ensure the transform flags have been aggregated for this NodeArray
-                if (elements.transformFlags === undefined) {
-                    aggregateChildrenFlags(elements as MutableNodeArray<T>);
+                if (hasTrailingComma === undefined || elements.hasTrailingComma === hasTrailingComma) {
+                    // Ensure the transform flags have been aggregated for this NodeArray
+                    if (elements.transformFlags === undefined) {
+                        aggregateChildrenFlags(elements as MutableNodeArray<T>);
+                    }
+                    Debug.attachNodeArrayDebugInfo(elements);
+                    return elements;
                 }
-                Debug.attachNodeArrayDebugInfo(elements);
-                return elements;
+
+                // This *was* a `NodeArray`, but the `hasTrailingComma` option differs. Recreate the
+                // array with the same elements, text range, and transform flags but with the updated
+                // value for `hasTrailingComma`
+                const array = elements.slice() as MutableNodeArray<T>;
+                array.pos = elements.pos;
+                array.end = elements.end;
+                array.hasTrailingComma = hasTrailingComma;
+                array.transformFlags = elements.transformFlags;
+                Debug.attachNodeArrayDebugInfo(array);
+                return array;
             }
 
             // Since the element list of a node array is typically created by starting with an empty array and
@@ -2184,7 +2197,12 @@ namespace ts {
         // @api
         function createArrayLiteralExpression(elements?: readonly Expression[], multiLine?: boolean) {
             const node = createBaseExpression<ArrayLiteralExpression>(SyntaxKind.ArrayLiteralExpression);
-            node.elements = parenthesizerRules().parenthesizeExpressionsOfCommaDelimitedList(createNodeArray(elements));
+            // Ensure we add a trailing comma for something like `[NumericLiteral(1), NumericLiteral(2), OmittedExpresion]` so that
+            // we end up with `[1, 2, ,]` instead of `[1, 2, ]` otherwise the `OmittedExpression` will just end up being treated like
+            // a trailing comma.
+            const lastElement = elements && lastOrUndefined(elements);
+            const elementsArray = createNodeArray(elements, lastElement && isOmittedExpression(lastElement) ? true : undefined);
+            node.elements = parenthesizerRules().parenthesizeExpressionsOfCommaDelimitedList(elementsArray);
             node.multiLine = multiLine;
             node.transformFlags |= propagateChildrenFlags(node.elements);
             return node;

src/compiler/transformers/es2015.ts

@@ -242,6 +242,21 @@ namespace ts {
         FunctionSubtreeExcludes = NewTarget | CapturedLexicalThis,
     }
 
+    const enum SpreadSegmentKind {
+        None,           // Not a spread segment
+        UnpackedSpread, // A spread segment that must be packed (i.e., converting `[...[1, , 2]]` into `[1, undefined, 2]`)
+        PackedSpread,   // A spread segment that is known to already be packed (i.e., `[...[1, 2]]` or `[...__read(a)]`)
+    }
+
+    interface SpreadSegment {
+        kind: SpreadSegmentKind;
+        expression: Expression;
+    }
+
+    function createSpreadSegment(kind: SpreadSegmentKind, expression: Expression): SpreadSegment {
+        return { kind, expression };
+    }
+
     export function transformES2015(context: TransformationContext) {
         const {
             factory,
@@ -3599,7 +3614,7 @@ namespace ts {
         function visitArrayLiteralExpression(node: ArrayLiteralExpression): Expression {
             if (some(node.elements, isSpreadElement)) {
                 // We are here because we contain a SpreadElementExpression.
-                return transformAndSpreadElements(node.elements, /*needsUniqueCopy*/ true, !!node.multiLine, /*hasTrailingComma*/ !!node.elements.hasTrailingComma);
+                return transformAndSpreadElements(node.elements, /*isArgumentList*/ false, !!node.multiLine, /*hasTrailingComma*/ !!node.elements.hasTrailingComma);
             }
             return visitEachChild(node, visitor, context);
         }
@@ -3820,7 +3835,7 @@ namespace ts {
                     resultingCall = factory.createFunctionApplyCall(
                         visitNode(target, callExpressionVisitor, isExpression),
                         node.expression.kind === SyntaxKind.SuperKeyword ? thisArg : visitNode(thisArg, visitor, isExpression),
-                        transformAndSpreadElements(node.arguments, /*needsUniqueCopy*/ false, /*multiLine*/ false, /*hasTrailingComma*/ false)
+                        transformAndSpreadElements(node.arguments, /*isArgumentList*/ true, /*multiLine*/ false, /*hasTrailingComma*/ false)
                     );
                 }
                 else {
@@ -3878,7 +3893,7 @@ namespace ts {
                     factory.createFunctionApplyCall(
                         visitNode(target, visitor, isExpression),
                         thisArg,
-                        transformAndSpreadElements(factory.createNodeArray([factory.createVoidZero(), ...node.arguments!]), /*needsUniqueCopy*/ false, /*multiLine*/ false, /*hasTrailingComma*/ false)
+                        transformAndSpreadElements(factory.createNodeArray([factory.createVoidZero(), ...node.arguments!]), /*isArgumentList*/ true, /*multiLine*/ false, /*hasTrailingComma*/ false)
                     ),
                     /*typeArguments*/ undefined,
                     []
@@ -3891,12 +3906,12 @@ namespace ts {
          * Transforms an array of Expression nodes that contains a SpreadExpression.
          *
          * @param elements The array of Expression nodes.
-         * @param needsUniqueCopy A value indicating whether to ensure that the result is a fresh array.
-         * This should be `true` when spreading into an `ArrayLiteral`, and `false` when spreading into an
+         * @param isArgumentList A value indicating whether to ensure that the result is a fresh array.
+         * This should be `false` when spreading into an `ArrayLiteral`, and `true` when spreading into an
          * argument list.
          * @param multiLine A value indicating whether the result should be emitted on multiple lines.
          */
-        function transformAndSpreadElements(elements: NodeArray<Expression>, needsUniqueCopy: boolean, multiLine: boolean, hasTrailingComma: boolean): Expression {
+        function transformAndSpreadElements(elements: NodeArray<Expression>, isArgumentList: boolean, multiLine: boolean, hasTrailingComma: boolean): Expression {
             // When there is no leading SpreadElement:
             //
             // [source]
@@ -3931,7 +3946,10 @@ namespace ts {
             // Map spans of spread expressions into their expressions and spans of other
             // expressions into an array literal.
             const numElements = elements.length;
-            const segments = flatten<Expression>(
+            const segments = flatten<SpreadSegment>(
+                // As we visit each element, we return one of two functions to use as the "key":
+                // - `visitSpanOfSpreads` for one or more contiguous `...` spread expressions, i.e. `...a, ...b` in `[1, 2, ...a, ...b]`
+                // - `visitSpanOfNonSpreads` for one or more contiguous non-spread elements, i.e. `1, 2`, in `[1, 2, ...a, ...b]`
                 spanMap(elements, partitionSpread, (partition, visitPartition, _start, end) =>
                     visitPartition(partition, multiLine, hasTrailingComma && end === numElements)
                 )
@@ -3943,26 +3961,26 @@ namespace ts {
                 // a CallExpression or NewExpression. When using `--downlevelIteration`, we need
                 // to coerce this into an array for use with `apply`, so we will use the code path
                 // that follows instead.
-                if (!needsUniqueCopy && !compilerOptions.downlevelIteration
-                    || isPackedArrayLiteral(firstSegment) // see NOTE (above)
-                    || isCallToHelper(firstSegment, "___spreadArray" as __String)) {
-                    return segments[0];
+                if (isArgumentList && !compilerOptions.downlevelIteration
+                    || isPackedArrayLiteral(firstSegment.expression) // see NOTE (above)
+                    || isCallToHelper(firstSegment.expression, "___spreadArray" as __String)) {
+                    return firstSegment.expression;
                 }
             }
 
             const helpers = emitHelpers();
-            const startsWithSpread = isSpreadElement(elements[0]);
+            const startsWithSpread = segments[0].kind !== SpreadSegmentKind.None;
             let expression: Expression =
                 startsWithSpread ? factory.createArrayLiteralExpression() :
-                segments[0];
+                segments[0].expression;
             for (let i = startsWithSpread ? 0 : 1; i < segments.length; i++) {
+                const segment = segments[i];
+                // If this is for an argument list, it doesn't matter if the array is packed or sparse
                 expression = helpers.createSpreadArrayHelper(
                     expression,
-                    compilerOptions.downlevelIteration && !isPackedArrayLiteral(segments[i]) ? // see NOTE (above)
-                        helpers.createReadHelper(segments[i], /*count*/ undefined) :
-                        segments[i]);
+                    segment.expression,
+                    segment.kind === SpreadSegmentKind.UnpackedSpread && !isArgumentList);
             }
-
             return expression;
         }
 
@@ -3972,27 +3990,38 @@ namespace ts {
                 : visitSpanOfNonSpreads;
         }
 
-        function visitSpanOfSpreads(chunk: Expression[]): VisitResult<Expression> {
+        function visitSpanOfSpreads(chunk: Expression[]): SpreadSegment[] {
             return map(chunk, visitExpressionOfSpread);
         }
 
-        function visitSpanOfNonSpreads(chunk: Expression[], multiLine: boolean, hasTrailingComma: boolean): VisitResult<Expression> {
-            return factory.createArrayLiteralExpression(
-                visitNodes(factory.createNodeArray(chunk, hasTrailingComma), visitor, isExpression),
-                multiLine
-            );
+        function visitExpressionOfSpread(node: SpreadElement): SpreadSegment {
+            let expression = visitNode(node.expression, visitor, isExpression);
+
+            // We don't need to pack already packed array literals, or existing calls to the `__read` helper.
+            const isCallToReadHelper = isCallToHelper(expression, "___read" as __String);
+            let kind = isCallToReadHelper || isPackedArrayLiteral(expression) ? SpreadSegmentKind.PackedSpread : SpreadSegmentKind.UnpackedSpread;
+
+            // We don't need the `__read` helper for array literals. Array packing will be performed by `__spreadArray`.
+            if (compilerOptions.downlevelIteration && kind === SpreadSegmentKind.UnpackedSpread && !isArrayLiteralExpression(expression) && !isCallToReadHelper) {
+                expression = emitHelpers().createReadHelper(expression, /*count*/ undefined);
+                // the `__read` helper returns a packed array, so we don't need to ensure a packed array
+                kind = SpreadSegmentKind.PackedSpread;
+            }
+
+            return createSpreadSegment(kind, expression);
         }
 
-        function visitSpreadElement(node: SpreadElement) {
-            return visitNode(node.expression, visitor, isExpression);
+        function visitSpanOfNonSpreads(chunk: Expression[], multiLine: boolean, hasTrailingComma: boolean): SpreadSegment {
+            const expression = factory.createArrayLiteralExpression(
+                visitNodes(factory.createNodeArray(chunk, hasTrailingComma), visitor, isExpression),
+                multiLine);
+
+            // We do not pack non-spread segments, this is so that `[1, , ...[2, , 3], , 4]` is properly downleveled to
+            // `[1, , 2, undefined, 3, , 4]`. See the NOTE in `transformAndSpreadElements`
+            return createSpreadSegment(SpreadSegmentKind.None, expression);
         }
 
-        /**
-         * Transforms the expression of a SpreadExpression node.
-         *
-         * @param node A SpreadExpression node.
-         */
-        function visitExpressionOfSpread(node: SpreadElement) {
+        function visitSpreadElement(node: SpreadElement) {
             return visitNode(node.expression, visitor, isExpression);
         }
 

src/compiler/types.ts

@@ -999,10 +999,13 @@ namespace ts {
         ;
 
     /* @internal */
-    export type MutableNodeArray<T extends Node> = NodeArray<T> & T[];
+    export interface MutableNodeArray<T extends Node> extends Array<T>, TextRange {
+        hasTrailingComma: boolean;
+        /* @internal */ transformFlags: TransformFlags;   // Flags for transforms, possibly undefined
+    }
 
     export interface NodeArray<T extends Node> extends ReadonlyArray<T>, ReadonlyTextRange {
-        hasTrailingComma?: boolean;
+        readonly hasTrailingComma: boolean;
         /* @internal */ transformFlags: TransformFlags;   // Flags for transforms, possibly undefined
     }
 

src/testRunner/tsconfig.json

@@ -88,6 +88,7 @@
         "unittests/config/showConfig.ts",
         "unittests/config/tsconfigParsing.ts",
         "unittests/config/tsconfigParsingWatchOptions.ts",
+        "unittests/evaluation/arraySpread.ts",
         "unittests/evaluation/asyncArrow.ts",
         "unittests/evaluation/asyncGenerator.ts",
         "unittests/evaluation/awaiter.ts",

src/testRunner/unittests/evaluation/arraySpread.ts

@@ -0,0 +1,39 @@
+describe("unittests:: evaluation:: arraySpread", () => {
+    it("array spread preserves side-effects", async () => {
+        const result = evaluator.evaluateTypeScript(`
+            const k = [1, 2];
+            const o = [3, ...k, k[0]++];
+            export const output = o;
+        `);
+        assert.deepEqual(result.output, [3, 1, 2, 1]);
+    });
+    it("array spread packs spread elements", async () => {
+        const result = evaluator.evaluateTypeScript(`
+            const k = [1, , 2];
+            const o = [3, ...k, 4];
+            export const output = o;
+        `);
+        assert.deepEqual(result.output, [3, 1, undefined, 2, 4]);
+        assert.hasAllKeys(result.output, ["0", "1", "2", "3", "4"]);
+    });
+    it("array spread does not pack non-spread elements", async () => {
+        const result = evaluator.evaluateTypeScript(`
+            const k = [1, 2];
+            const o = [3, , ...k, , 4];
+            export const output = o;
+        `);
+        assert.deepEqual(result.output, [3, , 1, 2, , 4]); // eslint-disable-line no-sparse-arrays
+        assert.hasAllKeys(result.output, ["0", "2", "3", "5"]);
+        assert.doesNotHaveAllKeys(result.output, ["1", "4"]);
+    });
+    it("argument spread pack does not matter", async () => {
+        const result = evaluator.evaluateTypeScript(`
+            const f = (...args) => args;
+            const k = [1, , 2];
+            const o = f(3, ...k, 4);
+            export const output = o;
+        `);
+        assert.deepEqual(result.output, [3, 1, undefined, 2,4]);
+        assert.hasAllKeys(result.output, ["0", "1", "2", "3", "4"]);
+    });
+});

tests/baselines/reference/api/tsserverlibrary.d.ts

@@ -574,7 +574,7 @@ declare namespace ts {
     export type HasInitializer = HasExpressionInitializer | ForStatement | ForInStatement | ForOfStatement | JsxAttribute;
     export type HasExpressionInitializer = VariableDeclaration | ParameterDeclaration | BindingElement | PropertySignature | PropertyDeclaration | PropertyAssignment | EnumMember;
     export interface NodeArray<T extends Node> extends ReadonlyArray<T>, ReadonlyTextRange {
-        hasTrailingComma?: boolean;
+        readonly hasTrailingComma: boolean;
     }
     export interface Token<TKind extends SyntaxKind> extends Node {
         readonly kind: TKind;

tests/baselines/reference/api/typescript.d.ts

@@ -574,7 +574,7 @@ declare namespace ts {
     export type HasInitializer = HasExpressionInitializer | ForStatement | ForInStatement | ForOfStatement | JsxAttribute;
     export type HasExpressionInitializer = VariableDeclaration | ParameterDeclaration | BindingElement | PropertySignature | PropertyDeclaration | PropertyAssignment | EnumMember;
     export interface NodeArray<T extends Node> extends ReadonlyArray<T>, ReadonlyTextRange {
-        hasTrailingComma?: boolean;
+        readonly hasTrailingComma: boolean;
     }
     export interface Token<TKind extends SyntaxKind> extends Node {
         readonly kind: TKind;