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 { printNode } from "../print";
8
 
9
// Replaces custom Iterator-like interfaces with built-in `Iterator` types:
10
//
11
// ```ts
12
// export class Thing {
13
//   readonly things: ThingIterator;
14
//   asyncThings(): AsyncThingIterator;
15
// }
16
//
17
// export interface ThingIterator extends Iterator {
18
//   next(): ThingIteratorNext;
19
//   [Symbol.iterator](): any;
20
// }
21
// export interface ThingIteratorNext {
22
//   done: boolean;
23
//   value?: string;
24
// }
25
//
26
// export interface AsyncThingIterator extends AsyncIterator {
27
//   next(): Promise<AsyncThingIteratorNext>;
28
//   return(value?: any): Promise<AsyncThingIteratorNext>;
29
//   [Symbol.asyncIterator](): any;
30
// }
31
// export interface AsyncThingIteratorNext {
32
//   done: boolean;
33
//   value?: number;
34
// }
35
// ```
36
//
37
// --- transforms to --->
38
//
39
// ```ts
40
// export class Thing {
41
//   readonly things: IterableIterator<string>;
42
//   asyncThings(): AsyncIterableIterator<number>;
43
// }
44
// ```
45
export function createIteratorTransformer(
46
  checker: ts.TypeChecker
47
): ts.TransformerFactory<ts.SourceFile> {
48
  return (ctx) => {
49
    return (node) => {
50
      const iteratorCtx: IteratorTransformContext = {
51
        types: new Map(),
52
        nextInterfaces: new Set(),
53
      };
54
      const v1 = createIteratorDeclarationsVisitor(ctx, checker, iteratorCtx);
55
      const v2 = createIteratorUsagesVisitor(ctx, checker, iteratorCtx);
56
      node = ts.visitEachChild(node, v1, ctx);
57
      return ts.visitEachChild(node, v2, ctx);
58
    };
59
  };
60
}
61
 
62
interface IteratorTransformContext {
63
  // Maps iterator-like interfaces to built-in `Iterator` types
64
  types: Map<ts.Symbol, ts.TypeNode>;
65
  // Set of iterator-next interfaces to remove
66
  nextInterfaces: Set<ts.Symbol>;
67
}
68
 
69
// Find all interfaces extending `Iterator`, record their next value type,
70
// and remove them. Also record the names of next interfaces for removal.
71
function createIteratorDeclarationsVisitor(
72
  ctx: ts.TransformationContext,
73
  checker: ts.TypeChecker,
74
  iteratorCtx: IteratorTransformContext
75
): ts.Visitor {
76
  const visitor: ts.Visitor = (node) => {
77
    // Visit interfaces inside module declarations too
78
    if (ts.isModuleDeclaration(node) || ts.isModuleBody(node)) {
79
      return ts.visitEachChild(node, visitor, ctx);
80
    }
81
 
82
    if (ts.isInterfaceDeclaration(node)) {
83
      // Check if interface extends `Iterator`
84
      const extendsNode = node.heritageClauses?.[0];
85
      if (
86
        extendsNode?.token === ts.SyntaxKind.ExtendsKeyword &&
87
        extendsNode.types.length === 1 &&
88
        ts.isIdentifier(extendsNode.types[0].expression) &&
89
        (extendsNode.types[0].expression.text === "Iterator" ||
90
          extendsNode.types[0].expression.text === "AsyncIterator")
91
      ) {
92
        const isAsync = extendsNode.types[0].expression.text !== "Iterator";
93
        // Check `node` has one of the following shapes:
94
        // ```ts
95
        // export interface ThingIterator extends Iterator {
96
        //   next(): ThingIteratorNext;
97
        //   [Symbol.iterator](): any;
98
        // }
99
        // export interface AsyncThingIterator extends AsyncIterator {
100
        //   next(): Promise<AsyncThingIteratorNext>;
101
        //   return(value?: any): Promise<AsyncThingIteratorNext>;
102
        //   [Symbol.asyncIterator](): any;
103
        // }
104
        // ```
105
        let nextTypeNode: ts.TypeNode | undefined;
106
        for (const member of node.members) {
107
          if (
108
            ts.isMethodSignature(member) &&
109
            ts.isIdentifier(member.name) &&
110
            member.name.text === "next" &&
111
            member.type !== undefined
112
          ) {
113
            nextTypeNode = member.type;
114
          }
115
        }
116
        assert(
117
          nextTypeNode !== undefined,
118
          `Expected iterator-like interface, got "${printNode(node)}"`
119
        );
120
 
121
        // Extract `IteratorBase_ThingIterator_...Next` type
122
        if (isAsync) {
123
          // Unwrap Promise type
124
          assert(
125
            ts.isTypeReferenceNode(nextTypeNode) &&
126
              ts.isIdentifier(nextTypeNode.typeName) &&
127
              nextTypeNode.typeName.text === "Promise" &&
128
              nextTypeNode.typeArguments?.length === 1,
129
            `Expected Promise, got "${printNode(nextTypeNode)}"`
130
          );
131
          nextTypeNode = nextTypeNode.typeArguments[0];
132
        }
133
 
134
        // Check `IteratorBase_ThingIterator_...Next` has the following shape,
135
        // and extract the `value?: T` declaration
136
        // ```ts
137
        // export interface ThingIteratorNext {
138
        //   done: boolean;
139
        //   value?: string;
140
        // }
141
        // ```
142
        const nextType = checker.getTypeFromTypeNode(nextTypeNode);
143
        const nextTypeSymbol = nextType.getSymbol();
144
        assert(nextTypeSymbol?.members !== undefined);
145
        let nextValueSymbol: ts.Symbol | undefined;
146
        nextTypeSymbol.members.forEach((value, key) => {
147
          // eslint-disable-next-line @typescript-eslint/no-unsafe-enum-comparison
148
          if (key === "value") nextValueSymbol = value;
149
        });
150
        assert(nextValueSymbol !== undefined);
151
        const nextValueDeclarations = nextValueSymbol.getDeclarations();
152
        assert.strictEqual(nextValueDeclarations?.length, 1);
153
        const nextValueDeclaration = nextValueDeclarations[0];
154
        assert(ts.isPropertySignature(nextValueDeclaration));
155
        // Mark this interface for removal
156
        iteratorCtx.nextInterfaces.add(nextTypeSymbol);
157
 
158
        // Extract `value`'s type
159
        const nextValueType = nextValueDeclaration.type;
160
        assert(nextValueType !== undefined);
161
 
162
        // Record this iterator type...
163
        const nodeType = checker.getTypeAtLocation(node);
164
        const nodeSymbol = nodeType.getSymbol();
165
        assert(nodeSymbol !== undefined);
166
        const iteratorType = ctx.factory.createTypeReferenceNode(
167
          isAsync ? "AsyncIterableIterator" : "IterableIterator",
168
          [nextValueType]
169
        );
170
        iteratorCtx.types.set(nodeSymbol, iteratorType);
171
        // ...and remove the node by returning `undefined`
172
        return;
173
      }
174
    }
175
 
176
    return node;
177
  };
178
  return visitor;
179
}
180
 
181
// Replace uses of iterator interfaces with built-in iterator type.
182
// Also remove all previously recorded next interfaces.
183
function createIteratorUsagesVisitor(
184
  ctx: ts.TransformationContext,
185
  checker: ts.TypeChecker,
186
  iteratorCtx: IteratorTransformContext
187
): ts.Visitor {
188
  // Find the built-in iterator type associated with a method's return type
189
  // or property's type
190
  function findIteratorType(
191
    node:
192
      | ts.MethodSignature
193
      | ts.MethodDeclaration
194
      | ts.PropertySignature
195
      | ts.PropertyDeclaration
196
      | ts.GetAccessorDeclaration
197
  ): ts.TypeNode | undefined {
198
    if (node.type === undefined) return;
199
    const type = checker.getTypeFromTypeNode(node.type);
200
    const typeSymbol = type.getSymbol();
201
    if (typeSymbol !== undefined) return iteratorCtx.types.get(typeSymbol);
202
  }
203
 
204
  const visitor: ts.Visitor = (node) => {
205
    // Remove all next interfaces by returning `undefined`
206
    if (ts.isInterfaceDeclaration(node)) {
207
      const type = checker.getTypeAtLocation(node);
208
      const symbol = type.getSymbol();
209
      if (symbol !== undefined && iteratorCtx.nextInterfaces.has(symbol)) {
210
        return;
211
      }
212
    }
213
 
214
    // Visit all interface/class/module declaration children
215
    if (
216
      ts.isInterfaceDeclaration(node) ||
217
      ts.isClassDeclaration(node) ||
218
      ts.isModuleDeclaration(node) ||
219
      ts.isModuleBody(node)
220
    ) {
221
      return ts.visitEachChild(node, visitor, ctx);
222
    }
223
 
224
    // Replace all method return types and property types referencing iterators
225
    // with the built-in type
226
    if (ts.isMethodSignature(node)) {
227
      const iteratorType = findIteratorType(node);
228
      if (iteratorType !== undefined) {
229
        return ctx.factory.updateMethodSignature(
230
          node,
231
          node.modifiers,
232
          node.name,
233
          node.questionToken,
234
          node.typeParameters,
235
          node.parameters,
236
          iteratorType
237
        );
238
      }
239
    }
240
    if (ts.isMethodDeclaration(node)) {
241
      const iteratorType = findIteratorType(node);
242
      if (iteratorType !== undefined) {
243
        return ctx.factory.updateMethodDeclaration(
244
          node,
245
          node.modifiers,
246
          node.asteriskToken,
247
          node.name,
248
          node.questionToken,
249
          node.typeParameters,
250
          node.parameters,
251
          iteratorType,
252
          node.body
253
        );
254
      }
255
    }
256
    if (ts.isPropertySignature(node)) {
257
      const iteratorType = findIteratorType(node);
258
      if (iteratorType !== undefined) {
259
        return ctx.factory.updatePropertySignature(
260
          node,
261
          node.modifiers,
262
          node.name,
263
          node.questionToken,
264
          iteratorType
265
        );
266
      }
267
    }
268
    if (ts.isPropertyDeclaration(node)) {
269
      const iteratorType = findIteratorType(node);
270
      if (iteratorType !== undefined) {
271
        return ctx.factory.updatePropertyDeclaration(
272
          node,
273
          node.modifiers,
274
          node.name,
275
          node.questionToken ?? node.exclamationToken,
276
          iteratorType,
277
          node.initializer
278
        );
279
      }
280
    }
281
    if (ts.isGetAccessorDeclaration(node)) {
282
      const iteratorType = findIteratorType(node);
283
      if (iteratorType !== undefined) {
284
        return ctx.factory.updateGetAccessorDeclaration(
285
          node,
286
          node.modifiers,
287
          node.name,
288
          node.parameters,
289
          iteratorType,
290
          node.body
291
        );
292
      }
293
    }
294
 
295
    return node;
296
  };
297
  return visitor;
298
}