Add supported endpoint types to predicates

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

@@ -33,7 +33,7 @@ export function decodeBqrsToMethods(
     let libraryVersion: string | undefined;
     let type: ModeledMethodType;
     let classification: CallClassification;
-    let endpointType = EndpointType.Method;
+    let endpointType: EndpointType | undefined = undefined;
 
     if (mode === Mode.Application) {
       [
@@ -67,8 +67,19 @@ export function decodeBqrsToMethods(
       type = "none";
     }
 
-    if (methodName === "") {
-      endpointType = EndpointType.Class;
+    if (definition.endpointTypeForEndpoint) {
+      endpointType = definition.endpointTypeForEndpoint({
+        endpointType,
+        packageName,
+        typeName,
+        methodName,
+        methodParameters,
+      });
+    }
+
+    if (endpointType === undefined) {
+      endpointType =
+        methodName === "" ? EndpointType.Class : EndpointType.Method;
     }
 
     const signature = definition.createMethodSignature({

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

@@ -1,4 +1,4 @@
-import type { MethodArgument, MethodDefinition } from "../method";
+import type { EndpointType, MethodArgument, MethodDefinition } from "../method";
 import type {
   ModeledMethod,
   NeutralModeledMethod,
@@ -23,6 +23,11 @@ type ReadModeledMethod = (row: DataTuple[]) => ModeledMethod;
 export type ModelsAsDataLanguagePredicate<T> = {
   extensiblePredicate: string;
   supportedKinds?: string[];
+  /**
+   * The endpoint types that this predicate supports. If not specified, the predicate supports all
+   * endpoint types.
+   */
+  supportedEndpointTypes?: EndpointType[];
   generateMethodDefinition: GenerateMethodDefinition<T>;
   readModeledMethod: ReadModeledMethod;
 };
@@ -76,6 +81,18 @@ export type ModelsAsDataLanguage = {
    */
   availableModes?: Mode[];
   createMethodSignature: (method: MethodDefinition) => string;
+  /**
+   * This allows modifying the endpoint type automatically assigned to an endpoint. The default
+   * endpoint type is undefined, and if this method returns undefined, the default endpoint type will
+   * be determined by heuristics.
+   * @param method The method to get the endpoint type for. The endpoint type can be undefined if the
+   *               query does not return an endpoint type.
+   */
+  endpointTypeForEndpoint?: (
+    method: Omit<MethodDefinition, "endpointType"> & {
+      endpointType: EndpointType | undefined;
+    },
+  ) => EndpointType | undefined;
   predicates: ModelsAsDataLanguagePredicates;
   modelGeneration?: ModelsAsDataLanguageModelGeneration;
   accessPathSuggestions?: ModelsAsDataLanguageAccessPathSuggestions;

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

@@ -64,6 +64,27 @@ export function rubyPath(methodName: string, path: string) {
   return `${methodPath}.${path}`;
 }
 
-export function rubyEndpointType(methodName: string) {
-  return methodName === "" ? EndpointType.Class : EndpointType.Method;
+/** For the purpose of the model editor, we are defining the endpoint types as follows:
+ * - Class: A class instance
+ * - Module: The class itself
+ * - Method: A method in a class
+ * - Constructor: A constructor method
+ * @param typeName
+ * @param methodName
+ */
+export function rubyEndpointType(typeName: string, methodName: string) {
+  if (typeName.endsWith("!") && methodName === "new") {
+    // This is a constructor
+    return EndpointType.Constructor;
+  }
+
+  if (typeName.endsWith("!") && methodName === "") {
+    return EndpointType.Module;
+  }
+
+  if (methodName === "") {
+    return EndpointType.Class;
+  }
+
+  return EndpointType.Method;
 }

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

@@ -3,7 +3,7 @@ import { sharedExtensiblePredicates, sharedKinds } from "../shared";
 import { Mode } from "../../shared/mode";
 import { parseGenerateModelResults } from "./generate";
 import type { MethodArgument } from "../../method";
-import { getArgumentsList } from "../../method";
+import { EndpointType, getArgumentsList } from "../../method";
 import {
   parseRubyAccessPath,
   parseRubyMethodFromPath,
@@ -19,10 +19,13 @@ export const ruby: ModelsAsDataLanguage = {
   availableModes: [Mode.Framework],
   createMethodSignature: ({ typeName, methodName }) =>
     `${typeName}#${methodName}`,
+  endpointTypeForEndpoint: ({ typeName, methodName }) =>
+    rubyEndpointType(typeName, methodName),
   predicates: {
     source: {
       extensiblePredicate: sharedExtensiblePredicates.source,
       supportedKinds: sharedKinds.source,
+      supportedEndpointTypes: [EndpointType.Method, EndpointType.Class],
       // extensible predicate sourceModel(
       //   string type, string path, string kind
       // );
@@ -42,7 +45,7 @@ export const ruby: ModelsAsDataLanguage = {
           kind: row[2] as string,
           provenance: "manual",
           signature: rubyMethodSignature(typeName, methodName),
-          endpointType: rubyEndpointType(methodName),
+          endpointType: rubyEndpointType(typeName, methodName),
           packageName: "",
           typeName,
           methodName,
@@ -53,6 +56,7 @@ export const ruby: ModelsAsDataLanguage = {
     sink: {
       extensiblePredicate: sharedExtensiblePredicates.sink,
       supportedKinds: sharedKinds.sink,
+      supportedEndpointTypes: [EndpointType.Method, EndpointType.Constructor],
       // extensible predicate sinkModel(
       //   string type, string path, string kind
       // );
@@ -74,7 +78,7 @@ export const ruby: ModelsAsDataLanguage = {
           kind: row[2] as string,
           provenance: "manual",
           signature: rubyMethodSignature(typeName, methodName),
-          endpointType: rubyEndpointType(methodName),
+          endpointType: rubyEndpointType(typeName, methodName),
           packageName: "",
           typeName,
           methodName,
@@ -85,6 +89,7 @@ export const ruby: ModelsAsDataLanguage = {
     summary: {
       extensiblePredicate: sharedExtensiblePredicates.summary,
       supportedKinds: sharedKinds.summary,
+      supportedEndpointTypes: [EndpointType.Method, EndpointType.Constructor],
       // extensible predicate summaryModel(
       //   string type, string path, string input, string output, string kind
       // );
@@ -105,7 +110,7 @@ export const ruby: ModelsAsDataLanguage = {
           kind: row[4] as string,
           provenance: "manual",
           signature: rubyMethodSignature(typeName, methodName),
-          endpointType: rubyEndpointType(methodName),
+          endpointType: rubyEndpointType(typeName, methodName),
           packageName: "",
           typeName,
           methodName,
@@ -132,7 +137,7 @@ export const ruby: ModelsAsDataLanguage = {
           kind: row[2] as string,
           provenance: "manual",
           signature: rubyMethodSignature(typeName, methodName),
-          endpointType: rubyEndpointType(methodName),
+          endpointType: rubyEndpointType(typeName, methodName),
           packageName: "",
           typeName,
           methodName,
@@ -157,7 +162,7 @@ export const ruby: ModelsAsDataLanguage = {
           relatedTypeName: row[0] as string,
           path,
           signature: rubyMethodSignature(typeName, methodName),
-          endpointType: rubyEndpointType(methodName),
+          endpointType: rubyEndpointType(typeName, methodName),
           packageName: "",
           typeName,
           methodName,

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

@@ -68,7 +68,7 @@ export function parseAccessPathSuggestionsResults(
       return {
         method: {
           packageName: "",
-          endpointType: rubyEndpointType(methodName),
+          endpointType: rubyEndpointType(type, methodName),
           typeName: type,
           methodName,
           methodParameters: "",

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

@@ -17,9 +17,17 @@ export type Usage = Call & {
   readonly classification: CallClassification;
 };
 
+/**
+ * Endpoint types are generic and can be used to represent different types of endpoints in different languages.
+ *
+ * For a reference of symbol kinds used in the LSP protocol (which is a good reference for widely supported features), see
+ * https://github.com/microsoft/vscode-languageserver-node/blob/4c8115f40b52f2e13adab41109c5b1208fc155ab/types/src/main.ts#L2890-L2920
+ */
 export enum EndpointType {
-  Method = "method",
+  Module = "module",
   Class = "class",
+  Method = "method",
+  Constructor = "constructor",
 }
 
 export interface MethodDefinition {

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

@@ -12,7 +12,8 @@ import type { Method } from "../../model-editor/method";
 import { createEmptyModeledMethod } from "../../model-editor/modeled-method-empty";
 import type { Mutable } from "../../common/mutable";
 import { ReadonlyDropdown } from "../common/ReadonlyDropdown";
-import { QueryLanguage } from "../../common/query-language";
+import type { QueryLanguage } from "../../common/query-language";
+import type { ModelsAsDataLanguagePredicates } from "../../model-editor/languages";
 import { getModelsAsDataLanguage } from "../../model-editor/languages";
 import type { ModelingStatus } from "../../model-editor/shared/modeling-status";
 import { InputDropdown } from "./InputDropdown";
@@ -25,6 +26,16 @@ type Props = {
   onChange: (modeledMethod: ModeledMethod) => void;
 };
 
+const typeLabels: Record<keyof ModelsAsDataLanguagePredicates, string> = {
+  source: "Source",
+  sink: "Sink",
+  summary: "Flow summary",
+  neutral: "Neutral",
+  type: "Type",
+};
+
+type Option = { value: ModeledMethodType; label: string };
+
 export const ModelTypeDropdown = ({
   language,
   method,
@@ -33,19 +44,31 @@ export const ModelTypeDropdown = ({
   onChange,
 }: Props): React.JSX.Element => {
   const options = useMemo(() => {
-    const baseOptions: Array<{ value: ModeledMethodType; label: string }> = [
+    const modelsAsDataLanguage = getModelsAsDataLanguage(language);
+
+    const baseOptions: Option[] = [
       { value: "none", label: "Unmodeled" },
-      { value: "source", label: "Source" },
-      { value: "sink", label: "Sink" },
-      { value: "summary", label: "Flow summary" },
-      { value: "neutral", label: "Neutral" },
+      ...Object.entries(modelsAsDataLanguage.predicates)
+        .map(([predicateKey, predicate]): Option | null => {
+          const type = predicateKey as keyof ModelsAsDataLanguagePredicates;
+
+          if (
+            predicate.supportedEndpointTypes &&
+            !predicate.supportedEndpointTypes.includes(method.endpointType)
+          ) {
+            return null;
+          }
+
+          return {
+            value: type,
+            label: typeLabels[type],
+          };
+        })
+        .filter((option): option is Option => option !== null),
     ];
-    if (language === QueryLanguage.Ruby) {
-      baseOptions.push({ value: "type", label: "Type" });
-    }
 
     return baseOptions;
-  }, [language]);
+  }, [language, method.endpointType]);
 
   const handleChange = useCallback(
     (e: ChangeEvent<HTMLSelectElement>) => {

extensions/ql-vscode/test/unit-tests/model-editor/languages/ruby/generate.test.ts

@@ -136,7 +136,7 @@ describe("parseGenerateModelResults", () => {
         typeName: "SQLite3::Database",
       },
       {
-        endpointType: EndpointType.Method,
+        endpointType: EndpointType.Constructor,
         input: "Argument[1]",
         kind: "value",
         methodName: "new",

extensions/ql-vscode/test/vscode-tests/no-workspace/model-editor/generate.test.ts

@@ -196,7 +196,7 @@ describe("runGenerateQueries", () => {
         typeName: "SQLite3::Database",
       },
       {
-        endpointType: EndpointType.Method,
+        endpointType: EndpointType.Constructor,
         input: "Argument[1]",
         kind: "value",
         methodName: "new",