decomp: auto generate docstring stub for methods/functions (#192)

src/extension.ts

@@ -13,12 +13,13 @@ import { activateTypeCastTools } from "./decomp/type-caster";
 import { IRInlayHintsProvider } from "./languages/ir2/ir2-inlay-hinter";
 import { OpenGOALDisasmRenameProvider } from "./languages/opengoal/disasm/opengoal-disasm-renamer";
 import { activateMiscDecompTools } from "./decomp/misc-tools";
-import { IR2RenameProvider } from "./languages/ir2/ir2-renamer";
+import { IRRenameProvider } from "./languages/ir2/ir2-renamer";
 import {
   onChangeSelection,
   onChangeTextDocument,
   registerParinferCommands,
 } from "./goal/parinfer/parinfer";
+import { IRCompletionItemProvider } from "./languages/ir2/ir2-completions";
 
 export async function activate(context: vscode.ExtensionContext) {
   try {
@@ -75,7 +76,12 @@ export async function activate(context: vscode.ExtensionContext) {
     );
     vscode.languages.registerRenameProvider(
       { scheme: "file", language: "opengoal-ir" },
-      new IR2RenameProvider()
+      new IRRenameProvider()
+    );
+    vscode.languages.registerCompletionItemProvider(
+      { scheme: "file", language: "opengoal-ir" },
+      new IRCompletionItemProvider(),
+      "@" // NOTE - can't use `"` without overriding a default setting https://github.com/microsoft/vscode/issues/131238#issuecomment-902519923
     );
 
     // Start the LSP

src/languages/ir2/ir2-completions.ts

@@ -0,0 +1,58 @@
+import * as vscode from "vscode";
+import { getArgumentsInSignature } from "../opengoal/opengoal-tools";
+
+export class IRCompletionItemProvider implements vscode.CompletionItemProvider {
+  provideCompletionItems(
+    document: vscode.TextDocument,
+    position: vscode.Position,
+    token: vscode.CancellationToken,
+    context: vscode.CompletionContext
+  ): vscode.ProviderResult<
+    vscode.CompletionItem[] | vscode.CompletionList<vscode.CompletionItem>
+  > {
+    // Currently, this is used to automatically generate docstrings for `defmethods` and `defun`
+    // NOTE - assumes single line signatures!
+
+    // Find the signature line
+    const prevLine = document.lineAt(position.line - 1).text;
+    if (!prevLine.includes("defmethod") && !prevLine.includes("defun")) {
+      return [];
+    }
+
+    const args = getArgumentsInSignature(prevLine);
+    if (args.length <= 0) {
+      return [];
+    }
+
+    const range = new vscode.Range(position, position);
+    const rangeToRemove = new vscode.Range(
+      position.line,
+      position.character - 1,
+      position.line,
+      position.character
+    );
+
+    let docstring = `"something\n`;
+    for (const arg of args) {
+      docstring += ` @param ${arg.name} something\n`;
+    }
+    docstring += ` @returns something"`;
+
+    return [
+      {
+        label: "Auto-Generated Docstring",
+        kind: vscode.CompletionItemKind.Text,
+        range: range,
+        insertText: docstring,
+        additionalTextEdits: [vscode.TextEdit.delete(rangeToRemove)],
+      },
+    ];
+  }
+
+  resolveCompletionItem?(
+    item: vscode.CompletionItem,
+    token: vscode.CancellationToken
+  ): vscode.ProviderResult<vscode.CompletionItem> {
+    throw new Error("Method not implemented.");
+  }
+}

src/languages/ir2/ir2-renamer.ts

@@ -4,7 +4,7 @@ import { getSymbolAtPosition } from "../common/utils";
 import { getSymbolsArgumentInfo } from "../opengoal/opengoal-tools";
 import { getFuncNameFromPosition, insideGoalCodeInIR } from "./ir2-utils";
 
-export class IR2RenameProvider implements vscode.RenameProvider {
+export class IRRenameProvider implements vscode.RenameProvider {
   public async provideRenameEdits(
     document: vscode.TextDocument,
     position: vscode.Position,

src/languages/opengoal/opengoal-tools.ts

@@ -6,6 +6,43 @@ export interface ArgumentMeta {
   isMethod: boolean;
 }
 
+export interface ArgumentDefinition {
+  name: string;
+  type: string;
+}
+
+export function getArgumentsInSignature(
+  signature: string
+): ArgumentDefinition[] {
+  const isArgument =
+    signature.includes("defun") ||
+    signature.includes("defmethod") ||
+    signature.includes("defbehavior");
+
+  if (!isArgument) {
+    return [];
+  }
+
+  const matches = [...signature.matchAll(/(\([^\s(]*\s[^\s)]*\))/g)];
+  if (matches.length == 0) {
+    return [];
+  }
+
+  const args: ArgumentDefinition[] = [];
+  for (const match of matches) {
+    const [argName, argType] = match[1]
+      .toString()
+      .replace("(", "")
+      .replace(")", "")
+      .split(" ");
+    args.push({
+      name: argName,
+      type: argType,
+    });
+  }
+  return args;
+}
+
 // This function can only currently figure out arguments if they are on the same line as
 // a defun/defmethod/etc
 //
@@ -15,35 +52,19 @@ export function getSymbolsArgumentInfo(
   line: string,
   symbol: string
 ): ArgumentMeta | undefined {
-  const isArgument =
-    line.includes("defun") ||
-    line.includes("defmethod") ||
-    line.includes("defbehavior");
   // TODO - 'new' method handling
-  const isMethod = line.includes("defmethod"); // if it's a method, make the first arg be `obj`
   // If it's an argument, we have to figure out the index
   let argumentIndex = undefined;
-  let argumentCount = undefined;
+  const args = getArgumentsInSignature(line);
-  if (isArgument) {
+  const argumentCount = args.length;
-    const matches = [...line.matchAll(/(\([^\s(]*\s[^\s)]*\))/g)];
+  if (argumentCount > 0) {
-    if (matches.length == 0) {
+    for (let i = 0; i < args.length; i++) {
-      return undefined;
+      if (args[i].name === symbol) {
-    }
+        argumentIndex = i;
-    let tempIdx = 0;
-    for (const match of matches) {
-      const [argName, argType] = match[1]
-        .toString()
-        .replace("(", "")
-        .replace(")", "")
-        .split(" ");
-      if (argName === symbol) {
-        argumentIndex = tempIdx;
-        argumentCount = matches.length;
         break;
       }
-      tempIdx++;
     }
-    if (argumentIndex === undefined || argumentCount === undefined) {
+    if (argumentIndex === undefined) {
       return undefined;
     }
   } else {
@@ -52,7 +73,7 @@ export function getSymbolsArgumentInfo(
   return {
     index: argumentIndex,
     totalCount: argumentCount,
-    isMethod: isMethod,
+    isMethod: line.includes("defmethod"), // if it's a method, make the first arg be `obj`
   };
 }