1
// Copyright (c) 2022-2023 Cloudflare, Inc.
2
// Licensed under the Apache 2.0 license found in the LICENSE file or at:
3
//     https://opensource.org/licenses/Apache-2.0
4
 
5
import assert from "node:assert";
6
import ts from "typescript";
7
import { isUnsatisfiable } from "../../generator/type";
8
import { printNode } from "../../print";
9
import { ensureStatementModifiers, hasModifier } from "../helpers";
10
import { maybeGetDefines, maybeGetOverride } from "./compiler";
11
 
12
export { compileOverridesDefines } from "./compiler";
13
 
14
// Applies handwritten partial TypeScript overrides to generate types to improve
15
// output fidelity. Also applies type renames and inserts additional handwritten
16
// definitions for non-generated types.
17
//
18
// See the `JSG_TS_OVERRIDE` macro's documentation in `src/workerd/jsg/jsg.h`
19
// for a full explanation of override rules and examples.
20
//
21
// `compileOverridesDefines()` must be used to compile overrides and defines
22
// into valid TypeScript source files. These should be included in the *same*
23
// TypeScript `Program` as the source file being transformed. This `Program`,
24
// along with the set of types that should be fully-replaced (also returned from
25
// `compileOverridesDefines()`), should be passed to this transformer factory.
26
//
27
// ```ts
28
// export declare class A {
29
//   thing: string;
30
// }
31
// export declare class B {
32
//   get(key: string, type: string): Promise<any>;
33
//   put(key: string, value: string): Promise<void>;
34
// }
35
// ```
36
//
37
// ...with the following overrides and defines...
38
//
39
// - `A`'s override: `RenamedA<Type extends string> { thing: Type; }`
40
// - `B`'s override: `{
41
//     get(key: string, type: "text"): Promise<string | null>;
42
//     get(key: string, type: "arrayBuffer"): Promise<ArrayBuffer | null>;
43
//   }`
44
// - `B`'s define: `interface C { foo: A; }`
45
//
46
// --- transforms to --->
47
//
48
// ```ts
49
// export declare class RenamedA<Type extends string> {
50
//   thing: Type;
51
// }
52
// export interface C {
53
//   foo: RenamedA;
54
// }
55
// export declare class B {
56
//   get(key: string, type: "text"): Promise<string | null>;
57
//   get(key: string, type: "arrayBuffer"): Promise<ArrayBuffer | null>;
58
//   put(key: string, value: string): Promise<void>;
59
// }
60
// ```
61
export function createOverrideDefineTransformer(
62
  program: ts.Program,
63
  replacements: Set<string>
64
): ts.TransformerFactory<ts.SourceFile> {
65
  return (ctx) => {
66
    return (node) => {
67
      const overrideCtx: OverrideTransformContext = {
68
        program,
69
        replacements,
70
        renames: new Map<string, string>(),
71
      };
72
      const v1 = createOverrideDefineVisitor(ctx, overrideCtx);
73
      const v2 = createRenameVisitor(ctx, overrideCtx.renames);
74
      node = ts.visitEachChild(node, v1, ctx);
75
      return ts.visitEachChild(node, v2, ctx);
76
    };
77
  };
78
}
79
 
80
interface OverrideTransformContext {
81
  program: ts.Program;
82
  replacements: Set<string>;
83
  renames: Map</* from */ string, /* to */ string>;
84
}
85
 
86
// Gets an identifying label for this member, shared between method overloads
87
function getMemberKey(member: ts.ClassElement | ts.TypeElement): string {
88
  if (ts.isConstructorDeclaration(member)) return "constructor$";
89
 
90
  const name = member.name;
91
  assert(
92
    name !== undefined,
93
    `Expected named member, got "${printNode(member)}"`
94
  );
95
 
96
  // Put static and instance members in different namespaces. For example, this
97
  // allows instance methods to be overridden without affecting static methods
98
  // of the same name.
99
  const isStatic =
100
    ts.canHaveModifiers(member) &&
101
    hasModifier(ts.getModifiers(member), ts.SyntaxKind.StaticKeyword);
102
  const keyNamespace = isStatic ? "static$" : "instance$";
103
 
104
  if (
105
    ts.isIdentifier(name) ||
106
    ts.isStringLiteral(name) ||
107
    ts.isNumericLiteral(name)
108
  ) {
109
    return keyNamespace + name.text;
110
  }
111
  if (ts.isComputedPropertyName(name)) {
112
    const expression = name.expression;
113
    if (ts.isStringLiteral(expression) || ts.isNumericLiteral(expression)) {
114
      return keyNamespace + expression.text;
115
    }
116
  }
117
  return keyNamespace + printNode(name);
118
}
119
 
120
// Groups override members by their identifying labels from `getMemberKey()`
121
function groupMembersByKey<Member extends ts.ClassElement | ts.TypeElement>(
122
  members: ts.NodeArray<Member>
123
): Map<string, Member[]> {
124
  const result = new Map<string, Member[]>();
125
  members.forEach((member) => {
126
    const key = getMemberKey(member);
127
    let array = result.get(key);
128
    if (array === undefined) result.set(key, (array = []));
129
    array.push(member);
130
  });
131
  return result;
132
}
133
 
134
// Returns the index of a member in `members` with the specified `key`, or -1 if
135
// none exists
136
// eslint-disable-next-line @typescript-eslint/no-unnecessary-type-parameters
137
function findMemberIndex<Member extends ts.ClassElement | ts.TypeElement>(
138
  members: Member[],
139
  key: string,
140
  fromIndex = 0
141
): number {
142
  return members.findIndex(
143
    (member, index) => fromIndex <= index && getMemberKey(member) === key
144
  );
145
}
146
 
147
// Merges generated members with overrides according to the following rules:
148
// 1. Members in the override but not in the generated type are inserted
149
// 2. If an override has the same key as a member in the generated type, the
150
//   generated member is removed, and the override is inserted instead
151
// 3. If an override member property is declared type `never`, it is not
152
//   inserted, but its presence may remove the generated member (as per 2)
153
function mergeMembers<Member extends ts.ClassElement | ts.TypeElement>(
154
  generated: ts.NodeArray<Member>,
155
  overrides: ts.NodeArray<ts.ClassElement>,
156
  transformer: (member: ts.ClassElement) => Member
157
): Member[] {
158
  const result = [...generated];
159
  const grouped = groupMembersByKey(overrides);
160
  for (const [key, overrideMembers] of grouped) {
161
    const filteredOverrideMembers = overrideMembers.filter((member) => {
162
      // Filter out `never` typed properties
163
      if (ts.isPropertyDeclaration(member) && member.type !== undefined) {
164
        return !isUnsatisfiable(member.type);
165
      }
166
      // Include all other members
167
      return true;
168
    });
169
    // Transform all class elements into the correct member type. If `Member` is
170
    // `ts.ClassElement` already, `transformer` will be the identify function.
171
    const transformedOverrideMembers = filteredOverrideMembers.map(transformer);
172
 
173
    // Try to find index of existing generated member with same key
174
    const index = findMemberIndex(result, key);
175
    if (index === -1) {
176
      // If the member couldn't be found, insert overrides at the end
177
      result.push(...transformedOverrideMembers);
178
    } else {
179
      const member = result[index];
180
      const nextIndex = findMemberIndex(result, key, index + 1);
181
      if (
182
        ts.isGetAccessorDeclaration(member) ||
183
        ts.isSetAccessorDeclaration(member)
184
      ) {
185
        // If this is a getter/setter, it's possible there's one other
186
        // getter/setter with the same key.
187
        if (nextIndex !== -1) {
188
          // Make sure this other member was a getter/setter
189
          const nextMember = result[nextIndex];
190
          assert(
191
            ts.isGetAccessorDeclaration(nextMember) ||
192
              ts.isSetAccessorDeclaration(nextMember),
193
            `Expected getter/setter, got "${printNode(nextMember)}"`
194
          );
195
 
196
          // Remove the other getter/setter. Because `nextIndex > index`, we'll
197
          // still be able to `splice(index)` later on.
198
          assert(nextIndex > index);
199
          result.splice(nextIndex, /* deleteCount */ 1);
200
 
201
          // Make sure this was the only other member with this key
202
          const nextNextIndex = findMemberIndex(result, key, nextIndex + 1);
203
          assert(nextNextIndex === -1);
204
        }
205
      } else {
206
        // Otherwise, make sure this was the only generated member with this key
207
        assert(nextIndex === -1);
208
      }
209
 
210
      // Remove the member at that index and replace it with overrides
211
      result.splice(index, /* deleteCount */ 1, ...transformedOverrideMembers);
212
    }
213
  }
214
  return result;
215
}
216
 
217
// Converts class members to interface members where possible. Used as a
218
// transformer when merging override members (which will always be class
219
// members) into an interface.
220
function classToTypeElement(
221
  ctx: ts.TransformationContext,
222
  member: ts.ClassElement
223
): ts.TypeElement {
224
  if (ts.isMethodDeclaration(member)) {
225
    return ctx.factory.createMethodSignature(
226
      ts.getModifiers(member),
227
      member.name,
228
      member.questionToken,
229
      member.typeParameters,
230
      member.parameters,
231
      member.type
232
    );
233
  }
234
  if (ts.isPropertyDeclaration(member)) {
235
    return ctx.factory.createPropertySignature(
236
      ts.getModifiers(member),
237
      member.name,
238
      member.questionToken,
239
      member.type
240
    );
241
  }
242
  if (
243
    ts.isGetAccessorDeclaration(member) ||
244
    ts.isSetAccessorDeclaration(member) ||
245
    ts.isIndexSignatureDeclaration(member)
246
  ) {
247
    return member;
248
  }
249
  assert.fail(
250
    `Expected interface-compatible member, got "${printNode(member)}".
251
You'll need to define a full-replacement override to a "class" if you wish to insert this member (i.e. "JSG_TS_OVERRIDE(class MyClass { <all_members> })").`
252
  );
253
}
254
 
255
// Finds and applies the override (if any) for a node, returning the new
256
// potentially overridden node
257
function applyOverride<
258
  Node extends ts.ClassDeclaration | ts.InterfaceDeclaration,
259
>(
260
  ctx: ts.TransformationContext,
261
  overrideCtx: OverrideTransformContext,
262
  node: Node,
263
  updateDeclaration: (node: Node, override: ts.ClassDeclaration) => Node
264
): ts.Node {
265
  assert(node.name !== undefined);
266
  const name = node.name.text;
267
  const override = maybeGetOverride(overrideCtx.program, name);
268
  const isReplacement = overrideCtx.replacements.has(name);
269
 
270
  // Full-type replacement may rename type too, so record renames now
271
  if (override !== undefined) {
272
    // If override's name is different to the node's name, rename it later
273
    const overrideIdentifier = maybeGetStatementName(override);
274
    if (overrideIdentifier !== undefined) {
275
      const overrideName = overrideIdentifier.text;
276
      if (name !== overrideName) overrideCtx.renames.set(name, overrideName);
277
    }
278
  }
279
 
280
  if (isReplacement) {
281
    assert(override !== undefined);
282
    return ensureStatementModifiers(ctx, override, {
283
      declare: true,
284
      export: false,
285
    });
286
  } else if (override !== undefined) {
287
    // Merge override into declaration. Whilst we convert all non-replacement
288
    // overrides to classes, this type classification is ignored when merging.
289
    // Classes just support all possible forms of override. See `./compiler.ts`
290
    // `compileOverride()` for details.
291
    assert(ts.isClassDeclaration(override));
292
    return updateDeclaration(node, override);
293
  } else {
294
    // No override, so return the node as is
295
    return node;
296
  }
297
}
298
 
299
// Apply all overrides, insert defines, and record type renames
300
function createOverrideDefineVisitor(
301
  ctx: ts.TransformationContext,
302
  overrideCtx: OverrideTransformContext
303
): ts.Visitor {
304
  // Copies all string and numeric literals. Without this, garbage would be
305
  // inserted in locations of literals instead.
306
  // TODO(soon): work out why this happens, something to do with source ranges
307
  //  and invalid source files/programs maybe?
308
  const copyLiteralsVisitor: ts.Visitor<ts.Node, ts.Node> = (node) => {
309
    node = ts.visitEachChild(node, copyLiteralsVisitor, ctx);
310
    if (ts.isStringLiteral(node)) {
311
      return ctx.factory.createStringLiteral(node.text);
312
    }
313
    if (ts.isNumericLiteral(node)) {
314
      return ctx.factory.createNumericLiteral(node.text);
315
    }
316
    return node;
317
  };
318
 
319
  const visitor: ts.Visitor = (node) => {
320
    // Visit classes and interfaces inside module declarations too
321
    if (ts.isModuleDeclaration(node) || ts.isModuleBody(node)) {
322
      return ts.visitEachChild(node, visitor, ctx);
323
    }
324
 
325
    let defines: ts.NodeArray<ts.Statement> | undefined;
326
 
327
    if (ts.isClassDeclaration(node) && node.name !== undefined) {
328
      defines = maybeGetDefines(overrideCtx.program, node.name.text);
329
      node = applyOverride(ctx, overrideCtx, node, (node, override) => {
330
        return ctx.factory.updateClassDeclaration(
331
          node,
332
          node.modifiers,
333
          override.name,
334
          override.typeParameters ?? node.typeParameters,
335
          override.heritageClauses ?? node.heritageClauses,
336
          mergeMembers(node.members, override.members, (member) => member)
337
        );
338
      });
339
    } else if (ts.isInterfaceDeclaration(node)) {
340
      defines = maybeGetDefines(overrideCtx.program, node.name.text);
341
      node = applyOverride(ctx, overrideCtx, node, (node, override) => {
342
        assert(override.name !== undefined);
343
        return ctx.factory.updateInterfaceDeclaration(
344
          node,
345
          node.modifiers,
346
          override.name,
347
          override.typeParameters ?? node.typeParameters,
348
          override.heritageClauses ?? node.heritageClauses,
349
          mergeMembers(node.members, override.members, (member) =>
350
            classToTypeElement(ctx, member)
351
          )
352
        );
353
      });
354
    }
355
 
356
    // Process node and defines if defined
357
    node = ts.visitNode(node, copyLiteralsVisitor);
358
    defines = ts.visitNodes(defines, copyLiteralsVisitor, ts.isStatement);
359
    defines = ts.visitNodes(
360
      defines,
361
      (node) =>
362
        ensureStatementModifiers(ctx, node, { declare: true, export: false }),
363
      ts.isStatement
364
    );
365
 
366
    if (ts.isTypeAliasDeclaration(node) && isUnsatisfiable(node.type)) {
367
      // If node was overridden to `type T = never`, delete it, and just insert
368
      // defines if any
369
      return defines === undefined ? undefined : [...defines];
370
    } else {
371
      // Otherwise, return potentially overridden node, inserting defines if any
372
      // before node
373
      return defines == undefined ? node : [...defines, node];
374
    }
375
  };
376
  return visitor;
377
}
378
 
379
// Apply previously-recorded type renames to all type references
380
export function createRenameVisitor(
381
  ctx: ts.TransformationContext,
382
  renames: Map</* from */ string, /* to */ string>,
383
  renameClassesInterfaces = false
384
): ts.Visitor {
385
  const visitor: ts.Visitor = (node) => {
386
    // Recursively visit all nodes
387
    node = ts.visitEachChild(node, visitor, ctx);
388
 
389
    // Rename all type references
390
    if (ts.isTypeReferenceNode(node) && ts.isIdentifier(node.typeName)) {
391
      const rename = renames.get(node.typeName.text);
392
      if (rename !== undefined) {
393
        return ctx.factory.updateTypeReferenceNode(
394
          node,
395
          ctx.factory.createIdentifier(rename),
396
          node.typeArguments
397
        );
398
      }
399
    }
400
 
401
    // Rename all type queries (e.g. nested types)
402
    if (ts.isTypeQueryNode(node) && ts.isIdentifier(node.exprName)) {
403
      const rename = renames.get(node.exprName.text);
404
      if (rename !== undefined) {
405
        return ctx.factory.updateTypeQueryNode(
406
          node,
407
          ctx.factory.createIdentifier(rename),
408
          node.typeArguments
409
        );
410
      }
411
    }
412
 
413
    // Rename all expressions with type arguments (e.g. heritage clauses)
414
    if (
415
      ts.isExpressionWithTypeArguments(node) &&
416
      ts.isIdentifier(node.expression)
417
    ) {
418
      const rename = renames.get(node.expression.text);
419
      if (rename !== undefined) {
420
        return ctx.factory.updateExpressionWithTypeArguments(
421
          node,
422
          ctx.factory.createIdentifier(rename),
423
          node.typeArguments
424
        );
425
      }
426
    }
427
 
428
    // Rename all class and interface names
429
    if (renameClassesInterfaces) {
430
      if (
431
        ts.isClassDeclaration(node) &&
432
        node.name !== undefined &&
433
        ts.isIdentifier(node.name)
434
      ) {
435
        const rename = renames.get(node.name.text);
436
        if (rename !== undefined) {
437
          return ctx.factory.updateClassDeclaration(
438
            node,
439
            node.modifiers,
440
            ctx.factory.createIdentifier(rename),
441
            node.typeParameters,
442
            node.heritageClauses,
443
            node.members
444
          );
445
        }
446
      }
447
      if (ts.isInterfaceDeclaration(node) && ts.isIdentifier(node.name)) {
448
        const rename = renames.get(node.name.text);
449
        if (rename !== undefined) {
450
          return ctx.factory.updateInterfaceDeclaration(
451
            node,
452
            node.modifiers,
453
            ctx.factory.createIdentifier(rename),
454
            node.typeParameters,
455
            node.heritageClauses,
456
            node.members
457
          );
458
        }
459
      }
460
    }
461
 
462
    return node;
463
  };
464
  return visitor;
465
}
466
 
467
// Returns a statement's identifier if it has one. This is used to get the
468
// name to rename too (if any), hence we ignore variables and functions, since
469
// replacing a reference to a type with a referencing to a variable/function is
470
// an error.
471
function maybeGetStatementName(node: ts.Statement): ts.Identifier | undefined {
472
  if (
473
    ts.isClassDeclaration(node) ||
474
    ts.isInterfaceDeclaration(node) ||
475
    ts.isEnumDeclaration(node) ||
476
    ts.isTypeAliasDeclaration(node)
477
  ) {
478
    return node.name;
479
  }
480
}