Merge pull request #3017 from hmac/hmac-model-editor-ruby

Add experimental model editor support for Ruby

Author: koesie10

extensions/ql-vscode/src/config.ts

@@ -707,12 +707,14 @@ const LLM_GENERATION_BATCH_SIZE = new Setting(
   MODEL_SETTING,
 );
 const EXTENSIONS_DIRECTORY = new Setting("extensionsDirectory", MODEL_SETTING);
+const ENABLE_RUBY = new Setting("enableRuby", MODEL_SETTING);
 
 export interface ModelConfig {
   flowGeneration: boolean;
   llmGeneration: boolean;
   getExtensionsDirectory(languageId: string): string | undefined;
   showMultipleModels: boolean;
+  enableRuby: boolean;
 }
 
 export class ModelConfigListener extends ConfigListener implements ModelConfig {
@@ -745,4 +747,8 @@ export class ModelConfigListener extends ConfigListener implements ModelConfig {
   public get showMultipleModels(): boolean {
     return isCanary();
   }
+
+  public get enableRuby(): boolean {
+    return !!ENABLE_RUBY.getValue<boolean>();
+  }
 }

extensions/ql-vscode/src/model-editor/languages/languages.ts

@@ -1,10 +1,12 @@
 import { QueryLanguage } from "../../common/query-language";
 import { ModelsAsDataLanguage } from "./models-as-data";
+import { ruby } from "./ruby";
 import { staticLanguage } from "./static";
 
 const languages: Partial<Record<QueryLanguage, ModelsAsDataLanguage>> = {
   [QueryLanguage.CSharp]: staticLanguage,
   [QueryLanguage.Java]: staticLanguage,
+  [QueryLanguage.Ruby]: ruby,
 };
 
 export function getModelsAsDataLanguage(

extensions/ql-vscode/src/model-editor/languages/models-as-data.ts

@@ -1,6 +1,7 @@
 import { MethodDefinition } from "../method";
 import { ModeledMethod, ModeledMethodType } from "../modeled-method";
 import { DataTuple } from "../model-extension-file";
+import { Mode } from "../shared/mode";
 
 type GenerateMethodDefinition = (method: ModeledMethod) => DataTuple[];
 type ReadModeledMethod = (row: DataTuple[]) => ModeledMethod;
@@ -20,6 +21,11 @@ export type ModelsAsDataLanguagePredicates = Record<
 >;
 
 export type ModelsAsDataLanguage = {
+  /**
+   * The modes that are available for this language. If not specified, all
+   * modes are available.
+   */
+  availableModes?: Mode[];
   createMethodSignature: (method: MethodDefinition) => string;
   predicates: ModelsAsDataLanguagePredicates;
 };

extensions/ql-vscode/src/model-editor/languages/ruby.ts

@@ -0,0 +1,153 @@
+import { ModelsAsDataLanguage } from "./models-as-data";
+import { sharedExtensiblePredicates, sharedKinds } from "./shared";
+import { Mode } from "../shared/mode";
+
+function parseRubyMethodFromPath(path: string): string {
+  const match = path.match(/Method\[([^\]]+)].*/);
+  if (match) {
+    return match[1];
+  } else {
+    return "";
+  }
+}
+
+function parseRubyAccessPath(path: string): {
+  methodName: string;
+  path: string;
+} {
+  const match = path.match(/Method\[([^\]]+)]\.(.*)/);
+  if (match) {
+    return { methodName: match[1], path: match[2] };
+  } else {
+    return { methodName: "", path: "" };
+  }
+}
+
+function rubyMethodSignature(typeName: string, methodName: string) {
+  return `${typeName}#${methodName}`;
+}
+
+export const ruby: ModelsAsDataLanguage = {
+  availableModes: [Mode.Framework],
+  createMethodSignature: ({ typeName, methodName }) =>
+    `${typeName}#${methodName}`,
+  predicates: {
+    source: {
+      extensiblePredicate: sharedExtensiblePredicates.source,
+      supportedKinds: sharedKinds.source,
+      // extensible predicate sourceModel(
+      //   string type, string path, string kind
+      // );
+      generateMethodDefinition: (method) => [
+        method.typeName,
+        `Method[${method.methodName}].${method.output}`,
+        method.kind,
+      ],
+      readModeledMethod: (row) => {
+        const typeName = row[0] as string;
+        const { methodName, path: output } = parseRubyAccessPath(
+          row[1] as string,
+        );
+        return {
+          type: "source",
+          input: "",
+          output,
+          kind: row[2] as string,
+          provenance: "manual",
+          signature: rubyMethodSignature(typeName, methodName),
+          packageName: "",
+          typeName,
+          methodName,
+          methodParameters: "",
+        };
+      },
+    },
+    sink: {
+      extensiblePredicate: sharedExtensiblePredicates.sink,
+      supportedKinds: sharedKinds.sink,
+      // extensible predicate sinkModel(
+      //   string type, string path, string kind
+      // );
+      generateMethodDefinition: (method) => {
+        const path = `Method[${method.methodName}].${method.input}`;
+        return [method.typeName, path, method.kind];
+      },
+      readModeledMethod: (row) => {
+        const typeName = row[0] as string;
+        const { methodName, path: input } = parseRubyAccessPath(
+          row[1] as string,
+        );
+        return {
+          type: "sink",
+          input,
+          output: "",
+          kind: row[2] as string,
+          provenance: "manual",
+          signature: rubyMethodSignature(typeName, methodName),
+          packageName: "",
+          typeName,
+          methodName,
+          methodParameters: "",
+        };
+      },
+    },
+    summary: {
+      extensiblePredicate: sharedExtensiblePredicates.summary,
+      supportedKinds: sharedKinds.summary,
+      // extensible predicate summaryModel(
+      //   string type, string path, string input, string output, string kind
+      // );
+      generateMethodDefinition: (method) => [
+        method.typeName,
+        `Method[${method.methodName}]`,
+        method.input,
+        method.output,
+        method.kind,
+      ],
+      readModeledMethod: (row) => {
+        const typeName = row[0] as string;
+        const methodName = parseRubyMethodFromPath(row[1] as string);
+        return {
+          type: "summary",
+          input: row[2] as string,
+          output: row[3] as string,
+          kind: row[4] as string,
+          provenance: "manual",
+          signature: rubyMethodSignature(typeName, methodName),
+          packageName: "",
+          typeName,
+          methodName,
+          methodParameters: "",
+        };
+      },
+    },
+    neutral: {
+      extensiblePredicate: sharedExtensiblePredicates.neutral,
+      supportedKinds: sharedKinds.neutral,
+      // extensible predicate neutralModel(
+      //   string type, string path, string kind
+      // );
+      generateMethodDefinition: (method) => [
+        method.typeName,
+        `Method[${method.methodName}]`,
+        method.kind,
+      ],
+      readModeledMethod: (row) => {
+        const typeName = row[0] as string;
+        const methodName = parseRubyMethodFromPath(row[1] as string);
+        return {
+          type: "neutral",
+          input: "",
+          output: "",
+          kind: row[2] as string,
+          provenance: "manual",
+          signature: rubyMethodSignature(typeName, methodName),
+          packageName: "",
+          typeName,
+          methodName,
+          methodParameters: "",
+        };
+      },
+    },
+  },
+};

extensions/ql-vscode/src/model-editor/model-editor-module.ts

@@ -22,8 +22,9 @@ import { showResolvableLocation } from "../databases/local-databases/locations";
 import { ModelEditorViewTracker } from "./model-editor-view-tracker";
 import { ModelConfigListener } from "../config";
 import { ModelingEvents } from "./modeling-events";
-
-const SUPPORTED_LANGUAGES: string[] = ["java", "csharp"];
+import { getModelsAsDataLanguage } from "./languages";
+import { INITIAL_MODE } from "./shared/mode";
+import { isSupportedLanguage } from "./supported-languages";
 
 export class ModelEditorModule extends DisposableObject {
   private readonly queryStorageDir: string;
@@ -32,6 +33,7 @@ export class ModelEditorModule extends DisposableObject {
   private readonly editorViewTracker: ModelEditorViewTracker<ModelEditorView>;
   private readonly methodsUsagePanel: MethodsUsagePanel;
   private readonly methodModelingPanel: MethodModelingPanel;
+  private readonly modelConfig: ModelConfigListener;
 
   private constructor(
     private readonly app: App,
@@ -56,6 +58,7 @@ export class ModelEditorModule extends DisposableObject {
         this.editorViewTracker,
       ),
     );
+    this.modelConfig = this.push(new ModelConfigListener());
 
     this.registerToModelingEvents();
   }
@@ -125,9 +128,10 @@ export class ModelEditorModule extends DisposableObject {
       }
 
       const language = db.language;
+
       if (
-        !SUPPORTED_LANGUAGES.includes(language) ||
-        !isQueryLanguage(language)
+        !isQueryLanguage(language) ||
+        !isSupportedLanguage(language, this.modelConfig)
       ) {
         void showAndLogErrorMessage(
           this.app.logger,
@@ -136,6 +140,10 @@ export class ModelEditorModule extends DisposableObject {
         return;
       }
 
+      const definition = getModelsAsDataLanguage(language);
+
+      const initialMode = definition.availableModes?.[0] ?? INITIAL_MODE;
+
       const existingView = this.editorViewTracker.getView(
         db.databaseUri.toString(),
       );
@@ -167,12 +175,10 @@ export class ModelEditorModule extends DisposableObject {
             return;
           }
 
-          const modelConfig = this.push(new ModelConfigListener());
-
           const modelFile = await pickExtensionPack(
             this.cliServer,
             db,
-            modelConfig,
+            this.modelConfig,
             this.app.logger,
             progress,
             maxStep,
@@ -196,7 +202,7 @@ export class ModelEditorModule extends DisposableObject {
             this.cliServer,
             queryDir,
             language,
-            modelConfig,
+            this.modelConfig,
           );
           if (!success) {
             await cleanupQueryDir();
@@ -225,7 +231,7 @@ export class ModelEditorModule extends DisposableObject {
             this.modelingStore,
             this.modelingEvents,
             this.editorViewTracker,
-            modelConfig,
+            this.modelConfig,
             this.databaseManager,
             this.cliServer,
             this.queryRunner,
@@ -234,6 +240,7 @@ export class ModelEditorModule extends DisposableObject {
             db,
             modelFile,
             language,
+            initialMode,
           );
 
           this.modelingEvents.onDbClosed(async (dbUri) => {

extensions/ql-vscode/src/model-editor/model-editor-queries.ts

@@ -10,7 +10,7 @@ import { redactableError } from "../common/errors";
 import { telemetryListener } from "../common/vscode/telemetry";
 import { join } from "path";
 import { Mode } from "./shared/mode";
-import { writeFile } from "fs-extra";
+import { outputFile, writeFile } from "fs-extra";
 import { QueryLanguage } from "../common/query-language";
 import { fetchExternalApiQueries } from "./queries";
 import { Method } from "./method";
@@ -57,7 +57,7 @@ export async function prepareModelEditorQueries(
   if (query.dependencies) {
     for (const [filename, contents] of Object.entries(query.dependencies)) {
       const dependencyFile = join(queryDir, filename);
-      await writeFile(dependencyFile, contents, "utf8");
+      await outputFile(dependencyFile, contents, "utf8");
     }
   }
   return true;

extensions/ql-vscode/src/model-editor/model-editor-view.ts

@@ -38,7 +38,7 @@ import { Method } from "./method";
 import { ModeledMethod } from "./modeled-method";
 import { ExtensionPack } from "./shared/extension-pack";
 import { ModelConfigListener } from "../config";
-import { INITIAL_MODE, Mode } from "./shared/mode";
+import { Mode } from "./shared/mode";
 import { loadModeledMethods, saveModeledMethods } from "./modeled-method-fs";
 import { pickExtensionPack } from "./extension-pack-picker";
 import {
@@ -50,12 +50,14 @@ import { telemetryListener } from "../common/vscode/telemetry";
 import { ModelingStore } from "./modeling-store";
 import { ModelEditorViewTracker } from "./model-editor-view-tracker";
 import { ModelingEvents } from "./modeling-events";
+import { getModelsAsDataLanguage, ModelsAsDataLanguage } from "./languages";
 
 export class ModelEditorView extends AbstractWebview<
   ToModelEditorMessage,
   FromModelEditorMessage
 > {
   private readonly autoModeler: AutoModeler;
+  private readonly languageDefinition: ModelsAsDataLanguage;
 
   public constructor(
     protected readonly app: App,
@@ -72,7 +74,7 @@ export class ModelEditorView extends AbstractWebview<
     private readonly extensionPack: ExtensionPack,
     // The language is equal to databaseItem.language but is properly typed as QueryLanguage
     private readonly language: QueryLanguage,
-    initialMode: Mode = INITIAL_MODE,
+    initialMode: Mode,
   ) {
     super(app);
 
@@ -95,6 +97,7 @@ export class ModelEditorView extends AbstractWebview<
         this.addModeledMethods(modeledMethods);
       },
     );
+    this.languageDefinition = getModelsAsDataLanguage(language);
   }
 
   public async openView() {
@@ -376,6 +379,10 @@ export class ModelEditorView extends AbstractWebview<
     const sourceArchiveAvailable =
       this.databaseItem.hasSourceArchiveInExplorer();
 
+    const showModeSwitchButton =
+      this.languageDefinition.availableModes === undefined ||
+      this.languageDefinition.availableModes.length > 1;
+
     await this.postMessage({
       t: "setModelEditorViewState",
       viewState: {
@@ -385,6 +392,7 @@ export class ModelEditorView extends AbstractWebview<
         showLlmButton,
         showMultipleModels: this.modelConfig.showMultipleModels,
         mode: this.modelingStore.getMode(this.databaseItem),
+        showModeSwitchButton,
         sourceArchiveAvailable,
       },
     });