Move model generation functions to language definition

Author: koesie10

extensions/ql-vscode/src/local-queries/index.ts

@@ -1,5 +1,6 @@
 export * from "./local-queries";
 export * from "./local-query-run";
+export * from "./query-constraints";
 export * from "./query-resolver";
 export * from "./quick-eval-code-lens-provider";
 export * from "./quick-query";

extensions/ql-vscode/src/local-queries/query-constraints.ts

@@ -0,0 +1,7 @@
+export interface QueryConstraints {
+  kind?: string;
+  "tags contain"?: string[];
+  "tags contain all"?: string[];
+  "query filename"?: string;
+  "query path"?: string;
+}

extensions/ql-vscode/src/local-queries/query-resolver.ts

@@ -14,6 +14,7 @@ import { showAndLogExceptionWithTelemetry } from "../common/logging";
 import { extLogger } from "../common/logging/vscode";
 import { telemetryListener } from "../common/vscode/telemetry";
 import { SuiteInstruction } from "../packaging/suite-instruction";
+import { QueryConstraints } from "./query-constraints";
 
 export async function qlpackOfDatabase(
   cli: Pick<CodeQLCliServer, "resolveQlpacks">,
@@ -27,14 +28,6 @@ export async function qlpackOfDatabase(
   return await getQlPackForDbscheme(cli, dbscheme);
 }
 
-export interface QueryConstraints {
-  kind?: string;
-  "tags contain"?: string[];
-  "tags contain all"?: string[];
-  "query filename"?: string;
-  "query path"?: string;
-}
-
 /**
  * Finds the queries with the specified kind and tags in a list of CodeQL packs.
  *

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

@@ -25,7 +25,7 @@ type GenerateQueriesQueryOptions = {
 /**
  * Options that are passed through by the caller of `runGenerateQueries`.
  */
-export type GenerateQueriesOptions = {
+type GenerateQueriesOptions = {
   cliServer: CodeQLCliServer;
   queryRunner: QueryRunner;
   queryStorageDir: string;

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

@@ -2,6 +2,9 @@ import { MethodDefinition } from "../method";
 import { ModeledMethod, ModeledMethodType } from "../modeled-method";
 import { DataTuple } from "../model-extension-file";
 import { Mode } from "../shared/mode";
+import type { QueryConstraints } from "../../local-queries/query-constraints";
+import { DecodedBqrs } from "../../common/bqrs-cli-types";
+import { BaseLogger } from "../../common/logging";
 
 type GenerateMethodDefinition = (method: ModeledMethod) => DataTuple[];
 type ReadModeledMethod = (row: DataTuple[]) => ModeledMethod;
@@ -15,6 +18,22 @@ export type ModelsAsDataLanguagePredicate = {
   readModeledMethod: ReadModeledMethod;
 };
 
+type ModelsAsDataLanguageModelGeneration = {
+  queryConstraints: QueryConstraints;
+  filterQueries?: (queryPath: string) => boolean;
+  parseResults: (
+    // The path to the query that generated the results.
+    queryPath: string,
+    // The results of the query.
+    bqrs: DecodedBqrs,
+    // The language-specific predicate that was used to generate the results. This is passed to allow
+    // sharing of code between different languages.
+    modelsAsDataLanguage: ModelsAsDataLanguage,
+    // The logger to use for logging.
+    logger: BaseLogger,
+  ) => ModeledMethod[];
+};
+
 export type ModelsAsDataLanguagePredicates = Record<
   ModelsAsDataLanguageModelType,
   ModelsAsDataLanguagePredicate
@@ -28,4 +47,5 @@ export type ModelsAsDataLanguage = {
   availableModes?: Mode[];
   createMethodSignature: (method: MethodDefinition) => string;
   predicates: ModelsAsDataLanguagePredicates;
+  modelGeneration?: ModelsAsDataLanguageModelGeneration;
 };

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

@@ -0,0 +1,50 @@
+import { BaseLogger } from "../../../common/logging";
+import { DecodedBqrs } from "../../../common/bqrs-cli-types";
+import { ModelsAsDataLanguage } from "../models-as-data";
+import { ModeledMethod } from "../../modeled-method";
+import { DataTuple } from "../../model-extension-file";
+
+export function parseGenerateModelResults(
+  _queryPath: string,
+  bqrs: DecodedBqrs,
+  modelsAsDataLanguage: ModelsAsDataLanguage,
+  logger: BaseLogger,
+): ModeledMethod[] {
+  const modeledMethods: ModeledMethod[] = [];
+
+  for (const resultSetName in bqrs) {
+    const definition = Object.values(modelsAsDataLanguage.predicates).find(
+      (definition) => definition.extensiblePredicate === resultSetName,
+    );
+    if (definition === undefined) {
+      void logger.log(`No predicate found for ${resultSetName}`);
+
+      continue;
+    }
+
+    const resultSet = bqrs[resultSetName];
+
+    if (
+      resultSet.tuples.some((tuple) =>
+        tuple.some((value) => typeof value === "object"),
+      )
+    ) {
+      void logger.log(
+        `Skipping ${resultSetName} because it contains undefined values`,
+      );
+      continue;
+    }
+
+    modeledMethods.push(
+      ...resultSet.tuples.map((tuple) => {
+        const row = tuple.filter(
+          (value): value is DataTuple => typeof value !== "object",
+        );
+
+        return definition.readModeledMethod(row);
+      }),
+    );
+  }
+
+  return modeledMethods;
+}

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

@@ -1,6 +1,7 @@
 import { ModelsAsDataLanguage } from "../models-as-data";
 import { sharedExtensiblePredicates, sharedKinds } from "../shared";
 import { Mode } from "../../shared/mode";
+import { parseGenerateModelResults } from "./generate";
 
 function parseRubyMethodFromPath(path: string): string {
   const match = path.match(/Method\[([^\]]+)].*/);
@@ -150,4 +151,10 @@ export const ruby: ModelsAsDataLanguage = {
       },
     },
   },
+  modelGeneration: {
+    queryConstraints: {
+      "query path": "queries/modeling/GenerateModel.ql",
+    },
+    parseResults: parseGenerateModelResults,
+  },
 };