@jpoly1219/context-extractor 0.2.6 → 0.2.8

package/dist/src/app.d.ts

@@ -7,7 +7,6 @@ export declare class App {
     private sketchPath;
     private repoPath;
     private result;
-    private timeout;
     constructor(language: Language, sketchPath: string, repoPath: string);
     init(): Promise<void>;
     run(): Promise<void>;

package/dist/src/app.js

@@ -40,6 +40,13 @@ const typescript_driver_1 = require("./typescript-driver");
 const ocaml_driver_1 = require("./ocaml-driver");
 const utils_1 = require("./utils");
 class App {
+    // Optional timeout for forced termination
+    // private timeout = setTimeout(() => {
+    //   if (!this.languageServer.killed) {
+    //     console.log('Forcibly killing the process...');
+    //     this.languageServer.kill('SIGKILL');
+    //   }
+    // }, 5000);
     constructor(language, sketchPath, repoPath) {
         // private result: {
         //   hole: string;
@@ -47,13 +54,6 @@ class App {
         //   relevantHeaders: string[];
         // } | null = null;
         this.result = null;
-        // Optional timeout for forced termination
-        this.timeout = setTimeout(() => {
-            if (!this.languageServer.killed) {
-                console.log('Forcibly killing the process...');
-                this.languageServer.kill('SIGKILL');
-            }
-        }, 5000);
         this.language = language;
         this.sketchPath = sketchPath;
         this.repoPath = repoPath;
@@ -93,6 +93,7 @@ class App {
         const c = new main_1.LspClient(e);
         this.languageServer = r;
         this.lspClient = c;
+        // console.log(r.pid)
         this.languageServer.on('close', (code) => {
             if (code !== 0) {
                 console.log(`ls process exited with code ${code}`);
@@ -100,7 +101,7 @@ class App {
         });
         // Clear timeout once the process exits
         this.languageServer.on('exit', () => {
-            clearTimeout(this.timeout);
+            // clearTimeout(this.timeout);
             console.log('Process terminated cleanly.');
         });
         // const logFile = fs.createWriteStream("log.txt");
@@ -113,9 +114,17 @@ class App {
         // const outputFile = fs.createWriteStream("output.txt");
         try {
             await this.init();
+            // console.time("getHoleContext");
             const holeContext = await this.languageDriver.getHoleContext(this.lspClient, this.sketchPath);
-            const relevantTypes = await this.languageDriver.extractRelevantTypes(this.lspClient, holeContext.fullHoverResult, holeContext.functionName, holeContext.range.start.line, holeContext.range.end.line, new Map(), holeContext.source);
+            // console.timeEnd("getHoleContext");
+            // console.time("extractRelevantTypes");
+            const relevantTypes = await this.languageDriver.extractRelevantTypes(this.lspClient, 
+            // NOTE: sometimes fullHoverResult isn't representative of the actual file contents, especially with generic functions.
+            holeContext.trueHoleFunction ? holeContext.trueHoleFunction : holeContext.fullHoverResult, holeContext.functionName, holeContext.range.start.line, new Map(), holeContext.source, new Map());
+            // console.timeEnd("extractRelevantTypes");
+            // console.dir(relevantTypes, { depth: null })
             // Postprocess the map.
+            // console.time("extractRelevantTypes postprocess");
             if (this.language === types_1.Language.TypeScript) {
                 relevantTypes.delete("_()");
                 for (const [k, { typeSpan: v, sourceFile: src }] of relevantTypes.entries()) {
@@ -125,7 +134,8 @@ class App {
             else if (this.language === types_1.Language.OCaml) {
                 relevantTypes.delete("_");
             }
-            console.log(path.join(path.dirname(this.sketchPath), `sketch${path.extname(this.sketchPath)}`));
+            // console.timeEnd("extractRelevantTypes postprocess");
+            // console.time("extractRelevantHeaders repo");
             let repo = [];
             if (this.language === types_1.Language.TypeScript) {
                 repo = (0, utils_1.getAllTSFiles)(this.repoPath);
@@ -133,8 +143,15 @@ class App {
             else if (this.language === types_1.Language.OCaml) {
                 repo = (0, utils_1.getAllOCamlFiles)(this.repoPath);
             }
+            // console.timeEnd("extractRelevantHeaders repo");
+            // console.time("extractRelevantHeaders");
             const relevantHeaders = await this.languageDriver.extractRelevantHeaders(this.lspClient, repo, relevantTypes, holeContext.functionTypeSpan);
+            // const relevantHeaders: { typeSpan: string, sourceFile: string }[] = []
+            // console.timeEnd("extractRelevantHeaders");
+            // console.log(relevantHeaders)
+            // console.log(relevantHeaders.size)
             // Postprocess the map.
+            // console.time("extractRelevantHaders postprocess");
             if (this.language === types_1.Language.TypeScript) {
                 relevantTypes.delete("");
                 for (const [k, { typeSpan: v, sourceFile: src }] of relevantTypes.entries()) {
@@ -144,6 +161,8 @@ class App {
                     obj.typeSpan += ";";
                 }
             }
+            // console.timeEnd("extractRelevantHeaders postprocess");
+            // console.time("toReturn");
             const relevantTypesToReturn = new Map();
             relevantTypes.forEach(({ typeSpan: v, sourceFile: src }, _) => {
                 if (relevantTypesToReturn.has(src)) {
@@ -168,6 +187,7 @@ class App {
                     relevantHeadersToReturn.set(src, [v]);
                 }
             });
+            // console.timeEnd("toReturn");
             this.result = {
                 holeType: holeContext.functionTypeSpan,
                 relevantTypes: relevantTypesToReturn,

package/dist/src/main.js

@@ -229,10 +229,12 @@ const extractWithCodeQL = async (sketchPath) => {
 };
 exports.extractWithCodeQL = extractWithCodeQL;
 const extractContext = async (language, sketchPath, repoPath) => {
+    // console.time("extractContext")
     const app = new app_1.App(language, sketchPath, repoPath);
     await app.run();
     const res = app.getSavedResult();
     app.close();
+    // console.timeEnd("extractContext")
     return res;
     // if (!getCompletion) {
     //   await app.close()

package/dist/src/ocaml-driver.d.ts

@@ -15,7 +15,7 @@ export declare class OcamlDriver implements LanguageDriver {
         range: Range;
         source: string;
     }>;
-    extractRelevantTypes(lspClient: LspClient, fullHoverResult: string, typeName: string, startLine: number, endLine: number, foundSoFar: Map<string, TypeSpanAndSourceFile>, currentFile: string): Promise<Map<string, TypeSpanAndSourceFile>>;
+    extractRelevantTypes(lspClient: LspClient, fullHoverResult: string, typeName: string, startLine: number, foundSoFar: Map<string, TypeSpanAndSourceFile>, currentFile: string): Promise<Map<string, TypeSpanAndSourceFile>>;
     extractRelevantHeaders(lspClient: LspClient, sources: string[], relevantTypes: Map<string, TypeSpanAndSourceFile>, holeType: string): Promise<Set<TypeSpanAndSourceFile>>;
     extractHeaderTypeSpans(lspClient: LspClient, preludeFilePath: string): Promise<{
         identifier: string;

package/dist/src/ocaml-driver.js

@@ -172,12 +172,12 @@ class OcamlDriver {
             source: `file://${sketchFilePath}`
         };
     }
-    async extractRelevantTypes(lspClient, fullHoverResult, typeName, startLine, endLine, foundSoFar, currentFile) {
+    async extractRelevantTypes(lspClient, fullHoverResult, typeName, startLine, foundSoFar, currentFile) {
         if (!foundSoFar.has(typeName)) {
             foundSoFar.set(typeName, { typeSpan: fullHoverResult, sourceFile: currentFile.slice(7) });
             // outputFile.write(`${fullHoverResult};\n`);
             const content = fs.readFileSync(currentFile.slice(7), "utf8");
-            for (let linePos = startLine; linePos <= endLine; ++linePos) {
+            for (let linePos = startLine; linePos <= fullHoverResult.length; ++linePos) {
                 const numOfCharsInLine = parseInt((0, child_process_1.execSync)(`wc -m <<< "${content.split("\n")[linePos]}"`, { shell: "/bin/bash" }).toString());
                 for (let charPos = 0; charPos < numOfCharsInLine; ++charPos) {
                     try {
@@ -224,7 +224,7 @@ class OcamlDriver {
                                 // TODO: this can potentially be its own method. the driver would require some way to get type context.
                                 // potentially, this type checker can be its own class.
                                 const identifier = this.typeChecker.getIdentifierFromDecl(snippetInRange);
-                                await this.extractRelevantTypes(lspClient, snippetInRange, identifier, matchingSymbolRange.start.line, matchingSymbolRange.end.line, foundSoFar, typeDefinitionResult[0].uri);
+                                await this.extractRelevantTypes(lspClient, snippetInRange, identifier, matchingSymbolRange.start.line, foundSoFar, typeDefinitionResult[0].uri);
                             }
                         }
                     }

package/dist/src/runner.js

@@ -9,8 +9,35 @@ const types_1 = require("./types");
 // extract("/home/jacob/projects/context-extractor/targets/emojipaint/sketch.ts").then(r => console.log("emojipaint\n", r));
 (async () => {
     try {
-        const x = await (0, main_1.extractContext)(types_1.Language.TypeScript, "/home/jacob/projects/context-extractor/targets/todo/sketch.ts", "/home/jacob/projects/context-extractor/targets/todo/");
+        let x;
+        // x = await extractContext(
+        //   Language.TypeScript,
+        //   "/home/jacob/projects/context-extractor/targets/todo/sketch.ts",
+        //   "/home/jacob/projects/context-extractor/targets/todo/",
+        // )
+        // console.dir(x, { depth: null })
+        //
+        // x = await extractContext(
+        //   Language.TypeScript,
+        //   "/home/jacob/projects/context-extractor/targets/playlist/sketch.ts",
+        //   "/home/jacob/projects/context-extractor/targets/playlist/",
+        // )
+        // console.dir(x, { depth: null })
+        //
+        // x = await extractContext(
+        //   Language.TypeScript,
+        //   "/home/jacob/projects/context-extractor/targets/passwords/sketch.ts",
+        //   "/home/jacob/projects/context-extractor/targets/passwords/",
+        // )
+        // console.dir(x, { depth: null })
+        x = await (0, main_1.extractContext)(types_1.Language.TypeScript, "/home/jacob/projects/context-extractor/targets/booking/sketch.ts", "/home/jacob/projects/context-extractor/targets/booking/");
         console.dir(x, { depth: null });
+        // x = await extractContext(
+        //   Language.TypeScript,
+        //   "/home/jacob/projects/context-extractor/targets/emojipaint/sketch.ts",
+        //   "/home/jacob/projects/context-extractor/targets/emojipaint/",
+        // )
+        // console.dir(x, { depth: null })
         // const y = await completeWithLLM(
         //   x!,
         //   Language.TypeScript,

package/dist/src/types.d.ts

@@ -72,8 +72,9 @@ interface LanguageDriver {
         holeTypeDefCharPos: number;
         range: Range;
         source: string;
+        trueHoleFunction?: string;
     }>;
-    extractRelevantTypes: (lspClient: LspClient, fullHoverResult: string, typeName: string, startLine: number, endLine: number, foundSoFar: Map<string, TypeSpanAndSourceFile>, currentFile: string) => Promise<Map<string, TypeSpanAndSourceFile>>;
+    extractRelevantTypes: (lspClient: LspClient, fullHoverResult: string, typeName: string, startLine: number, foundSoFar: Map<string, TypeSpanAndSourceFile>, currentFile: string, foundContents: Map<string, string>) => Promise<Map<string, TypeSpanAndSourceFile>>;
     extractRelevantHeaders: (lspClient: LspClient, sources: string[], relevantTypes: Map<string, TypeSpanAndSourceFile>, holeType: string) => Promise<Set<TypeSpanAndSourceFile>>;
 }
 type Filepath = string;

package/dist/src/typescript-driver.d.ts

@@ -1,4 +1,4 @@
-import { LspClient, Range } from "../ts-lsp-client-dist/src/main";
+import { LspClient, Location, Range, SymbolInformation } from "../ts-lsp-client-dist/src/main";
 import { LanguageDriver, TypeSpanAndSourceFile } from "./types";
 import { TypeScriptTypeChecker } from "./typescript-type-checker";
 export declare class TypeScriptDriver implements LanguageDriver {
@@ -14,14 +14,20 @@ export declare class TypeScriptDriver implements LanguageDriver {
         holeTypeDefCharPos: number;
         range: Range;
         source: string;
+        trueHoleFunction: string;
     }>;
-    extractRelevantTypes(lspClient: LspClient, fullHoverResult: string, typeName: string, startLine: number, endLine: number, foundSoFar: Map<string, TypeSpanAndSourceFile>, // identifier -> [full hover result, source]
-    currentFile: string): Promise<Map<string, TypeSpanAndSourceFile>>;
+    extractRelevantTypes1(lspClient: LspClient, fullHoverResult: string, typeName: string, startLine: number, endLine: number, foundSoFar: Map<string, TypeSpanAndSourceFile>, // identifier -> [full hover result, source]
+    currentFile: string, foundTypeDefinitions: Map<string, Location[]>, foundSymbols: Map<string, SymbolInformation[]>): Promise<Map<string, TypeSpanAndSourceFile>>;
+    extractRelevantTypes(lspClient: LspClient, fullHoverResult: string, typeName: string, startLine: number, foundSoFar: Map<string, TypeSpanAndSourceFile>, // identifier -> [full hover result, source]
+    currentFile: string, foundContents: Map<string, string>): Promise<Map<string, TypeSpanAndSourceFile>>;
+    extractRelevantTypesHelper(lspClient: LspClient, fullHoverResult: string, typeName: string, startLine: number, foundSoFar: Map<string, TypeSpanAndSourceFile>, // identifier -> [full hover result, source]
+    currentFile: string, foundContents: Map<string, string>, // uri -> contents
+    layer: number): Promise<void>;
     extractRelevantHeaders(_: LspClient, sources: string[], relevantTypes: Map<string, TypeSpanAndSourceFile>, holeType: string): Promise<Set<TypeSpanAndSourceFile>>;
     generateTargetTypes(relevantTypes: Map<string, TypeSpanAndSourceFile>, holeType: string): Set<string>;
     generateTargetTypesHelper(relevantTypes: Map<string, TypeSpanAndSourceFile>, currType: string, targetTypes: Set<string>): void;
-    extractRelevantHeadersHelper(typeSpan: string, targetTypes: Set<string>, relevantTypes: Map<string, TypeSpanAndSourceFile>, relevantContext: Set<TypeSpanAndSourceFile>, line: string, source: string): void;
-    isTypeEquivalent(t1: string, t2: string, relevantTypes: Map<string, TypeSpanAndSourceFile>): boolean;
+    extractRelevantHeadersHelper(typeSpan: string, targetTypes: Set<string>, relevantTypes: Map<string, TypeSpanAndSourceFile>, relevantContext: Set<TypeSpanAndSourceFile>, line: string, source: string, relevantContextMap: Map<string, TypeSpanAndSourceFile>, tag: boolean, trace: string[], foundNormalForms: Map<string, string>): void;
+    isTypeEquivalent(t1: string, t2: string, relevantTypes: Map<string, TypeSpanAndSourceFile>, foundNormalForms: Map<string, string>): boolean;
     normalize(typeSpan: string, relevantTypes: Map<string, TypeSpanAndSourceFile>): string;
     normalize2(typeSpan: string, relevantTypes: Map<string, TypeSpanAndSourceFile>): string;
 }

package/dist/src/typescript-driver.js

@@ -196,18 +196,26 @@ class TypeScriptDriver {
             holeTypeDefCharPos: "declare function _(): ".length,
             // range: { start: { line: 0, character: 0 }, end: { line: 0, character: 52 } }
             range: sketchSymbol[0].location.range,
-            source: `file://${injectedSketchFilePath}`
+            source: `file://${injectedSketchFilePath}`,
+            trueHoleFunction: trueHoleFunction
         };
     }
-    async extractRelevantTypes(lspClient, fullHoverResult, typeName, startLine, endLine, foundSoFar, // identifier -> [full hover result, source]
-    currentFile) {
+    async extractRelevantTypes1(lspClient, fullHoverResult, typeName, startLine, endLine, foundSoFar, // identifier -> [full hover result, source]
+    currentFile, foundTypeDefinitions, foundSymbols
+    // outputFile: fs.WriteStream,
+    ) {
         if (!foundSoFar.has(typeName)) {
             foundSoFar.set(typeName, { typeSpan: fullHoverResult, sourceFile: currentFile.slice(7) });
             // outputFile.write(`${fullHoverResult};\n`);
             const content = fs.readFileSync(currentFile.slice(7), "utf8");
             for (let linePos = startLine; linePos <= endLine; ++linePos) {
+                // TODO: use a platform-agnostic command here
                 const numOfCharsInLine = parseInt((0, child_process_1.execSync)(`wc -m <<< "${content.split("\n")[linePos]}"`, { shell: "/bin/bash" }).toString());
-                for (let charPos = 0; charPos < numOfCharsInLine; ++charPos) {
+                const numOfCharsInLine2 = content.split("\n")[linePos].length;
+                const numOfCharsInLine3 = [...content.split("\n")[linePos]].map(c => c.codePointAt(0)).length;
+                // console.log(numOfCharsInLine === numOfCharsInLine2, content.split("\n")[linePos], numOfCharsInLine, numOfCharsInLine2, numOfCharsInLine3)
+                // console.time(`===loop ${content.split("\n")[linePos]}===`);
+                for (let charPos = 0; charPos < numOfCharsInLine2; ++charPos) {
                     try {
                         const typeDefinitionResult = await lspClient.typeDefinition({
                             textDocument: {
@@ -218,41 +226,162 @@ class TypeScriptDriver {
                                 line: linePos
                             }
                         });
+                        // if (content.split("\n")[linePos] === `type Action = AddBooking | CancelBooking | ClearBookings;`) {
+                        //   console.dir(typeDefinitionResult, { depth: null })
+                        //   console.log(charPos)
+                        // }
                         if (typeDefinitionResult && typeDefinitionResult instanceof Array && typeDefinitionResult.length != 0) {
-                            // Use documentSymbol instead of hover.
-                            // This prevents type alias "squashing" done by tsserver.
-                            // This also allows for grabbing the entire definition range and not just the symbol range.
-                            // PERF: feels like this could be memoized to improve performance.
-                            const documentSymbolResult = await lspClient.documentSymbol({
-                                textDocument: {
-                                    uri: typeDefinitionResult[0].uri
+                            const tdResultStr = JSON.stringify(typeDefinitionResult);
+                            if (!foundTypeDefinitions.has(tdResultStr)) {
+                                foundTypeDefinitions.set(tdResultStr, typeDefinitionResult);
+                                // Use documentSymbol instead of hover.
+                                // This prevents type alias "squashing" done by tsserver.
+                                // This also allows for grabbing the entire definition range and not just the symbol range.
+                                // PERF: feels like this could be memoized to improve performance.
+                                // console.time("docSymbol")
+                                const tdUri = typeDefinitionResult[0].uri;
+                                let documentSymbolResult;
+                                if (foundSymbols.has(tdUri)) {
+                                    documentSymbolResult = foundSymbols.get(tdUri);
+                                }
+                                else {
+                                    documentSymbolResult = await lspClient.documentSymbol({
+                                        textDocument: {
+                                            uri: typeDefinitionResult[0].uri
+                                        }
+                                    });
+                                    foundSymbols.set(tdUri, documentSymbolResult);
+                                }
+                                // console.timeEnd("docSymbol")
+                                // console.time("dsMap")
+                                const dsMap = documentSymbolResult.reduce((m, obj) => {
+                                    m.set(obj.location.range.start.line, obj.location.range);
+                                    return m;
+                                }, new Map());
+                                // console.timeEnd("dsMap")
+                                // console.log("\n")
+                                // console.dir(typeDefinitionResult, { depth: null })
+                                // console.dir(documentSymbolResult, { depth: null })
+                                // console.log("\n")
+                                // grab if the line number of typeDefinitionResult and documentSymbolResult matches
+                                // const dsMap = documentSymbolResult!.reduce((m, obj) => {
+                                //   m.set((obj as SymbolInformation).location.range.start.line, (obj as SymbolInformation).location.range as unknown as Range);
+                                //   return m;
+                                // }, new Map<number, Range>());
+                                const matchingSymbolRange = dsMap.get(typeDefinitionResult[0].range.start.line);
+                                if (matchingSymbolRange) {
+                                    const snippetInRange = (0, utils_1.extractSnippet)(fs.readFileSync(typeDefinitionResult[0].uri.slice(7)).toString("utf8"), matchingSymbolRange.start, matchingSymbolRange.end);
+                                    // TODO: this can potentially be its own method. the driver would require some way to get type context.
+                                    // potentially, this type checker can be its own class.
+                                    const identifier = this.typeChecker.getIdentifierFromDecl(snippetInRange);
+                                    await this.extractRelevantTypes1(lspClient, snippetInRange, identifier, matchingSymbolRange.start.line, matchingSymbolRange.end.line, foundSoFar, typeDefinitionResult[0].uri, foundTypeDefinitions, foundSymbols
+                                    // outputFile,
+                                    );
                                 }
-                            });
-                            // grab if the line number of typeDefinitionResult and documentSymbolResult matches
-                            const dsMap = documentSymbolResult.reduce((m, obj) => {
-                                m.set(obj.location.range.start.line, obj.location.range);
-                                return m;
-                            }, new Map());
-                            const matchingSymbolRange = dsMap.get(typeDefinitionResult[0].range.start.line);
-                            if (matchingSymbolRange) {
-                                const snippetInRange = (0, utils_1.extractSnippet)(fs.readFileSync(typeDefinitionResult[0].uri.slice(7)).toString("utf8"), matchingSymbolRange.start, matchingSymbolRange.end);
-                                // TODO: this can potentially be its own method. the driver would require some way to get type context.
-                                // potentially, this type checker can be its own class.
-                                const identifier = this.typeChecker.getIdentifierFromDecl(snippetInRange);
-                                await this.extractRelevantTypes(lspClient, snippetInRange, identifier, matchingSymbolRange.start.line, matchingSymbolRange.end.line, foundSoFar, typeDefinitionResult[0].uri);
                             }
                         }
+                        // else {
+                        //   console.log(`already found ${tdResultStr}!`)
+                        // }
                     }
                     catch (err) {
                         console.log(`${err}`);
                     }
                 }
+                // console.timeEnd(`===loop ${content.split("\n")[linePos]}===`);
             }
         }
         return foundSoFar;
     }
+    async extractRelevantTypes(lspClient, fullHoverResult, typeName, startLine, foundSoFar, // identifier -> [full hover result, source]
+    currentFile, foundContents // uri -> contents
+    ) {
+        const content = fs.readFileSync(currentFile.slice(7), "utf8");
+        // console.log(content)
+        await this.extractRelevantTypesHelper(lspClient, fullHoverResult, typeName, startLine, foundSoFar, currentFile, foundContents, 0);
+        return foundSoFar;
+    }
+    async extractRelevantTypesHelper(lspClient, fullHoverResult, typeName, startLine, foundSoFar, // identifier -> [full hover result, source]
+    currentFile, foundContents, // uri -> contents
+    layer) {
+        // console.log("===", fullHoverResult, layer, startLine, "===")
+        // Split the type span into identifiers, where each include the text, line number, and character range.
+        // For each identifier, invoke go to type definition.
+        if (!foundSoFar.has(typeName)) {
+            foundSoFar.set(typeName, { typeSpan: fullHoverResult, sourceFile: currentFile.slice(7) });
+            const identifiers = this.typeChecker.extractIdentifiers(fullHoverResult);
+            // console.dir(identifiers, { depth: null })
+            for (const identifier of identifiers) {
+                // console.log(`== loop ==`)
+                // console.dir(identifier, { depth: null })
+                // console.time(`loop ${identifier.name} layer ${layer}`)
+                // if (identifier.name === "_") {
+                //   console.timeEnd(`loop ${identifier.name} layer ${layer}`)
+                //   continue;
+                // };
+                // console.log(identifier)
+                // console.log(foundSoFar.has(identifier.name))
+                if (!foundSoFar.has(identifier.name)) {
+                    try {
+                        const typeDefinitionResult = await lspClient.typeDefinition({
+                            textDocument: {
+                                uri: currentFile
+                            },
+                            position: {
+                                character: identifier.start,
+                                line: startLine + identifier.line - 1 // startLine is already 1-indexed
+                            }
+                        });
+                        // if (identifier.name == "Model") {
+                        //   console.dir(typeDefinitionResult, { depth: null })
+                        // }
+                        if (typeDefinitionResult && typeDefinitionResult instanceof Array && typeDefinitionResult.length != 0) {
+                            const tdLocation = typeDefinitionResult[0];
+                            let content = "";
+                            if (foundContents.has(tdLocation.uri.slice(7))) {
+                                content = foundContents.get(tdLocation.uri.slice(7));
+                            }
+                            else {
+                                content = fs.readFileSync(tdLocation.uri.slice(7), "utf8");
+                                foundContents.set(tdLocation.uri.slice(7), content);
+                            }
+                            const decl = this.typeChecker.findDeclarationForIdentifier(content, tdLocation.range.start.line, tdLocation.range.start.character, tdLocation.range.end.character);
+                            if (decl) {
+                                const ident = this.typeChecker.getIdentifierFromDecl(decl);
+                                // console.log(ident == identifier.name, ident, identifier.name, decl)
+                                // console.log(`Decl: ${decl} || Identifier: ${ident}`)
+                                // console.timeEnd(`loop ${identifier.name} layer ${layer}`)
+                                await this.extractRelevantTypesHelper(lspClient, decl, ident, tdLocation.range.start.line, foundSoFar, tdLocation.uri, foundContents, layer + 1);
+                            }
+                            else {
+                                console.log("decl not found");
+                                // console.timeEnd(`loop ${identifier.name} layer ${layer}`)
+                            }
+                        }
+                        else {
+                            console.log("td not found");
+                            // console.dir(typeDefinitionResult, { depth: null })
+                        }
+                    }
+                    catch (err) {
+                        console.log(err);
+                    }
+                }
+                else {
+                    // console.timeEnd(`loop ${identifier.name} layer ${layer}`)
+                }
+            }
+        }
+    }
     async extractRelevantHeaders(_, sources, relevantTypes, holeType) {
+        // console.time("extractRelevantHeaders");
         const relevantContext = new Set();
+        // NOTE: This is necessary because TypeScript sucks.
+        // There is no way to compare objects by value,
+        // so sets of objects starts to accumulate tons of duplicates.
+        const relevantContextMap = new Map();
+        const trace = [];
+        const foundNormalForms = new Map();
         const targetTypes = this.generateTargetTypes(relevantTypes, holeType);
         // only consider lines that start with let or const
         for (const source of sources) {
@@ -262,24 +391,41 @@ class TypeScriptDriver {
             });
             // check for relationship between each line and relevant types
             filteredLines.forEach(line => {
+                // console.time(`helper, line: ${line}`);
+                let tag = false;
+                if (line === `const initFormState: [[Weekday, TimeOfDay], string] = [["M", "AM"], ""];`) {
+                    tag = true;
+                }
                 // TODO: Use the compiler API to split this.
                 const splittedLine = line.split(" = ")[0];
                 const typeSpanPattern = /(^[^:]*: )(.+)/;
                 const regexMatch = splittedLine.match(typeSpanPattern);
                 if (regexMatch) {
                     const returnTypeSpan = regexMatch[2];
+                    // console.log(`returnTypeSpan: ${returnTypeSpan}`)
+                    // const typeAnalysisResult = this.typeChecker.analyzeTypeString(returnTypeSpan);
+                    // console.log(`typeAnalysisResult: ${JSON.stringify(typeAnalysisResult, null, 2)}`)
                     if (!this.typeChecker.isPrimitive(returnTypeSpan.split(" => ")[1])) {
-                        this.extractRelevantHeadersHelper(returnTypeSpan, targetTypes, relevantTypes, relevantContext, splittedLine, source);
+                        this.extractRelevantHeadersHelper(returnTypeSpan, targetTypes, relevantTypes, relevantContext, splittedLine, source, relevantContextMap, tag, trace, foundNormalForms);
                     }
                 }
+                // console.timeEnd(`helper, line: ${line}`);
             });
         }
+        // console.log(JSON.stringify(relevantContextMap, null, 2))
+        // console.log(relevantContextMap.keys())
+        for (const v of relevantContextMap.values()) {
+            relevantContext.add(v);
+        }
+        // console.timeEnd("extractRelevantHeaders");
         return relevantContext;
     }
     generateTargetTypes(relevantTypes, holeType) {
+        // console.time("generateTargetTypes");
         const targetTypes = new Set();
         targetTypes.add(holeType);
         this.generateTargetTypesHelper(relevantTypes, holeType, targetTypes);
+        // console.timeEnd("generateTargetTypes");
         return targetTypes;
     }
     // generateTargetTypesHelper(
@@ -358,31 +504,62 @@ class TypeScriptDriver {
     //     }
     //   });
     // }
-    extractRelevantHeadersHelper(typeSpan, targetTypes, relevantTypes, relevantContext, line, source) {
+    extractRelevantHeadersHelper(typeSpan, targetTypes, relevantTypes, relevantContext, line, source, relevantContextMap, tag, trace, foundNormalForms) {
+        if (tag) {
+            // console.time(`extractRelevantHeadersHelper, typeSpan: ${typeSpan}`)
+            // trace.push(typeSpan)
+            // console.log(trace)
+        }
+        // TODO: this can probably done at the top level.
+        // analyzeTypeString is recursive by itself.
         const typeAnalysisResult = this.typeChecker.analyzeTypeString(typeSpan);
         targetTypes.forEach(typ => {
-            if (this.isTypeEquivalent(typeSpan, typ, relevantTypes)) {
-                relevantContext.add({ typeSpan: line, sourceFile: source });
+            if (this.isTypeEquivalent(typeSpan, typ, relevantTypes, foundNormalForms)) {
+                // NOTE: This checks for dupes. ctx is an object so you need to check for each field.
+                // relevantContext.add({ typeSpan: line, sourceFile: source });
+                const ctx = { typeSpan: line, sourceFile: source };
+                relevantContextMap.set(JSON.stringify(ctx), ctx);
             }
             if (this.typeChecker.isFunction2(typeAnalysisResult)) {
                 const rettype = typeAnalysisResult.returnType;
-                this.extractRelevantHeadersHelper(rettype.text, targetTypes, relevantTypes, relevantContext, line, source);
+                this.extractRelevantHeadersHelper(rettype.text, targetTypes, relevantTypes, relevantContext, line, source, relevantContextMap, tag, trace, foundNormalForms);
             }
             else if (this.typeChecker.isTuple2(typeAnalysisResult)) {
                 typeAnalysisResult.constituents.forEach(constituent => {
-                    this.extractRelevantHeadersHelper(constituent.text, targetTypes, relevantTypes, relevantContext, line, source);
+                    this.extractRelevantHeadersHelper(constituent.text, targetTypes, relevantTypes, relevantContext, line, source, relevantContextMap, tag, trace, foundNormalForms);
                 });
             }
         });
+        if (tag) {
+            // console.log("\n\n\n")
+            // console.timeEnd(`extractRelevantHeadersHelper, typeSpan: ${typeSpan}`)
+        }
     }
     // two types are equivalent if they have the same normal forms
-    isTypeEquivalent(t1, t2, relevantTypes) {
+    isTypeEquivalent(t1, t2, relevantTypes, foundNormalForms) {
         // NOTE: BUGFIX
         // console.log(`isTypeEquivalent: ${t1}, ${t2}`)
         // console.log(t1 == undefined)
         // console.log(t2 == undefined)
-        const normT1 = this.normalize(t1, relevantTypes);
-        const normT2 = this.normalize(t2, relevantTypes);
+        let normT1 = "";
+        let normT2 = "";
+        if (foundNormalForms.has(t1)) {
+            normT1 = foundNormalForms.get(t1);
+        }
+        else {
+            normT1 = this.normalize2(t1, relevantTypes);
+            foundNormalForms.set(t1, normT1);
+        }
+        if (foundNormalForms.has(t2)) {
+            normT2 = foundNormalForms.get(t2);
+        }
+        else {
+            normT2 = this.normalize2(t2, relevantTypes);
+            foundNormalForms.set(t2, normT2);
+        }
+        // const normT1 = foundNormalForms.has(t1) ? foundNormalForms.get(t1) : this.normalize2(t1, relevantTypes);
+        // const normT2 = foundNormalForms.has(t2) ? foundNormalForms.get(t2) : this.normalize2(t2, relevantTypes);
+        // console.log(normT1, normT2)
         return normT1 === normT2;
     }
     // return the normal form given a type span and a set of relevant types
@@ -390,9 +567,16 @@ class TypeScriptDriver {
     normalize(typeSpan, relevantTypes) {
         // NOTE: BUGFIX
         // console.log(`normalize: ${typeSpan}`)
+        // console.log(`normalize: ${typeSpan}`)
+        // console.log(`normalize: ${typeSpan == undefined}`)
         if (typeSpan.slice(typeSpan.length - 2) == " =") {
-            typeSpan = typeSpan.slice(typeSpan.length - 2);
+            typeSpan = typeSpan.slice(0, typeSpan.length - 2);
         }
+        // DEBUG
+        // if (typeSpan.slice(typeSpan.length - 1) == ";") {
+        //   typeSpan = typeSpan.slice(0, typeSpan.length - 1);
+        // }
+        // console.log(typeSpan)
         let normalForm = "";
         // pattern matching for typeSpan
         if (this.typeChecker.isPrimitive(typeSpan)) {
@@ -462,11 +646,17 @@ class TypeScriptDriver {
     normalize2(typeSpan, relevantTypes) {
         // NOTE: BUGFIX
         // console.log(`normalize: ${typeSpan}`)
+        // console.log(`normalize: ${typeSpan == undefined}`)
         if (typeSpan.slice(typeSpan.length - 2) == " =") {
-            typeSpan = typeSpan.slice(typeSpan.length - 2);
+            typeSpan = typeSpan.slice(0, typeSpan.length - 2);
+        }
+        if (typeSpan.slice(typeSpan.length - 1) == ";") {
+            typeSpan = typeSpan.slice(0, typeSpan.length - 1);
         }
+        // console.log(typeSpan)
         let normalForm = "";
         const analysisResult = this.typeChecker.analyzeTypeString(typeSpan);
+        // console.dir(analysisResult, { depth: null })
         // pattern matching for typeSpan
         // if (this.typeChecker.isPrimitive(typeSpan)) {
         if (this.typeChecker.isPrimitive2(analysisResult)) {
@@ -480,7 +670,7 @@ class TypeScriptDriver {
             elements.forEach(element => {
                 if (element !== "") {
                     const kv = element.split(": ");
-                    normalForm += kv[0].slice(1, kv[0].length), ": ", this.normalize(kv[1], relevantTypes);
+                    normalForm += kv[0].slice(1, kv[0].length), ": ", this.normalize2(kv[1], relevantTypes);
                     normalForm += "; ";
                 }
             });
@@ -494,7 +684,7 @@ class TypeScriptDriver {
             const elements = this.typeChecker.parseTypeArrayString(typeSpan);
             normalForm += "[";
             elements.forEach((element, i) => {
-                normalForm += this.normalize(element, relevantTypes);
+                normalForm += this.normalize2(element, relevantTypes);
                 if (i < elements.length - 1) {
                     normalForm += ", ";
                 }
@@ -508,7 +698,7 @@ class TypeScriptDriver {
             const elements = typeSpan.split(" | ");
             elements.forEach((element, i) => {
                 normalForm += "(";
-                normalForm += this.normalize(element, relevantTypes);
+                normalForm += this.normalize2(element, relevantTypes);
                 normalForm += ")";
                 if (i < elements.length - 1) {
                     normalForm += " | ";
@@ -520,17 +710,23 @@ class TypeScriptDriver {
         else if (this.typeChecker.isArray2(analysisResult)) {
             // console.log(`isArray: ${typeSpan}`)
             const element = typeSpan.split("[]")[0];
-            normalForm += this.normalize(element, relevantTypes);
+            normalForm += this.normalize2(element, relevantTypes);
             normalForm += "[]";
             return normalForm;
             // } else if (this.typeChecker.isTypeAlias(typeSpan)) {
         }
         else if (this.typeChecker.isTypeAlias2(analysisResult)) {
-            const typ = relevantTypes.get(typeSpan).typeSpan.split(" = ")[1];
+            // console.log("ALERT!!!!!!")
+            // console.dir(relevantTypes, { depth: null })
+            // console.dir(analysisResult, { depth: null })
+            // console.log("typeSpan:", typeSpan)
+            // console.log("analysis:", analysisResult.text)
+            // console.log(relevantTypes.get(analysisResult.text))
+            const typ = relevantTypes.get(analysisResult.text).typeSpan.split(" = ")[1];
             if (typ === undefined) {
                 return typeSpan;
             }
-            normalForm += this.normalize(typ, relevantTypes);
+            normalForm += this.normalize2(typ, relevantTypes);
             return normalForm;
         }
         else {

package/dist/src/typescript-type-checker.d.ts

@@ -61,4 +61,20 @@ export declare class TypeScriptTypeChecker implements TypeChecker {
     isUnion2(typeAnalysisResult: TypeAnalysis): boolean;
     isArray2(typeAnalysisResult: TypeAnalysis): boolean;
     isTypeAlias2(typeAnalysisResult: TypeAnalysis): boolean;
+    extractIdentifiers(code: string): {
+        name: string;
+        start: number;
+        end: number;
+        line: number;
+        column: number;
+    }[];
+    extractIdentifiersHelper(node: ts.Node, identifiers: Set<string>): void;
+    extractIdentifiersWithPosHelper(sourceFile: ts.SourceFile, node: ts.Node, identifiers: {
+        name: string;
+        start: number;
+        end: number;
+        line: number;
+        column: number;
+    }[]): void;
+    findDeclarationForIdentifier(sourceCode: string, targetLine: number, targetCharStart: number, targetCharEnd: number): string | null;
 }

package/dist/src/typescript-type-checker.js

@@ -386,6 +386,9 @@ class TypeScriptTypeChecker {
                             type: this.analyzeTypeNode(element.type, checker)
                         };
                     }
+                    else if (ts.isNamedTupleMember(element)) {
+                        return this.analyzeTypeNode(element.type, checker);
+                    }
                     else {
                         return this.analyzeTypeNode(element, checker);
                     }
@@ -508,5 +511,71 @@ class TypeScriptTypeChecker {
     isTypeAlias2(typeAnalysisResult) {
         return typeAnalysisResult.kind === "TypeReference";
     }
+    extractIdentifiers(code) {
+        const sourceFile = ts.createSourceFile("sample.ts", code, ts.ScriptTarget.Latest, true, ts.ScriptKind.TS);
+        const identifiers = [];
+        this.extractIdentifiersWithPosHelper(sourceFile, sourceFile, identifiers);
+        return identifiers;
+    }
+    extractIdentifiersHelper(node, identifiers) {
+        if (ts.isIdentifier(node)) {
+            // console.log(node.kind, node.text)
+            if (ts.isTypeReferenceNode(node.parent) ||
+                ts.isTypeAliasDeclaration(node.parent) ||
+                ts.isInterfaceDeclaration(node.parent) ||
+                ts.isClassDeclaration(node.parent) ||
+                ts.isFunctionTypeNode(node.parent)) {
+                identifiers.add(node.getText());
+            }
+        }
+        node.forEachChild(child => this.extractIdentifiersHelper(child, identifiers));
+    }
+    extractIdentifiersWithPosHelper(sourceFile, node, identifiers) {
+        if (ts.isIdentifier(node)) {
+            // console.log(node.kind, node.text)
+            if (ts.isTypeReferenceNode(node.parent) ||
+                ts.isTypeAliasDeclaration(node.parent) ||
+                ts.isInterfaceDeclaration(node.parent) ||
+                ts.isClassDeclaration(node.parent) ||
+                ts.isFunctionTypeNode(node.parent)) {
+                const { line, character } = sourceFile.getLineAndCharacterOfPosition(node.getStart());
+                identifiers.push({
+                    name: node.getText(),
+                    start: node.getStart(),
+                    end: node.getEnd(),
+                    line: line + 1, // Convert 0-based to 1-based
+                    column: character + 1 // Convert 0-based to 1-based
+                });
+            }
+        }
+        node.forEachChild(child => this.extractIdentifiersWithPosHelper(sourceFile, child, identifiers));
+    }
+    findDeclarationForIdentifier(sourceCode, targetLine, targetCharStart, targetCharEnd) {
+        const sourceFile = ts.createSourceFile("sample.ts", sourceCode, ts.ScriptTarget.Latest, true, ts.ScriptKind.TS);
+        function findDeclarationForIdentifierHelper(node) {
+            if (ts.isTypeAliasDeclaration(node)) {
+                const identifier = node.name;
+                // Get start and end positions for identifier
+                const startPos = sourceFile.getLineAndCharacterOfPosition(identifier.getStart());
+                const endPos = sourceFile.getLineAndCharacterOfPosition(identifier.getEnd());
+                // Match the identifier position
+                if (startPos.line === targetLine &&
+                    startPos.character === targetCharStart &&
+                    endPos.character === targetCharEnd) {
+                    // Extract full declaration text from source code
+                    return sourceCode.slice(node.getStart(), node.getEnd());
+                }
+            }
+            return ts.forEachChild(node, findDeclarationForIdentifierHelper) || null;
+        }
+        const declarationText = findDeclarationForIdentifierHelper(sourceFile);
+        // console.log(`declarationText: ${declarationText}`)
+        if (declarationText) {
+            return declarationText;
+        }
+        else {
+            return null;
+        }
+    }
 }
 exports.TypeScriptTypeChecker = TypeScriptTypeChecker;

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@jpoly1219/context-extractor",
-  "version": "0.2.6",
+  "version": "0.2.8",
   "description": "Extract relevant context from an incomplete program sketch.",
   "repository": {
     "type": "git",
@@ -12,7 +12,7 @@
     "dist/"
   ],
   "scripts": {
-    "build": "npx tsc && npm run copy-deps",
+    "build": "rm -rf dist/ && npx tsc && npm run copy-deps",
     "copy-deps": "npm run copy-lsp-client && npm run copy-ocaml",
     "copy-lsp-client": "shx cp -r ts-lsp-client-dist dist/",
     "copy-ocaml": "shx mkdir -p dist/src/ocaml-utils/ && shx cp -r src/ocaml-utils/_build dist/src/ocaml-utils/_build/",