Extract creation of method argument options to languages

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

@@ -1,4 +1,4 @@
-import { MethodDefinition } from "../method";
+import { MethodArgument, MethodDefinition } from "../method";
 import {
   ModeledMethod,
   NeutralModeledMethod,
@@ -45,6 +45,11 @@ export type ModelsAsDataLanguagePredicates = {
   neutral?: ModelsAsDataLanguagePredicate<NeutralModeledMethod>;
 };
 
+export type MethodArgumentOptions = {
+  options: MethodArgument[];
+  defaultArgumentPath: string;
+};
+
 export type ModelsAsDataLanguage = {
   /**
    * The modes that are available for this language. If not specified, all
@@ -54,4 +59,9 @@ export type ModelsAsDataLanguage = {
   createMethodSignature: (method: MethodDefinition) => string;
   predicates: ModelsAsDataLanguagePredicates;
   modelGeneration?: ModelsAsDataLanguageModelGeneration;
+  /**
+   * Returns the list of valid arguments that can be selected for the given method.
+   * @param method The method to get the valid arguments for.
+   */
+  getArgumentOptions: (method: MethodDefinition) => MethodArgumentOptions;
 };

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

@@ -2,6 +2,7 @@ import { ModelsAsDataLanguage } from "../models-as-data";
 import { sharedExtensiblePredicates, sharedKinds } from "../shared";
 import { Mode } from "../../shared/mode";
 import { parseGenerateModelResults } from "./generate";
+import { getArgumentsList, MethodArgument } from "../../method";
 
 function parseRubyMethodFromPath(path: string): string {
   const match = path.match(/Method\[([^\]]+)].*/);
@@ -157,4 +158,25 @@ export const ruby: ModelsAsDataLanguage = {
     },
     parseResults: parseGenerateModelResults,
   },
+  getArgumentOptions: (method) => {
+    const argumentsList = getArgumentsList(method.methodParameters).map(
+      (argument, index): MethodArgument => ({
+        path: `Argument[${index}]`,
+        label: `Argument[${index}]: ${argument}`,
+      }),
+    );
+
+    return {
+      options: [
+        {
+          path: "Argument[self]",
+          label: "Argument[self]",
+        },
+        ...argumentsList,
+      ],
+      // If there are no arguments, we will default to "Argument[self]"
+      defaultArgumentPath:
+        argumentsList.length > 0 ? argumentsList[0].path : "Argument[self]",
+    };
+  },
 };

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

@@ -3,6 +3,7 @@ import { Provenance } from "../../modeled-method";
 import { DataTuple } from "../../model-extension-file";
 import { sharedExtensiblePredicates, sharedKinds } from "../shared";
 import { filterFlowModelQueries, parseFlowModelResults } from "./generate";
+import { getArgumentsList, MethodArgument } from "../../method";
 
 function readRowToMethod(row: DataTuple[]): string {
   return `${row[0]}.${row[1]}#${row[3]}${row[4]}`;
@@ -145,4 +146,25 @@ export const staticLanguage: ModelsAsDataLanguage = {
     filterQueries: filterFlowModelQueries,
     parseResults: parseFlowModelResults,
   },
+  getArgumentOptions: (method) => {
+    const argumentsList = getArgumentsList(method.methodParameters).map(
+      (argument, index): MethodArgument => ({
+        path: `Argument[${index}]`,
+        label: `Argument[${index}]: ${argument}`,
+      }),
+    );
+
+    return {
+      options: [
+        {
+          path: "Argument[this]",
+          label: "Argument[this]",
+        },
+        ...argumentsList,
+      ],
+      // If there are no arguments, we will default to "Argument[this]"
+      defaultArgumentPath:
+        argumentsList.length > 0 ? argumentsList[0].path : "Argument[this]",
+    };
+  },
 };

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

@@ -61,6 +61,11 @@ export interface Method extends MethodSignature {
   readonly usages: readonly Usage[];
 }
 
+export interface MethodArgument {
+  path: string;
+  label: string;
+}
+
 export function getArgumentsList(methodParameters: string): string[] {
   if (methodParameters === "()") {
     return [];

extensions/ql-vscode/src/view/method-modeling/MethodModelingInputs.tsx

@@ -39,6 +39,7 @@ export const MethodModelingInputs = ({
   onChange,
 }: MethodModelingInputsProps): JSX.Element => {
   const inputProps = {
+    language,
     method,
     modeledMethod,
     onChange,
@@ -82,7 +83,7 @@ export const MethodModelingInputs = ({
           {isModelingInProgress ? (
             <InProgressDropdown />
           ) : (
-            <ModelKindDropdown language={language} {...inputProps} />
+            <ModelKindDropdown {...inputProps} />
           )}
         </Input>
       </Container>

extensions/ql-vscode/src/view/model-editor/MethodRow.tsx

@@ -234,20 +234,23 @@ const ModelableMethodRow = forwardRef<HTMLElement | undefined, MethodRowProps>(
               <DataGridRow key={index} focused={focusedIndex === index}>
                 <DataGridCell>
                   <ModelTypeDropdown
+                    language={viewState.language}
                     method={method}
                     modeledMethod={modeledMethod}
                     onChange={modeledMethodChangedHandlers[index]}
                   />
                 </DataGridCell>
                 <DataGridCell>
                   <ModelInputDropdown
+                    language={viewState.language}
                     method={method}
                     modeledMethod={modeledMethod}
                     onChange={modeledMethodChangedHandlers[index]}
                   />
                 </DataGridCell>
                 <DataGridCell>
                   <ModelOutputDropdown
+                    language={viewState.language}
                     method={method}
                     modeledMethod={modeledMethod}
                     onChange={modeledMethodChangedHandlers[index]}

extensions/ql-vscode/src/view/model-editor/ModelInputDropdown.tsx

@@ -5,34 +5,33 @@ import {
   ModeledMethod,
   modeledMethodSupportsInput,
 } from "../../model-editor/modeled-method";
-import { Method, getArgumentsList } from "../../model-editor/method";
+import { Method } from "../../model-editor/method";
+import { QueryLanguage } from "../../common/query-language";
+import { getModelsAsDataLanguage } from "../../model-editor/languages";
 
 type Props = {
+  language: QueryLanguage;
   method: Method;
   modeledMethod: ModeledMethod | undefined;
   onChange: (modeledMethod: ModeledMethod) => void;
 };
 
 export const ModelInputDropdown = ({
+  language,
   method,
   modeledMethod,
   onChange,
 }: Props): JSX.Element => {
-  const argumentsList = useMemo(
-    () => getArgumentsList(method.methodParameters),
-    [method.methodParameters],
-  );
+  const options = useMemo(() => {
+    const modelsAsDataLanguage = getModelsAsDataLanguage(language);
 
-  const options = useMemo(
-    () => [
-      { value: "Argument[this]", label: "Argument[this]" },
-      ...argumentsList.map((argument, index) => ({
-        value: `Argument[${index}]`,
-        label: `Argument[${index}]: ${argument}`,
-      })),
-    ],
-    [argumentsList],
-  );
+    return modelsAsDataLanguage
+      .getArgumentOptions(method)
+      .options.map((option) => ({
+        value: option.path,
+        label: option.label,
+      }));
+  }, [language, method]);
 
   const enabled = useMemo(
     () => modeledMethod && modeledMethodSupportsInput(modeledMethod),

extensions/ql-vscode/src/view/model-editor/ModelOutputDropdown.tsx

@@ -5,35 +5,34 @@ import {
   ModeledMethod,
   modeledMethodSupportsOutput,
 } from "../../model-editor/modeled-method";
-import { Method, getArgumentsList } from "../../model-editor/method";
+import { Method } from "../../model-editor/method";
+import { getModelsAsDataLanguage } from "../../model-editor/languages";
+import { QueryLanguage } from "../../common/query-language";
 
 type Props = {
+  language: QueryLanguage;
   method: Method;
   modeledMethod: ModeledMethod | undefined;
   onChange: (modeledMethod: ModeledMethod) => void;
 };
 
 export const ModelOutputDropdown = ({
+  language,
   method,
   modeledMethod,
   onChange,
 }: Props): JSX.Element => {
-  const argumentsList = useMemo(
-    () => getArgumentsList(method.methodParameters),
-    [method.methodParameters],
-  );
+  const options = useMemo(() => {
+    const modelsAsDataLanguage = getModelsAsDataLanguage(language);
 
-  const options = useMemo(
-    () => [
-      { value: "ReturnValue", label: "ReturnValue" },
-      { value: "Argument[this]", label: "Argument[this]" },
-      ...argumentsList.map((argument, index) => ({
-        value: `Argument[${index}]`,
-        label: `Argument[${index}]: ${argument}`,
-      })),
-    ],
-    [argumentsList],
-  );
+    const options = modelsAsDataLanguage
+      .getArgumentOptions(method)
+      .options.map((option) => ({
+        value: option.path,
+        label: option.label,
+      }));
+    return [{ value: "ReturnValue", label: "ReturnValue" }, ...options];
+  }, [language, method]);
 
   const enabled = useMemo(
     () => modeledMethod && modeledMethodSupportsOutput(modeledMethod),

extensions/ql-vscode/src/view/model-editor/ModelTypeDropdown.tsx

@@ -1,15 +1,17 @@
 import * as React from "react";
-import { ChangeEvent, useCallback, useMemo } from "react";
+import { ChangeEvent, useCallback } from "react";
 import { Dropdown } from "../common/Dropdown";
 import {
   ModeledMethod,
   modeledMethodSupportsProvenance,
   ModeledMethodType,
   Provenance,
 } from "../../model-editor/modeled-method";
-import { Method, getArgumentsList } from "../../model-editor/method";
+import { Method } from "../../model-editor/method";
 import { createEmptyModeledMethod } from "../../model-editor/modeled-method-empty";
 import { Mutable } from "../../common/mutable";
+import { QueryLanguage } from "../../common/query-language";
+import { getModelsAsDataLanguage } from "../../model-editor/languages";
 
 const options: Array<{ value: ModeledMethodType; label: string }> = [
   { value: "none", label: "Unmodeled" },
@@ -20,23 +22,22 @@ const options: Array<{ value: ModeledMethodType; label: string }> = [
 ];
 
 type Props = {
+  language: QueryLanguage;
   method: Method;
   modeledMethod: ModeledMethod | undefined;
   onChange: (modeledMethod: ModeledMethod) => void;
 };
 
 export const ModelTypeDropdown = ({
+  language,
   method,
   modeledMethod,
   onChange,
 }: Props): JSX.Element => {
-  const argumentsList = useMemo(
-    () => getArgumentsList(method.methodParameters),
-    [method.methodParameters],
-  );
-
   const handleChange = useCallback(
     (e: ChangeEvent<HTMLSelectElement>) => {
+      const modelsAsDataLanguage = getModelsAsDataLanguage(language);
+
       let newProvenance: Provenance = "manual";
       if (modeledMethod && modeledMethodSupportsProvenance(modeledMethod)) {
         if (modeledMethod.provenance === "df-generated") {
@@ -54,9 +55,8 @@ export const ModelTypeDropdown = ({
         ...emptyModeledMethod,
       };
       if ("input" in updatedModeledMethod) {
-        // If there are no arguments, we will default to "Argument[this]"
         updatedModeledMethod.input =
-          argumentsList.length === 0 ? "Argument[this]" : "Argument[0]";
+          modelsAsDataLanguage.getArgumentOptions(method).defaultArgumentPath;
       }
       if ("output" in updatedModeledMethod) {
         updatedModeledMethod.output = "ReturnValue";
@@ -70,7 +70,7 @@ export const ModelTypeDropdown = ({
 
       onChange(updatedModeledMethod);
     },
-    [onChange, method, modeledMethod, argumentsList],
+    [onChange, method, modeledMethod, language],
   );
 
   return (