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
16import * as ts from 'typescript';
17import { LinterConfig } from '../TypeScriptLinterConfig';
18import { forEachNodeInSubtree } from '../utils/functions/ForEachNodeInSubtree';
19import { isStructDeclaration } from '../utils/functions/IsStruct';
20import type { TsUtils } from '../utils/TsUtils';
21
22export class SymbolCache {
23  constructor(
24    private readonly typeChecker: ts.TypeChecker,
25    private readonly utils: TsUtils,
26    readonly sourceFile: ts.SourceFile,
27    readonly cancellationToken?: ts.CancellationToken
28  ) {
29    const callback = (node: ts.Node): void => {
30      if (isStructDeclaration(node)) {
31        // early exit via exception if cancellation was requested
32        this.cancellationToken?.throwIfCancellationRequested();
33      }
34
35      const symbol = this.handlersMap.get(node.kind)?.call(this, node);
36      if (symbol !== undefined) {
37        this.addReference(symbol, node);
38      }
39    };
40
41    const stopCondition = (node: ts.Node): boolean => {
42      return !node || LinterConfig.terminalTokens.has(node.kind);
43    };
44
45    forEachNodeInSubtree(sourceFile, callback, stopCondition);
46  }
47
48  getReferences(symbol: ts.Symbol): ts.Node[] {
49    return this.cache.get(symbol) ?? [];
50  }
51
52  private handleElementAccessExpression(node: ts.Node): ts.Symbol | undefined {
53    const elementAccessExpr = node as ts.ElementAccessExpression;
54    return this.typeChecker.getSymbolAtLocation(elementAccessExpr.argumentExpression);
55  }
56
57  private handleEnumDeclaration(node: ts.Node): ts.Symbol | undefined {
58    const enumDeclaration = node as ts.EnumDeclaration;
59    return this.utils.trueSymbolAtLocation(enumDeclaration.name);
60  }
61
62  private handlePrivateIdentifier(node: ts.Node): ts.Symbol | undefined {
63    const privateIdentifier = node as ts.PrivateIdentifier;
64    return this.typeChecker.getSymbolAtLocation(privateIdentifier);
65  }
66
67  private handlePropertyAssignment(node: ts.Node): ts.Symbol | undefined {
68    const propertyAssignment = node as ts.PropertyAssignment;
69    const contextualType = this.typeChecker.getContextualType(propertyAssignment.parent);
70    return contextualType === undefined ? undefined : this.utils.getPropertySymbol(contextualType, propertyAssignment);
71  }
72
73  private handlePropertyDeclaration(node: ts.Node): ts.Symbol | undefined {
74    const propertyDeclaration = node as ts.PropertyDeclaration;
75    return this.typeChecker.getSymbolAtLocation(propertyDeclaration.name);
76  }
77
78  private handlePropertySignature(node: ts.Node): ts.Symbol | undefined {
79    const propertySignature = node as ts.PropertySignature;
80    return this.typeChecker.getSymbolAtLocation(propertySignature.name);
81  }
82
83  private handleFunctionDeclaration(node: ts.Node): ts.Symbol | undefined {
84    const functionDeclaration = node as ts.FunctionDeclaration;
85    return functionDeclaration.name ? this.typeChecker.getSymbolAtLocation(functionDeclaration.name) : undefined;
86  }
87
88  private handleCallExpression(node: ts.Node): ts.Symbol | undefined {
89    const callExpression = node as ts.CallExpression;
90    return this.typeChecker.getSymbolAtLocation(callExpression.expression);
91  }
92
93  private handleIdentifier(node: ts.Node): ts.Symbol | undefined {
94    const identifier = node as ts.Identifier;
95    const symbol = this.typeChecker.getSymbolAtLocation(identifier);
96    if (symbol?.flags) {
97      return (symbol.flags & ts.SymbolFlags.Variable) !== 0 ? symbol : undefined;
98    }
99    return undefined;
100  }
101
102  private addReference(symbol: ts.Symbol, node: ts.Node): void {
103    let nodes = this.cache.get(symbol);
104    if (nodes === undefined) {
105      nodes = [];
106      this.cache.set(symbol, nodes);
107    }
108    nodes.push(node);
109  }
110
111  private readonly handlersMap = new Map([
112    [ts.SyntaxKind.ElementAccessExpression, this.handleElementAccessExpression],
113    [ts.SyntaxKind.EnumDeclaration, this.handleEnumDeclaration],
114    [ts.SyntaxKind.PrivateIdentifier, this.handlePrivateIdentifier],
115    [ts.SyntaxKind.PropertyAssignment, this.handlePropertyAssignment],
116    [ts.SyntaxKind.PropertyDeclaration, this.handlePropertyDeclaration],
117    [ts.SyntaxKind.PropertySignature, this.handlePropertySignature],
118    [ts.SyntaxKind.FunctionDeclaration, this.handleFunctionDeclaration],
119    [ts.SyntaxKind.CallExpression, this.handleCallExpression],
120    [ts.SyntaxKind.Identifier, this.handleIdentifier]
121  ]);
122
123  private readonly cache: Map<ts.Symbol, ts.Node[]> = new Map<ts.Symbol, ts.Node[]>();
124}
125