decomp: allow renaming variables and arguments inside the IR2 file's OpenGOAL output

src/decomp/decomp-tools.ts

@@ -1,5 +1,5 @@
 import { exec, execFile } from "child_process";
-import { existsSync, promises as fs, readFileSync } from "fs";
+import { existsSync, promises as fs } from "fs";
 import * as vscode from "vscode";
 import { determineGameFromPath, GameName, openFile } from "../utils/file-utils";
 import { open_in_pdf } from "./man-page";

src/decomp/type-caster.ts

@@ -3,7 +3,8 @@ import * as vscode from "vscode";
 import { basename, join } from "path";
 import { readFileSync, writeFileSync } from "fs";
 import { parse, stringify } from "comment-json";
-import { getDecompilerConfigDirectory } from "../utils/decomp-tools";
+import { getFuncNameFromSelection } from "../languages/ir2/ir2-utils";
+import { getDecompilerConfigDirectory } from "./utils";
 
 enum CastKind {
   Label,
@@ -20,7 +21,6 @@ let lastTypeCastType: string | undefined;
 
 const opNumRegex = /.*;; \[\s*(\d+)\]/g;
 const registerRegex = /[a|s|t|v|f]\d+|gp|fp|r0|ra/g;
-const funcNameRegex = /; \.function (.*).*/g;
 const stackOffsetRegex = /sp, (\d+)/g;
 const labelRefRegex = /(L\d+).*;;/g;
 
@@ -42,23 +42,6 @@ async function getOpNumber(line: string): Promise<number | undefined> {
   return undefined;
 }
 
-async function getFuncName(
-  document: vscode.TextDocument,
-  selection: vscode.Selection
-): Promise<string | undefined> {
-  for (let i = selection.start.line; i >= 0; i--) {
-    const line = document.lineAt(i).text;
-    const matches = [...line.matchAll(funcNameRegex)];
-    if (matches.length == 1) {
-      return matches[0][1].toString();
-    }
-  }
-  await vscode.window.showErrorMessage(
-    "Couldn't determine function or method name"
-  );
-  return undefined;
-}
-
 async function getLabelReference(line: string): Promise<string | undefined> {
   const matches = [...line.matchAll(labelRefRegex)];
   if (matches.length == 1) {
@@ -203,7 +186,10 @@ async function stackCastSelection() {
   }
 
   // Get the relevant function/method name
-  const funcName = await getFuncName(editor.document, editor.selection);
+  const funcName = await getFuncNameFromSelection(
+    editor.document,
+    editor.selection
+  );
   if (funcName === undefined) {
     return;
   }
@@ -312,7 +298,10 @@ async function typeCastSelection() {
   }
 
   // Get the relevant function/method name
-  const funcName = await getFuncName(editor.document, editor.selection);
+  const funcName = await getFuncNameFromSelection(
+    editor.document,
+    editor.selection
+  );
   if (funcName === undefined) {
     return;
   }
@@ -387,7 +376,10 @@ async function repeatLastCast() {
       lastLabelCastSize
     );
   } else if (lastCastKind === CastKind.Stack) {
-    const funcName = await getFuncName(editor.document, editor.selection);
+    const funcName = await getFuncNameFromSelection(
+      editor.document,
+      editor.selection
+    );
     if (funcName === undefined) {
       return;
     }
@@ -401,7 +393,10 @@ async function repeatLastCast() {
     }
     await applyStackCast(editor, funcName, stackOffset, lastStackCastType);
   } else if (lastCastKind === CastKind.TypeCast) {
-    const funcName = await getFuncName(editor.document, editor.selection);
+    const funcName = await getFuncNameFromSelection(
+      editor.document,
+      editor.selection
+    );
     if (funcName === undefined) {
       return;
     }

src/utils/decomp-tools.ts → src/decomp/utils.ts

@@ -1,17 +1,19 @@
-import { parse } from "comment-json";
-import { existsSync, readFileSync } from "fs";
+import { parse, stringify } from "comment-json";
+import { existsSync, readFileSync, writeFileSync } from "fs";
+import { join } from "path";
 import * as vscode from "vscode";
 import {
   getConfig,
   updateJak1DecompConfigDirectory,
   updateJak2DecompConfigDirectory,
 } from "../config/config";
+import { ArgumentMeta } from "../languages/opengoal/opengoal-tools";
 import {
   determineGameFromPath,
   GameName,
   getDirectoriesInDir,
-} from "./file-utils";
-import { getWorkspaceFolderByName } from "./workspace";
+} from "../utils/file-utils";
+import { getWorkspaceFolderByName } from "../utils/workspace";
 
 export function getCastFileData(
   projectRoot: vscode.Uri,
@@ -45,7 +47,6 @@ export function getCastFileData(
     return undefined;
   }
 
-  // TODO - would be performant to cache these files, requires listening to them as well though
   return parse(readFileSync(path).toString(), undefined, true);
 }
 
@@ -132,3 +133,81 @@ export async function getDecompilerConfigDirectory(
     }
   }
 }
+
+export async function updateVarCasts(
+  document: vscode.TextDocument,
+  funcName: string,
+  argMeta: ArgumentMeta | undefined,
+  currSymbol: string,
+  newName: 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] = {};
+  }
+  if (argMeta) {
+    if (argMeta.index === undefined || argMeta.totalCount === undefined) {
+      return;
+    }
+    if ("args" in varNameData[funcName]) {
+      // We assume that all slots are filled up already, as this is required
+      varNameData[funcName].args[argMeta.index] = newName;
+    } else {
+      // Otherwise, we initialize it properly
+      varNameData[funcName].args = [];
+      for (let i = 0; i < argMeta.totalCount; i++) {
+        if (i == argMeta.index) {
+          varNameData[funcName].args[i] = newName;
+        } else {
+          if (argMeta.isMethod && i == 0) {
+            varNameData[funcName].args[i] = "obj";
+          } else {
+            varNameData[funcName].args[i] = `arg${i}`;
+          }
+        }
+      }
+    }
+  } else {
+    // if "vars" is in
+    if ("vars" in varNameData[funcName]) {
+      // Check to see if the current symbol has already been renamed
+      let internalVar = undefined;
+      for (const [key, value] of Object.entries(varNameData[funcName].vars)) {
+        if (value === currSymbol) {
+          internalVar = key;
+          break;
+        }
+      }
+      if (internalVar !== undefined) {
+        varNameData[funcName].vars[internalVar] = newName;
+      } else {
+        varNameData[funcName].vars[currSymbol] = newName;
+      }
+    } else {
+      varNameData[funcName]["vars"] = {};
+      varNameData[funcName]["vars"][currSymbol] = 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/extension.ts

@@ -13,6 +13,7 @@ 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";
 
 export async function activate(context: vscode.ExtensionContext) {
   try {
@@ -67,6 +68,10 @@ export async function activate(context: vscode.ExtensionContext) {
       { scheme: "file", language: "opengoal", pattern: "**/*_disasm.gc" },
       new OpenGOALDisasmRenameProvider()
     );
+    vscode.languages.registerRenameProvider(
+      { scheme: "file", language: "opengoal-ir" },
+      new IR2RenameProvider()
+    );
 
     // Start the LSP
     lsp.activate(context);

src/languages/common/utils.ts

@@ -0,0 +1,12 @@
+import * as vscode from "vscode";
+
+export function getSymbolAtPosition(
+  document: vscode.TextDocument,
+  position: vscode.Position
+) {
+  const symbolRange = document.getWordRangeAtPosition(position, /[^\s()]+/g);
+  if (symbolRange === undefined) {
+    return;
+  }
+  return document.getText(symbolRange);
+}

src/languages/ir2/ir2-inlay-hinter.ts

@@ -1,6 +1,6 @@
 import path = require("path");
 import * as vscode from "vscode";
-import { getCastFileData } from "../../utils/decomp-tools";
+import { getCastFileData } from "../../decomp/utils";
 import { getWorkspaceFolderByName } from "../../utils/workspace";
 
 class HintCacheEntry {

src/languages/ir2/ir2-renamer.ts

@@ -0,0 +1,39 @@
+import * as vscode from "vscode";
+import { updateVarCasts } from "../../decomp/utils";
+import { getSymbolAtPosition } from "../common/utils";
+import { getSymbolsArgumentInfo } from "../opengoal/opengoal-tools";
+import { getFuncNameFromPosition, insideGoalCodeInIR } from "./ir2-utils";
+
+export class IR2RenameProvider implements vscode.RenameProvider {
+  public async provideRenameEdits(
+    document: vscode.TextDocument,
+    position: vscode.Position,
+    newName: string,
+    token: vscode.CancellationToken
+  ): Promise<vscode.WorkspaceEdit | undefined> {
+    const symbol = getSymbolAtPosition(document, position);
+    if (symbol === undefined) {
+      return;
+    }
+
+    // Check that we are actually inside an OpenGOAL function
+    if (!insideGoalCodeInIR(document, position)) {
+      return;
+    }
+
+    const line = document.lineAt(position.line).text;
+    const argMeta = getSymbolsArgumentInfo(line, symbol);
+    // 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(document, position);
+    if (funcName === undefined) {
+      return;
+    }
+
+    await updateVarCasts(document, funcName, argMeta, symbol, newName);
+
+    // The actual renaming is done by the decompiler, which happens automatically if
+    // auto decompilation is enabled
+    return;
+  }
+}

src/languages/ir2/ir2-utils.ts

@@ -0,0 +1,47 @@
+import * as vscode from "vscode";
+
+// Checks if we are inside the decompiler's output within the IR2 file
+// This is somewhat guaranteed by the fact that this is how the embded syntax highlighting works
+export function insideGoalCodeInIR(
+  document: vscode.TextDocument,
+  position: vscode.Position
+): boolean {
+  // Somewhat primitive, walk back until we find a `;;-*-OpenGOAL-Start-*-` before we find a `.function`
+  let idx = position.line;
+  while (idx > 0) {
+    const line = document.lineAt(idx).text;
+    if (line.includes(";;-*-OpenGOAL-Start-*-")) {
+      return true;
+    }
+    if (line.includes(".function")) {
+      return false;
+    }
+    idx--;
+  }
+  return false;
+}
+
+export async function getFuncNameFromPosition(
+  document: vscode.TextDocument,
+  position: vscode.Position
+): Promise<string | undefined> {
+  const funcNameRegex = /; \.function (.*).*/g;
+  for (let i = position.line; i >= 0; i--) {
+    const line = document.lineAt(i).text;
+    const matches = [...line.matchAll(funcNameRegex)];
+    if (matches.length == 1) {
+      return matches[0][1].toString();
+    }
+  }
+  await vscode.window.showErrorMessage(
+    "Couldn't determine function or method name"
+  );
+  return undefined;
+}
+
+export async function getFuncNameFromSelection(
+  document: vscode.TextDocument,
+  selection: vscode.Selection
+): Promise<string | undefined> {
+  return await getFuncNameFromPosition(document, selection.start);
+}