Merge pull request #2486 from github/starcke/candidate-selection

Tweak candidate selection

Author: starcke

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

@@ -8,6 +8,12 @@ import {
 } from "./auto-model-api";
 import type { UsageSnippetsBySignature } from "./auto-model-usages-query";
 
+// Soft limit on the number of candidates to send to the model.
+// Note that the model may return fewer than this number of candidates.
+const candidateLimit = 20;
+// Soft limit on the number of samples to send to the model.
+const sampleLimit = 100;
+
 export function createAutoModelRequest(
   language: string,
   externalApiUsages: ExternalApiUsage[],
@@ -40,11 +46,15 @@ export function createAutoModelRequest(
         ? 0
         : externalApiUsage.methodParameters.split(",").length;
 
+    const candidates: Method[] = [];
+    const samples: Method[] = [];
     for (
-      let argumentIndex = 0;
+      let argumentIndex = -1; // Start at -1 which means `this` as in `this.method()`
       argumentIndex < numberOfArguments;
       argumentIndex++
     ) {
+      const argumentInput: string =
+        argumentIndex === -1 ? "Argument[this]" : `Argument[${argumentIndex}]`;
       const method: Method = {
         package: externalApiUsage.packageName,
         type: externalApiUsage.typeName,
@@ -54,21 +64,35 @@ export function createAutoModelRequest(
           modeledMethod.type === "none"
             ? undefined
             : toMethodClassification(modeledMethod),
-        usages: usagesForMethod.slice(0, 10),
-        input: `Argument[${argumentIndex}]`,
+        usages: usagesForMethod.slice(0, 6), // At most 6 usages per argument
+        input: argumentInput,
       };
 
+      // A method that is supported is modeled outside of the model file, so it is not a candidate.
+      // We also do not want it as a sample because we do not know the classification.
+      if (modeledMethod.type === "none" && externalApiUsage.supported) {
+        continue;
+      }
+
+      // Candidates are methods that are not currently modeled
       if (modeledMethod.type === "none") {
-        request.candidates.push(method);
+        candidates.push(method);
       } else {
-        request.samples.push(method);
+        samples.push(method);
       }
     }
+    // If there is room for at least one candidate, add all candidates.
+    // This ensures that we send all arguments for a method together.
+    // NOTE: this might go above the candidate limit, but that's okay.
+    if (request.candidates.length < candidateLimit) {
+      request.candidates.push(...candidates);
+    }
+    // Same for samples
+    if (request.samples.length < sampleLimit) {
+      request.samples.push(...samples);
+    }
   }
 
-  request.candidates = request.candidates.slice(0, 20);
-  request.samples = request.samples.slice(0, 100);
-
   return request;
 }
 

extensions/ql-vscode/test/unit-tests/data-extensions-editor/auto-model.test.ts

@@ -39,7 +39,7 @@ describe("createAutoModelRequest", () => {
       typeName: "Connection",
       methodName: "createQuery",
       methodParameters: "(String)",
-      supported: true,
+      supported: false,
       usages: [
         {
           label: "createQuery(...)",
@@ -69,7 +69,7 @@ describe("createAutoModelRequest", () => {
       typeName: "Query",
       methodName: "executeScalar",
       methodParameters: "(Class)",
-      supported: true,
+      supported: false,
       usages: [
         {
           label: "executeScalar(...)",
@@ -99,7 +99,7 @@ describe("createAutoModelRequest", () => {
       typeName: "Sql2o",
       methodName: "open",
       methodParameters: "()",
-      supported: true,
+      supported: false,
       usages: [
         {
           label: "open(...)",
@@ -129,7 +129,7 @@ describe("createAutoModelRequest", () => {
       typeName: "PrintStream",
       methodName: "println",
       methodParameters: "(String)",
-      supported: true,
+      supported: false,
       usages: [
         {
           label: "println(...)",
@@ -149,7 +149,7 @@ describe("createAutoModelRequest", () => {
       typeName: "Sql2o",
       methodName: "Sql2o",
       methodParameters: "(String,String,String)",
-      supported: true,
+      supported: false,
       usages: [
         {
           label: "new Sql2o(...)",
@@ -169,7 +169,7 @@ describe("createAutoModelRequest", () => {
       typeName: "Sql2o",
       methodName: "Sql2o",
       methodParameters: "(String)",
-      supported: true,
+      supported: false,
       usages: [
         {
           label: "new Sql2o(...)",
@@ -183,6 +183,26 @@ describe("createAutoModelRequest", () => {
         },
       ],
     },
+    {
+      signature: "org.test.MyClass#test()",
+      packageName: "org.test",
+      typeName: "MyClass",
+      methodName: "test",
+      methodParameters: "()",
+      supported: true,
+      usages: [
+        {
+          label: "abc.test(...)",
+          url: {
+            uri: "file:/home/runner/work/test/Test.java",
+            startLine: 23,
+            startColumn: 23,
+            endLine: 23,
+            endColumn: 36,
+          },
+        },
+      ],
+    },
   ];
 
   const modeledMethods: Record<string, ModeledMethod> = {
@@ -235,6 +255,32 @@ describe("createAutoModelRequest", () => {
     ).toEqual({
       language: "java",
       samples: [
+        {
+          package: "org.sql2o",
+          type: "Sql2o",
+          name: "open",
+          signature: "()",
+          classification: {
+            type: "CLASSIFICATION_TYPE_NEUTRAL",
+            kind: "",
+            explanation: "",
+          },
+          usages: usages["org.sql2o.Sql2o#open()"],
+          input: "Argument[this]",
+        },
+        {
+          package: "org.sql2o",
+          type: "Sql2o",
+          name: "Sql2o",
+          signature: "(String)",
+          classification: {
+            type: "CLASSIFICATION_TYPE_SINK",
+            kind: "jndi-injection",
+            explanation: "",
+          },
+          usages: usages["org.sql2o.Sql2o#Sql2o(String)"],
+          input: "Argument[this]",
+        },
         {
           package: "org.sql2o",
           type: "Sql2o",
@@ -250,6 +296,15 @@ describe("createAutoModelRequest", () => {
         },
       ],
       candidates: [
+        {
+          package: "org.sql2o",
+          type: "Connection",
+          name: "createQuery",
+          signature: "(String)",
+          usages: usages["org.sql2o.Connection#createQuery(String)"],
+          input: "Argument[this]",
+          classification: undefined,
+        },
         {
           package: "org.sql2o",
           type: "Connection",
@@ -259,6 +314,15 @@ describe("createAutoModelRequest", () => {
           input: "Argument[0]",
           classification: undefined,
         },
+        {
+          package: "org.sql2o",
+          type: "Query",
+          name: "executeScalar",
+          signature: "(Class)",
+          usages: usages["org.sql2o.Query#executeScalar(Class)"],
+          input: "Argument[this]",
+          classification: undefined,
+        },
         {
           package: "org.sql2o",
           type: "Query",
@@ -268,6 +332,18 @@ describe("createAutoModelRequest", () => {
           input: "Argument[0]",
           classification: undefined,
         },
+        {
+          package: "org.springframework.boot",
+          type: "SpringApplication",
+          name: "run",
+          signature: "(Class,String[])",
+          usages:
+            usages[
+              "org.springframework.boot.SpringApplication#run(Class,String[])"
+            ],
+          input: "Argument[this]",
+          classification: undefined,
+        },
         {
           package: "org.springframework.boot",
           type: "SpringApplication",
@@ -292,6 +368,15 @@ describe("createAutoModelRequest", () => {
           input: "Argument[1]",
           classification: undefined,
         },
+        {
+          package: "java.io",
+          type: "PrintStream",
+          name: "println",
+          signature: "(String)",
+          usages: usages["java.io.PrintStream#println(String)"],
+          input: "Argument[this]",
+          classification: undefined,
+        },
         {
           package: "java.io",
           type: "PrintStream",
@@ -301,6 +386,15 @@ describe("createAutoModelRequest", () => {
           input: "Argument[0]",
           classification: undefined,
         },
+        {
+          package: "org.sql2o",
+          type: "Sql2o",
+          name: "Sql2o",
+          signature: "(String,String,String)",
+          usages: usages["org.sql2o.Sql2o#Sql2o(String,String,String)"],
+          input: "Argument[this]",
+          classification: undefined,
+        },
         {
           package: "org.sql2o",
           type: "Sql2o",