decomp: add new feature that lets you bulk rename variables in a selection (#332)

package.json

@@ -135,6 +135,10 @@
         "command": "opengoal.decomp.misc.applyDecompilerSuggestions",
         "title": "OpenGOAL - Misc - Apply Decompiler Suggestions to Selection"
       },
+      {
+        "command": "opengoal.decomp.misc.batchRenameUnnamedVars",
+        "title": "OpenGOAL - Misc - Batch Rename Unnamed Vars"
+      },
       {
         "command": "opengoal.decomp.typeSearcher.open",
         "title": "OpenGOAL - Misc - Type Searcher"
@@ -414,6 +418,11 @@
           "command": "opengoal.decomp.openManPage",
           "group": "z_commands"
         }
+      ],
+      "editor/title": [
+        {
+          "when": "resourceScheme == opengoalBatchRename"
+        }
       ]
     },
     "languages": [

src/decomp/misc-tools.ts

@@ -7,6 +7,13 @@ import {
   updateFileBeforeDecomp,
 } from "../utils/file-utils";
 import { getWorkspaceFolderByName } from "../utils/workspace";
+import { getFuncNameFromPosition } from "../languages/ir2/ir2-utils";
+import {
+  ArgumentMeta,
+  getArgumentsInSignature,
+  getSymbolsArgumentInfo,
+} from "../languages/opengoal/opengoal-tools";
+import { bulkUpdateVarCasts } from "./utils";
 
 async function addToOffsets() {
   const editor = vscode.window.activeTextEditor;
@@ -455,7 +462,118 @@ async function applyDecompilerSuggestions() {
   });
 }
 
+let originalDocumentForRename: vscode.TextDocument | undefined = undefined;
+let currentRenameWindow: vscode.TextEditor | undefined = undefined;
+let currentRenameLines: string[] = [];
+let currentRenameFunctionName: string | undefined = undefined;
+let currentRenameArgMeta: ArgumentMeta | undefined = undefined;
+let currentRenameFileVersion = 0;
+
+async function batchRenameUnnamedVars() {
+  const editor = vscode.window.activeTextEditor;
+  if (editor === undefined || editor.selection.isEmpty) {
+    return;
+  }
+  const currentSelection = editor.document.getText(editor.selection);
+
+  // We can determine the function in a more consistent way here, that will also allow
+  // for renaming anon-functions / states / etc
+  const funcName = await getFuncNameFromPosition(
+    editor.document,
+    editor.selection.active,
+  );
+  if (funcName === undefined) {
+    return;
+  }
+  currentRenameFunctionName = funcName;
+  currentRenameArgMeta = {
+    index: 0,
+    totalCount: getArgumentsInSignature(currentSelection.split("\n")[0]).length,
+    isMethod: currentSelection.split("\n")[0].includes("defmethod"),
+  };
+
+  const unnamedVarRegex =
+    /(?:(?:arg\d+)|(?:f\d+|at|v[0-1]|a[0-3]|t[0-9]|s[0-7]|k[0-1]|gp|sp|sv|fp|ra)-\d+)/g;
+
+  const vars = new Set(
+    [...currentSelection.matchAll(unnamedVarRegex)].map((match) => match[0]),
+  );
+
+  currentRenameLines = [];
+  currentRenameLines.push(`Renaming Vars in - ${funcName}:`);
+  for (const variable of vars) {
+    currentRenameLines.push(`${variable} => `);
+  }
+
+  originalDocumentForRename = editor.document;
+  currentRenameFileVersion++;
+  currentRenameWindow = await vscode.window.showTextDocument(
+    vscode.Uri.from({
+      scheme: "opengoalBatchRename",
+      path: "/opengoalBatchRename",
+    }),
+    { preview: false, viewColumn: vscode.ViewColumn.Beside },
+  );
+}
+
+async function processOpengoalBatchRename() {
+  if (
+    originalDocumentForRename === undefined ||
+    currentRenameFunctionName === undefined ||
+    currentRenameArgMeta === undefined
+  ) {
+    return;
+  }
+
+  const renameMap: Record<string, string> = {};
+
+  for (let i = 0; i < currentRenameLines.length; i++) {
+    const tokens = currentRenameLines[i].split("=>");
+    if (tokens.length !== 2) {
+      continue;
+    }
+    const oldName = tokens[0].trim();
+    const newName = tokens[1].trim();
+    renameMap[oldName] = newName;
+  }
+
+  await bulkUpdateVarCasts(
+    originalDocumentForRename,
+    currentRenameFunctionName,
+    currentRenameArgMeta,
+    renameMap,
+  );
+  await vscode.commands.executeCommand(
+    "workbench.action.revertAndCloseActiveEditor",
+  );
+}
+
 export async function activateMiscDecompTools() {
+  const emitter = new vscode.EventEmitter<vscode.FileChangeEvent[]>();
+  vscode.workspace.registerFileSystemProvider("opengoalBatchRename", {
+    createDirectory() {},
+    delete() {},
+    onDidChangeFile: emitter.event,
+    readDirectory() {
+      return [];
+    },
+    readFile() {
+      return new TextEncoder().encode(currentRenameLines.join("\n"));
+    },
+    rename() {},
+    stat() {
+      return { ctime: 0, mtime: currentRenameFileVersion, size: 0, type: 0 };
+    },
+    watch(uri) {
+      return new vscode.Disposable(() => {});
+    },
+    writeFile(uri, content) {
+      currentRenameLines = new TextDecoder().decode(content).split("\n");
+      processOpengoalBatchRename();
+      currentRenameFileVersion++;
+    },
+  });
+
   getExtensionContext().subscriptions.push(
     vscode.commands.registerCommand(
       "opengoal.decomp.misc.addToOffsets",
@@ -504,4 +622,10 @@ export async function activateMiscDecompTools() {
       applyDecompilerSuggestions,
     ),
   );
+  getExtensionContext().subscriptions.push(
+    vscode.commands.registerCommand(
+      "opengoal.decomp.misc.batchRenameUnnamedVars",
+      batchRenameUnnamedVars,
+    ),
+  );
 }

src/decomp/utils.ts

@@ -144,7 +144,7 @@ export async function updateVarCasts(
           }
         } else {
           if (argMeta.isMethod && i == 0) {
-            varNameData[funcName].args[i] = "obj";
+            varNameData[funcName].args[i] = "this";
           } else {
             varNameData[funcName].args[i] = `arg${i}`;
           }
@@ -215,3 +215,64 @@ export async function updateVarCasts(
 
   writeFileSync(filePath, stringify(varNameData, null, 2));
 }
+
+export async function bulkUpdateVarCasts(
+  document: vscode.TextDocument,
+  funcName: string,
+  argMeta: ArgumentMeta,
+  renameMap: Record<string, string>,
+) {
+  // Update the var-names file
+  const projectRoot = getWorkspaceFolderByName("jak-project");
+  if (projectRoot === undefined) {
+    vscode.window.showErrorMessage(
+      "OpenGOAL - Unable to locate 'jak-project' workspace folder",
+    );
+    return;
+  }
+
+  const varNameData = getCastFileData(projectRoot, document, "var_names.jsonc");
+  if (varNameData === undefined) {
+    return;
+  }
+
+  if (!(funcName in varNameData)) {
+    varNameData[funcName] = {};
+  }
+
+  for (const [oldName, newName] of Object.entries(renameMap)) {
+    if (oldName.startsWith("arg")) {
+      // initialize if not already done
+      if (!("args" in varNameData[funcName])) {
+        varNameData[funcName].args = [];
+        for (let i = 0; i < argMeta.totalCount; i++) {
+          if (argMeta.isMethod && i == 0) {
+            varNameData[funcName].args[i] = "this";
+          } else {
+            varNameData[funcName].args[i] = `arg${i}`;
+          }
+        }
+      }
+      let argIndex = parseInt(oldName.substring(3));
+      if (argMeta.isMethod) {
+        argIndex++;
+      }
+      varNameData[funcName].args[argIndex] = newName;
+    } else {
+      if (!("vars" in varNameData[funcName])) {
+        varNameData[funcName].vars = {};
+      }
+      // NOTE - omitting check for duplicate names, just know what you're doing
+      varNameData[funcName].vars[oldName] = newName;
+    }
+  }
+
+  // Write out cast file change
+  const configDir = await getDecompilerConfigDirectory(document.uri);
+  if (configDir === undefined) {
+    return;
+  }
+  const filePath = join(configDir, "var_names.jsonc");
+
+  writeFileSync(filePath, stringify(varNameData, null, 2));
+}

src/languages/opengoal/opengoal-tools.ts

@@ -14,12 +14,12 @@ export interface ArgumentDefinition {
 export function getArgumentsInSignature(
   signature: string,
 ): ArgumentDefinition[] {
-  const isArgument =
+  const isSignature =
     signature.includes("defun") ||
     signature.includes("defmethod") ||
     signature.includes("defbehavior");
 
-  if (!isArgument) {
+  if (!isSignature) {
     return [];
   }