Switch ModeledMethod to union of types

extensions/ql-vscode/src/common/mutable.ts

@@ -0,0 +1,6 @@
+/**
+ * Remove all readonly modifiers from a type.
+ */
+export type Mutable<T> = {
+  -readonly [P in keyof T]: T[P];
+};

extensions/ql-vscode/src/model-editor/auto-modeler.ts

@@ -218,8 +218,6 @@ export class AutoModeler {
           {
             type: "neutral",
             kind: "sink",
-            input: "",
-            output: "",
             provenance: "ai-generated",
             signature: candidate.signature,
             packageName: candidate.packageName,

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

@@ -7,10 +7,13 @@ import {
   NotificationLogger,
   showAndLogExceptionWithTelemetry,
 } from "../common/logging";
-import { getModelsAsDataLanguage } from "./languages";
+import {
+  getModelsAsDataLanguageModel,
+  ModelsAsDataLanguagePredicates,
+} from "./languages";
 import { ProgressCallback } from "../common/vscode/progress";
 import { getOnDiskWorkspaceFolders } from "../common/vscode/workspace-folders";
-import { ModeledMethod, ModeledMethodType } from "./modeled-method";
+import { ModeledMethod } from "./modeled-method";
 import { redactableError } from "../common/errors";
 import { telemetryListener } from "../common/vscode/telemetry";
 import { runQuery } from "../local-queries/run-query";
@@ -112,7 +115,7 @@ async function resolveFlowQueries(
 }
 
 async function runSingleFlowQuery(
-  type: Exclude<ModeledMethodType, "none">,
+  type: keyof ModelsAsDataLanguagePredicates,
   queryPath: string | undefined,
   queryStep: number,
   {
@@ -158,9 +161,7 @@ async function runSingleFlowQuery(
   }
 
   // Interpret the results
-  const modelsAsDataLanguage = getModelsAsDataLanguage(language);
-
-  const definition = modelsAsDataLanguage.predicates[type];
+  const definition = getModelsAsDataLanguageModel(language, type);
 
   const bqrsPath = completedQuery.outputDir.bqrsPath;
 

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

@@ -1,5 +1,8 @@
 import { QueryLanguage } from "../../common/query-language";
-import { ModelsAsDataLanguage } from "./models-as-data";
+import {
+  ModelsAsDataLanguage,
+  ModelsAsDataLanguagePredicates,
+} from "./models-as-data";
 import { ruby } from "./ruby";
 import { staticLanguage } from "./static";
 
@@ -18,3 +21,16 @@ export function getModelsAsDataLanguage(
   }
   return definition;
 }
+
+export function getModelsAsDataLanguageModel<
+  T extends keyof ModelsAsDataLanguagePredicates,
+>(
+  language: QueryLanguage,
+  model: T,
+): NonNullable<ModelsAsDataLanguagePredicates[T]> {
+  const definition = getModelsAsDataLanguage(language).predicates[model];
+  if (!definition) {
+    throw new Error(`No models-as-data predicate for ${model}`);
+  }
+  return definition;
+}

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

@@ -1,24 +1,30 @@
 import { MethodDefinition } from "../method";
-import { ModeledMethod, ModeledMethodType } from "../modeled-method";
+import {
+  ModeledMethod,
+  NeutralModeledMethod,
+  SinkModeledMethod,
+  SourceModeledMethod,
+  SummaryModeledMethod,
+} from "../modeled-method";
 import { DataTuple } from "../model-extension-file";
 import { Mode } from "../shared/mode";
 
-type GenerateMethodDefinition = (method: ModeledMethod) => DataTuple[];
+type GenerateMethodDefinition<T> = (method: T) => DataTuple[];
 type ReadModeledMethod = (row: DataTuple[]) => ModeledMethod;
 
-export type ModelsAsDataLanguageModelType = Exclude<ModeledMethodType, "none">;
-
-export type ModelsAsDataLanguagePredicate = {
+export type ModelsAsDataLanguagePredicate<T> = {
   extensiblePredicate: string;
   supportedKinds: string[];
-  generateMethodDefinition: GenerateMethodDefinition;
+  generateMethodDefinition: GenerateMethodDefinition<T>;
   readModeledMethod: ReadModeledMethod;
 };
 
-export type ModelsAsDataLanguagePredicates = Record<
-  ModelsAsDataLanguageModelType,
-  ModelsAsDataLanguagePredicate
->;
+export type ModelsAsDataLanguagePredicates = {
+  source?: ModelsAsDataLanguagePredicate<SourceModeledMethod>;
+  sink?: ModelsAsDataLanguagePredicate<SinkModeledMethod>;
+  summary?: ModelsAsDataLanguagePredicate<SummaryModeledMethod>;
+  neutral?: ModelsAsDataLanguagePredicate<NeutralModeledMethod>;
+};
 
 export type ModelsAsDataLanguage = {
   /**

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

@@ -1,5 +1,5 @@
 import { ModelsAsDataLanguage } from "./models-as-data";
-import { ModeledMethodType, Provenance } from "../modeled-method";
+import { Provenance } from "../modeled-method";
 import { DataTuple } from "../model-extension-file";
 import { sharedExtensiblePredicates, sharedKinds } from "./shared";
 
@@ -34,7 +34,7 @@ export const staticLanguage: ModelsAsDataLanguage = {
         method.provenance,
       ],
       readModeledMethod: (row) => ({
-        type: "source" as ModeledMethodType,
+        type: "source",
         input: "",
         output: row[6] as string,
         kind: row[7] as string,

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

@@ -0,0 +1,88 @@
+import { ModeledMethod, SinkModeledMethod } from "./modeled-method";
+import { MethodSignature } from "./method";
+import { assertNever } from "../common/helpers-pure";
+
+export function createEmptyModeledMethod(
+  type: ModeledMethod["type"],
+  methodSignature: MethodSignature,
+) {
+  const canonicalMethodSignature: MethodSignature = {
+    packageName: methodSignature.packageName,
+    typeName: methodSignature.typeName,
+    methodName: methodSignature.methodName,
+    methodParameters: methodSignature.methodParameters,
+    signature: methodSignature.signature,
+  };
+
+  switch (type) {
+    case "none":
+      return createEmptyNoneModeledMethod(canonicalMethodSignature);
+    case "source":
+      return createEmptySourceModeledMethod(canonicalMethodSignature);
+    case "sink":
+      return createEmptySinkModeledMethod(canonicalMethodSignature);
+    case "summary":
+      return createEmptySummaryModeledMethod(canonicalMethodSignature);
+    case "neutral":
+      return createEmptyNeutralModeledMethod(canonicalMethodSignature);
+    default:
+      assertNever(type);
+  }
+}
+
+function createEmptyNoneModeledMethod(
+  methodSignature: MethodSignature,
+): ModeledMethod {
+  return {
+    ...methodSignature,
+    type: "none",
+  };
+}
+
+function createEmptySourceModeledMethod(
+  methodSignature: MethodSignature,
+): ModeledMethod {
+  return {
+    ...methodSignature,
+    type: "source",
+    output: "",
+    kind: "",
+    provenance: "manual",
+  };
+}
+
+function createEmptySinkModeledMethod(
+  methodSignature: MethodSignature,
+): SinkModeledMethod {
+  return {
+    ...methodSignature,
+    type: "sink",
+    input: "",
+    kind: "",
+    provenance: "manual",
+  };
+}
+
+function createEmptySummaryModeledMethod(
+  methodSignature: MethodSignature,
+): ModeledMethod {
+  return {
+    ...methodSignature,
+    type: "summary",
+    input: "",
+    output: "",
+    kind: "",
+    provenance: "manual",
+  };
+}
+
+function createEmptyNeutralModeledMethod(
+  methodSignature: MethodSignature,
+): ModeledMethod {
+  return {
+    ...methodSignature,
+    type: "neutral",
+    kind: "",
+    provenance: "manual",
+  };
+}

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

@@ -19,12 +19,85 @@ export type Provenance =
   // Entered by the user in the editor manually
   | "manual";
 
-export interface ModeledMethod extends MethodSignature {
-  readonly type: ModeledMethodType;
+export interface NoneModeledMethod extends MethodSignature {
+  readonly type: "none";
+}
+
+export interface SourceModeledMethod extends MethodSignature {
+  readonly type: "source";
+  readonly output: string;
+  readonly kind: ModeledMethodKind;
+  readonly provenance: Provenance;
+}
+
+export interface SinkModeledMethod extends MethodSignature {
+  readonly type: "sink";
+  readonly input: string;
+  readonly kind: ModeledMethodKind;
+  readonly provenance: Provenance;
+}
+
+export interface SummaryModeledMethod extends MethodSignature {
+  readonly type: "summary";
   readonly input: string;
   readonly output: string;
   readonly kind: ModeledMethodKind;
   readonly provenance: Provenance;
 }
 
+export interface NeutralModeledMethod extends MethodSignature {
+  readonly type: "neutral";
+  readonly kind: ModeledMethodKind;
+  readonly provenance: Provenance;
+}
+
+export type ModeledMethod =
+  | NoneModeledMethod
+  | SourceModeledMethod
+  | SinkModeledMethod
+  | SummaryModeledMethod
+  | NeutralModeledMethod;
+
 export type ModeledMethodKind = string;
+
+export function modeledMethodSupportsKind(
+  modeledMethod: ModeledMethod,
+): modeledMethod is
+  | SourceModeledMethod
+  | SinkModeledMethod
+  | SummaryModeledMethod
+  | NeutralModeledMethod {
+  return (
+    modeledMethod.type === "source" ||
+    modeledMethod.type === "sink" ||
+    modeledMethod.type === "summary" ||
+    modeledMethod.type === "neutral"
+  );
+}
+
+export function modeledMethodSupportsInput(
+  modeledMethod: ModeledMethod,
+): modeledMethod is SinkModeledMethod | SummaryModeledMethod {
+  return modeledMethod.type === "sink" || modeledMethod.type === "summary";
+}
+
+export function modeledMethodSupportsOutput(
+  modeledMethod: ModeledMethod,
+): modeledMethod is SourceModeledMethod | SummaryModeledMethod {
+  return modeledMethod.type === "source" || modeledMethod.type === "summary";
+}
+
+export function modeledMethodSupportsProvenance(
+  modeledMethod: ModeledMethod,
+): modeledMethod is
+  | SourceModeledMethod
+  | SinkModeledMethod
+  | SummaryModeledMethod
+  | NeutralModeledMethod {
+  return (
+    modeledMethod.type === "source" ||
+    modeledMethod.type === "sink" ||
+    modeledMethod.type === "summary" ||
+    modeledMethod.type === "neutral"
+  );
+}

extensions/ql-vscode/src/model-editor/shared/validation.ts

@@ -1,4 +1,4 @@
-import { ModeledMethod } from "../modeled-method";
+import { ModeledMethod, NeutralModeledMethod } from "../modeled-method";
 import { MethodSignature } from "../method";
 import { assertNever } from "../../common/helpers-pure";
 
@@ -37,16 +37,11 @@ function canonicalizeModeledMethod(
       return {
         ...methodSignature,
         type: "none",
-        input: "",
-        output: "",
-        kind: "",
-        provenance: "manual",
       };
     case "source":
       return {
         ...methodSignature,
         type: "source",
-        input: "",
         output: modeledMethod.output,
         kind: modeledMethod.kind,
         provenance: "manual",
@@ -56,7 +51,6 @@ function canonicalizeModeledMethod(
         ...methodSignature,
         type: "sink",
         input: modeledMethod.input,
-        output: "",
         kind: modeledMethod.kind,
         provenance: "manual",
       };
@@ -73,13 +67,11 @@ function canonicalizeModeledMethod(
       return {
         ...methodSignature,
         type: "neutral",
-        input: "",
-        output: "",
         kind: modeledMethod.kind,
         provenance: "manual",
       };
     default:
-      assertNever(modeledMethod.type);
+      assertNever(modeledMethod);
   }
 }
 
@@ -118,7 +110,8 @@ export function validateModeledMethods(
   }
 
   const neutralModeledMethods = consideredModeledMethods.filter(
-    (modeledMethod) => modeledMethod.type === "neutral",
+    (modeledMethod): modeledMethod is NeutralModeledMethod =>
+      modeledMethod.type === "neutral",
   );
 
   const neutralModeledMethodsByKind = new Map<string, ModeledMethod[]>();