horror code

index.ts

import assert = require("assert");
import { collectVariableUsage, getTokenAtPosition, VariableInfo } from "tsutils";
import { createSourceFile, Identifier, isVariableDeclaration, ScriptKind, ScriptTarget, SourceFile, VariableDeclaration } from "typescript";
import * as vscode from "vscode";

import { formatModule, mkStringUri, sendAndGetData } from "../webSocketServer";
import { findExportLocation, findObjectLiteralByKey, findParrent, findWebpackArg, findWreq_d, getLeadingIdentifier, getModuleId, makeRange, zeroRange } from "./util";
import ts = require("typescript");
import format from "../format";
import { findReturnIdentifier } from "./util";


type Definitions = Promise<vscode.Definition | vscode.LocationLink[] | null | undefined>;

export class DefinitionProvider implements vscode.DefinitionProvider {
    async provideDefinition(document: vscode.TextDocument, position: vscode.Position): Definitions {
        try {
            const text = document.getText();
            // not sure if substring is a good idea here
            // just dont want to search really long webpack modules
            if (!text.startsWith("//WebpackModule") || text.substring(0, 100).includes("//OPEN FULL MODULE:")) return;
            return await new WebpackAstParser(text).generateDefinitions(document, position);
        } catch (e) {
            console.error(e);
        }
    }

}

class WebpackAstParser {
    private text: string;
    private sourceFile: SourceFile;
    /** All vars in the file */
    private vars: Map<Identifier, VariableInfo>;
    /** The webpack instanse */
    private wreq: Identifier;
    /** where {@link WebpackAstParser.wreq this.wreq} is used*/
    private uses: VariableInfo | undefined;
    public constructor(text: string) {
        this.text = text;
        this.sourceFile = createSourceFile("module.js", this.text, ScriptTarget.ESNext, true, ScriptKind.JS);
        this.vars = collectVariableUsage(this.sourceFile);

        const wreq = findWebpackArg(this.sourceFile);
        if (!wreq) throw new Error("Couldn't find wreq!");
        else this.wreq = wreq;

        const uses = this.vars.get(this.wreq);
        this.uses = uses;
    }

    public async generateDefinitions(document: vscode.TextDocument, position: vscode.Position): Definitions {
        if (!this.uses) throw new Error("Wreq isnt used anywhere");
        const modules = new Map([...this.vars.entries()].filter(([k, _v]) => {
            return this.uses!.uses.some(e => {
                const node = findParrent<VariableDeclaration>(e.location, isVariableDeclaration);
                return node?.name === k;
            });
        }));
        const x = getTokenAtPosition(this.sourceFile, document.offsetAt(position), this.sourceFile);
        const [y, q] = getLeadingIdentifier(x);
        if (!y) return;
        const [_, dec] = [...modules.entries()].find(([k, v]) => {
            return v.uses.some(({ location }) => y === location);
        }) ?? [];
        const moduleId = getModuleId(dec, q);
        if (!moduleId) return;
        const res = await sendAndGetData<{ data: string }>({
            type: "rawId",
            data: {
                id: moduleId
            }
        })
            .catch(console.error);
        if (res?.data === undefined) return;
        res.data = await format(formatModule(res.data, moduleId));
        return {
            range: q ? new WebpackAstParser(res.data).findExportLocation(q.text) : zeroRange,
            uri: mkStringUri(res.data),
        };
    }

    public findExportLocation(exportName: string): vscode.Range {
        if (this.uses) {
            const wreq_d = this.uses.uses.find(findWreq_d);
            const call = wreq_d?.location.parent.parent;
            if (!call || !ts.isCallExpression(call)) return zeroRange

            if (call.arguments.length !== 2
                || !ts.isIdentifier(call.arguments[0])
                || !ts.isObjectLiteralExpression(call.arguments[1])
            ) return zeroRange;
            const moduleExports = call.arguments[1];

            const exportProperty = findObjectLiteralByKey(moduleExports, exportName);

            if (!exportProperty || !ts.isPropertyAssignment(exportProperty) || !ts.isFunctionExpression(exportProperty.initializer)) return zeroRange

            const exportVar = findReturnIdentifier(exportProperty.initializer)
            
            const [exportDec] = [...this.vars.entries()].find(([_k, v]) => {
                return v.uses.some(u => u.location === exportVar)
            }) ?? [] 
            
            if(!exportDec) return zeroRange;

            return makeRange(exportDec, this.text);
        }
        return zeroRange;
    }
}