From d1dea90847cfa52d58bc7b3f91829a64fa4d36f5 Mon Sep 17 00:00:00 2001 From: Ivan Goncharov Date: Wed, 10 Aug 2022 15:04:37 +0300 Subject: [PATCH] build/npm: add TS transformation that optimise for-of on arrays Context: I wanted to try more automatic alternative to #3687 that allows us to continue using ES6 features without perfomance lost. --- resources/build-npm.ts | 2 + resources/optimise-for-of.ts | 183 +++++++++++++++++++++++++++++++++++ src/language/visitor.ts | 2 +- 3 files changed, 186 insertions(+), 1 deletion(-) create mode 100644 resources/optimise-for-of.ts diff --git a/resources/build-npm.ts b/resources/build-npm.ts index badbecc906..8fdeb57580 100644 --- a/resources/build-npm.ts +++ b/resources/build-npm.ts @@ -6,6 +6,7 @@ import * as ts from 'typescript'; import { addExtensionToImportPaths } from './add-extension-to-import-paths'; import { inlineInvariant } from './inline-invariant'; +import { optimiseForOf } from './optimise-for-of'; import { localRepoPath, readdirRecursive, @@ -55,6 +56,7 @@ tsHost.writeFile = (filepath, body) => { const tsProgram = ts.createProgram(['src/index.ts'], tsOptions, tsHost); const tsResult = tsProgram.emit(undefined, undefined, undefined, undefined, { + before: [optimiseForOf(tsProgram)], after: [addExtensionToImportPaths({ extension: '.js' }), inlineInvariant], }); assert( diff --git a/resources/optimise-for-of.ts b/resources/optimise-for-of.ts new file mode 100644 index 0000000000..34e085d775 --- /dev/null +++ b/resources/optimise-for-of.ts @@ -0,0 +1,183 @@ +import * as assert from 'node:assert'; + +import * as ts from 'typescript'; + +/** + * The following ES6 code: + * + * ``` + * for (let v of expr) { } + * ``` + * + * should be emitted as + * + * ``` + * for (let _i = 0, _a = expr; _i < _a.length; _i++) { + * let v = _a[_i]; + * } + * ``` + * + * where _a and _i are temps emitted to capture the RHS and the counter, respectively. + * When the left hand side is a let/const, the v is renamed if there is another v in scope. + * Note that all assignments to the LHS are emitted in the body, including all destructuring. + * + * Code is based on TS ES5 transpilation: + * https://github.com/microsoft/TypeScript/blob/71e852922888337ef51a0e48416034a94a6c34d9/src/compiler/transformers/es2015.ts#L2521 + */ +export function optimiseForOf(program: ts.Program) { + return (context: ts.TransformationContext) => { + const typeChecker = program.getTypeChecker(); + const { factory } = context; + + return visitSourceFile; + + function visitSourceFile(sourceFile: ts.SourceFile) { + return ts.visitNode(sourceFile, visitNode); + + function visitNode(node: ts.Node): ts.Node { + if (isArrayForOfStatement(node)) { + return convertForOfStatementForArray(node); + } + return ts.visitEachChild(node, visitNode, context); + } + + function isArrayForOfStatement(node: ts.Node): node is ts.ForOfStatement { + if (!ts.isForOfStatement(node) || node.awaitModifier != null) { + return false; + } + + const { expression } = node; + + const expressionType = typeChecker.getTypeAtLocation(expression); + + for (const subType of unionTypeParts(expressionType)) { + assert( + !(subType.flags & (ts.TypeFlags.Any | ts.TypeFlags.Unknown)), + 'Can not use any or uknown values in for-of loop: ' + + nodeLocationString(node), + ); + + if (subType.flags & ts.TypeFlags.StringLike) { + continue; + } + + const typeName = subType.getSymbol()?.getName(); + assert(typeName != null); + + if (typeName === 'Array' || typeName === 'ReadonlyArray') { + continue; + } + + return false; + } + + return true; + } + + function nodeLocationString(node: ts.Node): string { + const position = sourceFile.getLineAndCharacterOfPosition( + node.getStart(), + ); + return sourceFile.fileName + ':' + position.line; + } + + function unionTypeParts(type: ts.Type): ReadonlyArray { + return isUnionType(type) ? type.types : [type]; + } + + function isUnionType(type: ts.Type): type is ts.UnionType { + return (type.flags & ts.TypeFlags.Union) !== 0; + } + + function convertForOfStatementForArray( + forOfNode: ts.ForOfStatement, + ): ts.Statement { + const counter = factory.createLoopVariable(); + const forDeclarations = [ + factory.createVariableDeclaration( + counter, + undefined, // exclamationToken + undefined, // type + factory.createNumericLiteral(0), + ), + ]; + + // In the case where the user wrote an identifier as the RHS, like this: + // + // for (let v of arr) { } + // + // we don't want to emit a temporary variable for the RHS, just use it directly. + let rhsReference; + if (ts.isIdentifier(forOfNode.expression)) { + rhsReference = forOfNode.expression; + } else { + rhsReference = factory.createTempVariable( + undefined, // recordTempVariable + ); + forDeclarations.push( + factory.createVariableDeclaration( + rhsReference, + undefined, // exclamationToken + undefined, // type + forOfNode.expression, + ), + ); + } + + const forIntiliazer = factory.createVariableDeclarationList( + forDeclarations, + ts.NodeFlags.Let, + ); + + const forCondition = factory.createLessThan( + counter, + factory.createPropertyAccessExpression(rhsReference, 'length'), + ); + const forIncrementor = factory.createPostfixIncrement(counter); + + assert(ts.isVariableDeclarationList(forOfNode.initializer)); + // It will use rhsIterationValue _a[_i] as the initializer. + const itemAssigment = convertForOfInitializer( + forOfNode.initializer, + factory.createElementAccessExpression(rhsReference, counter), + ); + + assert(ts.isBlock(forOfNode.statement)); + const forBody = factory.updateBlock(forOfNode.statement, [ + itemAssigment, + ...forOfNode.statement.statements, + ]); + + return factory.createForStatement( + forIntiliazer, + forCondition, + forIncrementor, + forBody, + ); + } + + function convertForOfInitializer( + forOfDeclarationList: ts.VariableDeclarationList, + itemAccessExpression: ts.Expression, + ) { + assert(forOfDeclarationList.declarations.length === 1); + const [forOfDeclaration] = forOfDeclarationList.declarations; + + const updatedDeclaration = factory.updateVariableDeclaration( + forOfDeclaration, + forOfDeclaration.name, + forOfDeclaration.exclamationToken, + forOfDeclaration.type, + itemAccessExpression, + ); + + return factory.createVariableStatement( + undefined, // modifiers + factory.updateVariableDeclarationList(forOfDeclarationList, [ + updatedDeclaration, + ]), + ); + } + } + }; +} diff --git a/src/language/visitor.ts b/src/language/visitor.ts index 0d5001f34f..fbab39b3b4 100644 --- a/src/language/visitor.ts +++ b/src/language/visitor.ts @@ -189,7 +189,7 @@ export function visit( let inArray = Array.isArray(root); let keys: any = [root]; let index = -1; - let edits = []; + let edits: Array = []; let node: any = root; let key: any = undefined; let parent: any = undefined;