idlize

1
/*
2
 * Copyright (c) 2022-2024 Huawei Device Co., Ltd.
3
 * Licensed under the Apache License, Version 2.0 (the "License");
4
 * you may not use this file except in compliance with the License.
5
 * You may obtain a copy of the License at
6
 *
7
 * http://www.apache.org/licenses/LICENSE-2.0
8
 *
9
 * Unless required by applicable law or agreed to in writing, software
10
 * distributed under the License is distributed on an "AS IS" BASIS,
11
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
 * See the License for the specific language governing permissions and
13
 * limitations under the License.
14
 */
15

16
import * as ts from 'typescript';
17
import * as fs from "fs"
18
import { UniqueId } from "@koalaui/common"
19
import { Rewrite } from './transformation-context';
20

21

22
export enum FunctionKind {
23
    REGULAR,
24
    MEMO,
25
    MEMO_INTRINSIC,
26
}
27

28
export type FunctionTable = Map<ts.SignatureDeclarationBase, FunctionKind>
29
export type CallTable = Map<ts.CallExpression, FunctionKind>
30
export type EntryTable = Set<ts.CallExpression>
31
export type VariableTable = Map<ts.VariableLikeDeclaration, FunctionKind>
32

33
export function isNamedDeclaration(node: ts.Node): node is ts.NamedDeclaration {
34
    return ("name" in node )
35
}
36

37
export function asString(node: ts.Node|undefined): string {
38
    if (node === undefined) return "undefined node"
39
    if (ts.isIdentifier(node)) return ts.idText(node)
40
    if (isNamedDeclaration(node)) {
41
        if (node.name === undefined) {
42
            return `${ts.SyntaxKind[node.kind]}(undefined name)`
43
        } else {
44

45
            return `${ts.SyntaxKind[node.kind]}(${asString(node.name)})`
46
        }
47
    } else {
48
        return `${ts.SyntaxKind[node.kind]}`
49
    }
50
}
51

52
export function isFunctionOrMethod(node: ts.Node): node is ts.FunctionLikeDeclaration {
53
    return ts.isFunctionDeclaration(node) ||
54
           ts.isMethodDeclaration(node) ||
55
           ts.isFunctionExpression(node) ||
56
           ts.isArrowFunction(node)
57
}
58

59
export enum RuntimeNames {
60
    COMPUTE = "compute",
61
    CONTEXT = "__memo_context",
62
    ID = "__memo_id",
63
    SCOPE = "__memo_scope",
64
    INTERNAL_PARAMETER_STATE = "param",
65
    INTERNAL_VALUE = "cached",
66
    INTERNAL_VALUE_NEW = "recache",
67
    INTERNAL_SCOPE = "scope",
68
    INTERNAL_VALUE_OK = "unchanged",
69
    CONTENT = "content",
70
    VALUE = "value",
71
    __CONTEXT = "__context",
72
    __ID = "__id",
73
    __KEY = "__key",
74
    __STATE = "__state",
75
    CONTEXT_TYPE = "__memo_context_type",
76
    ID_TYPE = "__memo_id_type",
77
    TRANSFORMED_TYPE = "__memo_transformed",
78
    SYNTHETIC_RETURN_MARK = "__synthetic_return_value",
79
    CONTEXT_TYPE_DEFAULT_IMPORT = "@koalaui/runtime",
80
    ANNOTATION = "@memo",
81
    ANNOTATION_INTRINSIC = "@memo:intrinsic",
82
    ANNOTATION_ENTRY = "@memo:entry",
83
    ANNOTATION_SKIP = "@skip:memo", // skip binding to parameter changes
84
    ANNOTATION_STABLE = "@memo:stable", // assume this should not be tracked
85
}
86

87
export function Undefined(): ts.Identifier {
88
    return ts.factory.createIdentifier("undefined")
89
}
90

91
export function runtimeIdentifier(name: RuntimeNames): ts.Identifier {
92
    return ts.factory.createIdentifier(name)
93
}
94

95
export function componentTypeName(functionName: ts.PropertyName|undefined): ts.Identifier {
96
    if (!functionName || !ts.isIdentifier(functionName) && !ts.isPrivateIdentifier(functionName)) {
97
        throw new Error("Expected an identifier: " + asString(functionName))
98
    }
99
    return ts.factory.createIdentifier(ts.idText(functionName)+"Node")
100
}
101

102
export function isSpecialContentParameter(param: ts.ParameterDeclaration): boolean {
103
    if (!param.type) return false
104
    if (!param.name) return false
105
    return (
106
        ts.isFunctionTypeNode(param.type) &&
107
        ts.isIdentifier(param.name) &&
108
        ts.idText(param.name) == RuntimeNames.CONTENT
109
    )
110
}
111

112
export function isSpecialContextParameter(parameter: ts.ParameterDeclaration): boolean {
113
    return !!parameter.name && ts.isIdentifier(parameter.name) && ts.idText(parameter.name) == RuntimeNames.CONTEXT
114
}
115

116
export function isSpecialIdParameter(parameter: ts.ParameterDeclaration): boolean {
117
    return !!parameter.name && ts.isIdentifier(parameter.name) && ts.idText(parameter.name) == RuntimeNames.ID
118
}
119

120
export function isSkippedParameter(sourceFile: ts.SourceFile, parameter: ts.ParameterDeclaration): boolean {
121
    return getComment(sourceFile, parameter).includes(RuntimeNames.ANNOTATION_SKIP)
122
}
123

124
export function isStableClass(sourceFile: ts.SourceFile, clazz: ts.ClassDeclaration): boolean {
125
    return getComment(sourceFile, clazz).includes(RuntimeNames.ANNOTATION_STABLE)
126
}
127

128
export function hasMemoEntry(sourceFile: ts.SourceFile, node: ts.Node): boolean {
129
    const comment = getComment(sourceFile, node)
130
    return comment.includes(RuntimeNames.ANNOTATION_ENTRY)
131
}
132

133
export function isMemoEntry(sourceFile: ts.SourceFile, func: ts.FunctionDeclaration): boolean {
134
    const name = func.name
135
    if (!name) return false
136
    if (!ts.isIdentifier(name)) return false
137

138
    return hasMemoEntry(sourceFile, func)
139
}
140

141
export function isTrackableParameter(sourceFile: ts.SourceFile, parameter: ts.ParameterDeclaration): boolean {
142
    return !isSpecialContentParameter(parameter)
143
        && !isSpecialContextParameter(parameter)
144
        && !isSpecialIdParameter(parameter)
145
        && !isSkippedParameter(sourceFile, parameter)
146
}
147

148
export function parameterStateName(original: string): string {
149
    return `__memo_parameter_${original}`
150
}
151

152
export function getSymbolByNode(typechecker: ts.TypeChecker, node: ts.Node) : ts.Symbol|undefined {
153
      return typechecker.getSymbolAtLocation(node)
154
}
155

156
export function getDeclarationsByNode(typechecker: ts.TypeChecker, node: ts.Node) : ts.Declaration[] {
157
    const symbol = getSymbolByNode(typechecker, node)
158
    const declarations =  symbol?.getDeclarations() ?? []
159
    return declarations
160
}
161

162
export function findSourceFile(node: ts.Node): ts.SourceFile|undefined {
163
    let element = node
164
    while(element) {
165
        if (ts.isSourceFile(element)) return element
166
        element = element.parent
167
    }
168
    return undefined
169
}
170

171
export function findFunctionDeclaration(node: ts.Node): ts.FunctionDeclaration|undefined {
172
    let element = node
173
    while(element) {
174
        if (ts.isFunctionDeclaration(element)) return element
175
        element = element.parent
176
    }
177
    return undefined
178
}
179

180
export function isMethodOfStableClass(sourceFile: ts.SourceFile, method: ts.MethodDeclaration): boolean {
181
    const original = ts.getOriginalNode(method)
182
    const parent = original.parent
183
    if (!parent) return false
184
    if (!ts.isClassDeclaration(parent)) return false
185
    return isStableClass(sourceFile, parent)
186
}
187

188
export function arrayAt<T>(array: T[] | undefined, index: number): T|undefined {
189
    return array ? array[index >= 0 ? index : array.length + index] : undefined
190
}
191

192
export function getComment(sourceFile: ts.SourceFile, node: ts.Node): string {
193
    const commentRanges = ts.getLeadingCommentRanges(
194
        sourceFile.getFullText(),
195
        node.getFullStart()
196
    )
197

198
    const commentRange = arrayAt(commentRanges, -1)
199
    if (!commentRange) return ""
200

201
    const comment = sourceFile.getFullText()
202
        .slice(commentRange.pos, commentRange.end)
203
    return comment
204
}
205

206
export function isVoidOrNotSpecified(type: ts.TypeNode | undefined): boolean {
207
    return type === undefined
208
        || type.kind === ts.SyntaxKind.VoidKeyword
209
        || ts.isTypeReferenceNode(type)
210
        && ts.isIdentifier(type.typeName)
211
        && ts.idText(type.typeName) == "void"
212
}
213

214
export class PositionalIdTracker {
215
    // Global for the whole program.
216
    static callCount: number = 0
217

218
    // Set `stable` to true if you want to have more predictable values.
219
    // For example for tests.
220
    // Don't use it in production!
221
    constructor(public sourceFile: ts.SourceFile, public stableForTests: boolean = false) {
222
        if (stableForTests) PositionalIdTracker.callCount = 0
223
    }
224

225
    sha1Id(callName: string, fileName: string): string {
226
        const uniqId = new UniqueId()
227
        uniqId.addString("memo call uniqid")
228
        uniqId.addString(fileName)
229
        uniqId.addString(callName)
230
        uniqId.addI32(PositionalIdTracker.callCount++)
231
        return uniqId.compute().substring(0,10)
232
    }
233

234
    stringId(callName: string, fileName: string): string {
235
        return `${PositionalIdTracker.callCount++}_${callName}_id_DIRNAME/${fileName}`
236
    }
237

238
    id(callName: string): ts.Expression {
239

240
        const fileName = this.stableForTests ?
241
            baseName(this.sourceFile.fileName) :
242
            this.sourceFile.fileName
243

244
        const positionId = (this.stableForTests) ?
245
            this.stringId(callName, fileName) :
246
            this.sha1Id(callName, fileName)
247

248

249
        return this.stableForTests ?
250
            ts.factory.createStringLiteral(positionId) :
251
            ts.factory.createNumericLiteral("0x"+positionId)
252

253
    }
254
}
255

256
export function wrapInCompute(node: ts.ConciseBody, id: ts.Expression): ts.Expression {
257
    return ts.factory.createCallExpression(
258
        ts.factory.createPropertyAccessExpression(
259
            runtimeIdentifier(RuntimeNames.CONTEXT),
260
            runtimeIdentifier(RuntimeNames.COMPUTE)
261
        ),
262
        /*typeArguments*/ undefined,
263
        [
264
            id,
265
            ts.factory.createArrowFunction(
266
                /*modifiers*/ undefined,
267
                /*typeParameters*/ undefined,
268
                /*parameters*/ [],
269
                /*type*/ undefined,
270
                ts.factory.createToken(ts.SyntaxKind.EqualsGreaterThanToken),
271
                node
272
            )
273
        ]
274
    )
275
}
276

277
export function createComputeScope(pseudoStateCount: number, id: ts.Expression, typeArgument: ts.TypeNode | undefined): ts.VariableStatement {
278
    return constVariable(RuntimeNames.SCOPE,
279
        ts.factory.createCallExpression(
280
            ts.factory.createPropertyAccessExpression(
281
                runtimeIdentifier(RuntimeNames.CONTEXT),
282
                runtimeIdentifier(RuntimeNames.INTERNAL_SCOPE)
283
            ),
284
            typeArgument ? [typeArgument] : undefined,
285
            pseudoStateCount > 0
286
                // TODO: these are because es2panda is not goot at handling default values in complex signatures
287
                ? [id, ts.factory.createNumericLiteral(pseudoStateCount), Undefined(), Undefined(), Undefined(), Undefined()]
288
                : [id, Undefined(), Undefined(), Undefined(), Undefined(), Undefined()]
289
        )
290
    )
291
}
292

293
export function constVariable(name: string, initializer: ts.Expression): ts.VariableStatement {
294
    return ts.factory.createVariableStatement(
295
        undefined,
296
        ts.factory.createVariableDeclarationList(
297
            [ts.factory.createVariableDeclaration(
298
                ts.factory.createIdentifier(name),
299
                undefined,
300
                undefined,
301
                initializer
302
            )],
303
            ts.NodeFlags.Const
304
        )
305
    )
306
}
307

308
export function idPlusKey(positionalIdTracker: PositionalIdTracker): ts.Expression {
309
    return ts.factory.createBinaryExpression(
310
        runtimeIdentifier(RuntimeNames.ID),
311
        ts.factory.createToken(ts.SyntaxKind.PlusToken),
312
        ts.factory.createAsExpression(
313
            positionalIdTracker.id(RuntimeNames.__KEY),
314
            ts.factory.createTypeReferenceNode(RuntimeNames.ID_TYPE)
315
        )
316
    )
317
}
318

319
export function isAnyMemoKind(kind: FunctionKind|undefined): boolean {
320
    switch(kind) {
321
        case FunctionKind.MEMO:
322
        case FunctionKind.MEMO_INTRINSIC:
323
            return true
324
    }
325
    return false
326
}
327

328
export function isMemoKind(kind: FunctionKind|undefined): boolean {
329
    switch(kind) {
330
        case FunctionKind.MEMO:
331
            return true
332
    }
333
    return false
334
}
335

336
export function createContextType(): ts.TypeNode {
337
    return ts.factory.createTypeReferenceNode(
338
        runtimeIdentifier(RuntimeNames.CONTEXT_TYPE),
339
        undefined
340
    )
341
}
342

343
export function createIdType(): ts.TypeNode {
344
    return ts.factory.createTypeReferenceNode(
345
        runtimeIdentifier(RuntimeNames.ID_TYPE),
346
        undefined
347
    )
348
}
349

350
export function createHiddenParameters(): ts.ParameterDeclaration[] {
351
    const context = ts.factory.createParameterDeclaration(
352
        /*modifiers*/ undefined,
353
        /*dotDotDotToken*/ undefined,
354
        RuntimeNames.CONTEXT,
355
        /* questionToken */ undefined,
356
        createContextType(),
357
        /* initializer */ undefined
358

359
    )
360
    const id = ts.factory.createParameterDeclaration(
361
        /*modifiers*/ undefined,
362
        /*dotDotDotToken*/ undefined,
363
        RuntimeNames.ID,
364
        /* questionToken */ undefined,
365
        createIdType(),
366
        /* initializer */ undefined
367
    )
368

369
    return [context, id]
370
}
371

372
export function createContextTypeImport(importSource: string): ts.Statement {
373
    return ts.factory.createImportDeclaration(
374
        undefined,
375
        ts.factory.createImportClause(
376
          false,
377
          undefined,
378
            ts.factory.createNamedImports(
379
                [
380
                    ts.factory.createImportSpecifier(
381
                        false,
382
                        undefined,
383
                        ts.factory.createIdentifier(RuntimeNames.CONTEXT_TYPE)
384
                    ),
385
                    ts.factory.createImportSpecifier(
386
                        false,
387
                        undefined,
388
                        ts.factory.createIdentifier(RuntimeNames.ID_TYPE)
389
                    )
390
                ]
391
            )
392
        ),
393
        ts.factory.createStringLiteral(importSource),
394
        undefined
395
      )
396
}
397

398
export function setNeedTypeImports(rewrite: Rewrite) {
399
    rewrite.importTypesFrom = rewrite.pluginOptions.contextImport ?? RuntimeNames.CONTEXT_TYPE_DEFAULT_IMPORT
400
}
401

402
export function hiddenParameters(rewriter: Rewrite): ts.ParameterDeclaration[] {
403
    setNeedTypeImports(rewriter)
404
    return createHiddenParameters()
405
}
406

407
export function skipParenthesizedType(type: ts.TypeNode|undefined): ts.TypeNode|undefined {
408
    let current: ts.TypeNode|undefined = type
409
    while (current && ts.isParenthesizedTypeNode(current)) {
410
        current = current.type
411
    }
412
    return current
413
}
414

415
export function localStateStatement(
416
    stateName: string,
417
    referencedEntity: ts.Expression,
418
    parameterIndex: number
419
): ts.Statement {
420
    return ts.factory.createVariableStatement(
421
        undefined,
422
        ts.factory.createVariableDeclarationList(
423
            [
424
                ts.factory.createVariableDeclaration(
425
                    parameterStateName(stateName),
426
                    undefined,
427
                    undefined,
428
                    ts.factory.createCallExpression(
429
                        ts.factory.createPropertyAccessExpression(
430
                            runtimeIdentifier(RuntimeNames.SCOPE),
431
                            runtimeIdentifier(RuntimeNames.INTERNAL_PARAMETER_STATE)
432
                        ),
433
                        undefined,
434
                        // TODO: these are because es2panda is not goot at handling default values in complex signatures
435
                        [ts.factory.createNumericLiteral(parameterIndex), referencedEntity, Undefined(), Undefined(), Undefined()]
436
                    )
437
                )
438
            ],
439
            ts.NodeFlags.Const
440
        )
441
    )
442
}
443

444
export function hasStaticModifier(node: ts.MethodDeclaration): boolean {
445
    return node.modifiers?.find(it => it.kind == ts.SyntaxKind.StaticKeyword) != undefined
446
}
447

448
export function error(message: any): any {
449
    console.log(message)
450
    console.trace()
451
    return undefined
452
}
453

454
export interface TransformerOptions {
455
    // Emit transformed functions to the console.
456
    trace?: boolean,
457
    // Store the transformed functions to this directory.
458
    keepTransformed?: string,
459
    // Use human readable callsite IDs without directory paths.
460
    stableForTest?: boolean,
461
    // Import context and id types from alternative source
462
    contextImport?: string,
463
    // Dump sources with resolved memo annotations to unmemoized directory
464
    only_unmemoize?: boolean,
465
}
466

467
function baseName(path: string): string {
468
    return path.replace(/^.*\/(.*)$/, "$1")
469
}
470

471
export class Tracer {
472
    constructor (public options: TransformerOptions, public printer: ts.Printer) {
473
    }
474

475
    trace(msg: any) {
476
        if (!this.options.trace) return
477
        console.log(msg)
478
    }
479

480
    writeTextToFile(text: string, file: string) {
481
        fs.writeFileSync(file, text, 'utf8')
482
        this.trace("DUMP TO: " + file)
483
    }
484

485
    createDirs(dirs: string) {
486
        fs.mkdirSync(dirs, { recursive: true });
487
    }
488

489
    dumpFileName(sourceFile: ts.SourceFile, transformed: ts.FunctionLikeDeclarationBase): string|undefined {
490
        if (!this.options.keepTransformed) return undefined
491

492
        const outDir = (this.options.keepTransformed[0] == "/") ?
493
            this.options.keepTransformed :
494
            `${__dirname}/${this.options.keepTransformed}`
495

496
        this.createDirs(outDir)
497

498
        const sourceBaseName = baseName(sourceFile.fileName)
499
        if (!transformed.name) return
500
        if (!ts.isIdentifier(transformed.name)) return
501
        const fileName = `${ts.idText(transformed.name)}_${sourceBaseName}`
502
        return `${outDir}/${fileName}_dump`
503
    }
504

505
    keepTransformedFunction(transformed: ts.FunctionLikeDeclarationBase, sourceFile: ts.SourceFile) {
506
        const fileName = this.dumpFileName(sourceFile, transformed)
507
        if (!fileName) return
508

509
        const content = this.printer.printNode(ts.EmitHint.Unspecified, transformed, sourceFile)
510
        this.writeTextToFile(content+"\n", fileName)
511
    }
512
}
513