Skip to content

Commit 66665bf

Browse files
jcreedcmujcreedcmu
committed
Refactor definitions/references finding
1 parent 1c6b4a6 commit 66665bf

File tree

2 files changed

+95
-92
lines changed

2 files changed

+95
-92
lines changed

‎extensions/ql-vscode/src/definitions.ts‎

Lines changed: 92 additions & 89 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,13 @@ import * as messages from "./messages";
1212
import { QueryServerClient } from "./queryserver-client";
1313
import { compileAndRunQueryAgainstDatabase, QueryWithResults } from "./run-queries";
1414

15+
/**
16+
* Run templated CodeQL queries to find definitions and references in
17+
* source-language files. We may eventually want to find a way to
18+
* generalize this to other custom queries, e.g. showing dataflow to
19+
* or from a selected identifier.
20+
*/
21+
1522
const TEMPLATE_NAME = "selectedSourceFile";
1623
const SELECT_QUERY_NAME = "#select";
1724

@@ -46,104 +53,62 @@ async function qlpackOfDatabase(cli: CodeQLCliServer, db: DatabaseItem): Promise
4653
return qlpack;
4754
}
4855

49-
export async function createDefinitionsHandler(cli: CodeQLCliServer, qs: QueryServerClient, dbm: DatabaseManager): Promise<vscode.DefinitionProvider> {
50-
let fileCache = new CachedOperation<vscode.LocationLink[]>(async (uriString: string) => {
51-
const uri = decodeSourceArchiveUri(vscode.Uri.parse(uriString));
52-
const sourceArchiveUri = vscode.Uri.file(uri.sourceArchiveZipPath).with({ scheme: zipArchiveScheme });
53-
54-
const db = dbm.findDatabaseItemBySourceArchive(sourceArchiveUri);
55-
if (db) {
56-
const qlpack = await qlpackOfDatabase(cli, db);
57-
if (qlpack === undefined) {
58-
throw new Error("Can't infer qlpack from database source archive");
59-
}
60-
const links: vscode.DefinitionLink[] = []
61-
for (const query of await resolveQueries(cli, qlpack, KeyType.DefinitionQuery)) {
62-
const templates: messages.TemplateDefinitions = {
63-
[TEMPLATE_NAME]: {
64-
values: {
65-
tuples: [[{
66-
stringValue: uri.pathWithinSourceArchive
67-
}]]
68-
}
69-
}
70-
};
71-
const results = await compileAndRunQueryAgainstDatabase(cli, qs, db, false, vscode.Uri.file(query), templates);
72-
if (results.result.resultType == messages.QueryResultType.SUCCESS) {
73-
links.push(...await getLinksFromResults(results, cli, db, (src, _dest) => src === uriString));
74-
}
75-
}
76-
return links;
77-
} else {
78-
return [];
79-
}
80-
});
81-
82-
return {
83-
async provideDefinition(document: vscode.TextDocument, position: vscode.Position, _token: vscode.CancellationToken): Promise<vscode.LocationLink[]> {
84-
const fileLinks = await fileCache.get(document.uri.toString());
85-
let locLinks: vscode.LocationLink[] = [];
86-
for (const link of fileLinks) {
87-
if (link.originSelectionRange!.contains(position)) {
88-
locLinks.push(link);
89-
}
90-
}
91-
return locLinks;
92-
}
93-
94-
};
95-
}
96-
9756
interface FullLocationLink extends vscode.LocationLink {
9857
originUri: vscode.Uri;
9958
}
10059

101-
export async function createReferencesHander(cli: CodeQLCliServer, qs: QueryServerClient, dbm: DatabaseManager): Promise<vscode.ReferenceProvider> {
102-
let fileCache = new CachedOperation<FullLocationLink[]>(async (uriString: string) => {
103-
const uri = decodeSourceArchiveUri(vscode.Uri.parse(uriString));
104-
const sourceArchiveUri = vscode.Uri.file(uri.sourceArchiveZipPath).with({ scheme: zipArchiveScheme });
60+
export class TemplateQueryDefinitionProvider implements vscode.DefinitionProvider {
61+
private cache: CachedOperation<vscode.LocationLink[]>;
10562

106-
const db = dbm.findDatabaseItemBySourceArchive(sourceArchiveUri);
107-
if (db) {
108-
const qlpack = await qlpackOfDatabase(cli, db);
109-
if (qlpack === undefined) {
110-
throw new Error("Can't infer qlpack from database source archive");
111-
}
112-
const links: FullLocationLink[] = []
113-
for (const query of await resolveQueries(cli, qlpack, KeyType.ReferenceQuery)) {
114-
const templates: messages.TemplateDefinitions = {
115-
[TEMPLATE_NAME]: {
116-
values: {
117-
tuples: [[{
118-
stringValue: uri.pathWithinSourceArchive
119-
}]]
120-
}
121-
}
122-
};
123-
const results = await compileAndRunQueryAgainstDatabase(cli, qs, db, false, vscode.Uri.file(query), templates);
124-
if (results.result.resultType == messages.QueryResultType.SUCCESS) {
125-
links.push(...await getLinksFromResults(results, cli, db, (_src, dest) => dest === uriString));
126-
}
63+
constructor(
64+
private cli: CodeQLCliServer,
65+
private qs: QueryServerClient,
66+
private dbm: DatabaseManager,
67+
) {
68+
this.cache = new CachedOperation<vscode.LocationLink[]>(this.getDefinitions.bind(this));
69+
}
70+
71+
async getDefinitions(uriString: string): Promise<vscode.LocationLink[]> {
72+
return getLinksForUriString(this.cli, this.qs, this.dbm, uriString, (src, _dest) => src === uriString);
73+
}
74+
75+
async provideDefinition(document: vscode.TextDocument, position: vscode.Position, _token: vscode.CancellationToken): Promise<vscode.LocationLink[]> {
76+
const fileLinks = await this.cache.get(document.uri.toString());
77+
let locLinks: vscode.LocationLink[] = [];
78+
for (const link of fileLinks) {
79+
if (link.originSelectionRange!.contains(position)) {
80+
locLinks.push(link);
12781
}
128-
return links;
129-
} else {
130-
return [];
13182
}
132-
})
133-
134-
return {
135-
async provideReferences(document: vscode.TextDocument, position: vscode.Position, _context: vscode.ReferenceContext, _token: vscode.CancellationToken): Promise<vscode.Location[]> {
136-
const fileLinks = await fileCache.get(document.uri.toString());
137-
let locLinks: vscode.Location[] = [];
138-
for (const link of fileLinks) {
139-
if (link.targetRange!.contains(position)) {
140-
locLinks.push({ range: link.originSelectionRange!, uri: link.originUri });
141-
}
83+
return locLinks;
84+
}
85+
}
86+
87+
export class TemplateQueryReferenceProvider implements vscode.ReferenceProvider {
88+
private cache: CachedOperation<FullLocationLink[]>;
89+
90+
constructor(
91+
private cli: CodeQLCliServer,
92+
private qs: QueryServerClient,
93+
private dbm: DatabaseManager,
94+
) {
95+
this.cache = new CachedOperation<FullLocationLink[]>(this.getReferences.bind(this));
96+
}
97+
98+
async getReferences(uriString: string): Promise<FullLocationLink[]> {
99+
return getLinksForUriString(this.cli, this.qs, this.dbm, uriString, (_src, dest) => dest === uriString);
100+
}
101+
102+
async provideReferences(document: vscode.TextDocument, position: vscode.Position, _context: vscode.ReferenceContext, _token: vscode.CancellationToken): Promise<vscode.Location[]> {
103+
const fileLinks = await this.cache.get(document.uri.toString());
104+
let locLinks: vscode.Location[] = [];
105+
for (const link of fileLinks) {
106+
if (link.targetRange!.contains(position)) {
107+
locLinks.push({ range: link.originSelectionRange!, uri: link.originUri });
142108
}
143-
return locLinks;
144109
}
145-
146-
};
110+
return locLinks;
111+
}
147112
}
148113

149114
interface FileRange {
@@ -175,6 +140,44 @@ async function getLinksFromResults(results: QueryWithResults, cli: CodeQLCliServ
175140
return localLinks;
176141
}
177142

143+
async function getLinksForUriString(
144+
cli: CodeQLCliServer,
145+
qs: QueryServerClient,
146+
dbm: DatabaseManager,
147+
uriString: string,
148+
filter: (src: string, dest: string) => boolean
149+
) {
150+
const uri = decodeSourceArchiveUri(vscode.Uri.parse(uriString));
151+
const sourceArchiveUri = vscode.Uri.file(uri.sourceArchiveZipPath).with({ scheme: zipArchiveScheme });
152+
153+
const db = dbm.findDatabaseItemBySourceArchive(sourceArchiveUri);
154+
if (db) {
155+
const qlpack = await qlpackOfDatabase(cli, db);
156+
if (qlpack === undefined) {
157+
throw new Error("Can't infer qlpack from database source archive");
158+
}
159+
const links: FullLocationLink[] = []
160+
for (const query of await resolveQueries(cli, qlpack, KeyType.ReferenceQuery)) {
161+
const templates: messages.TemplateDefinitions = {
162+
[TEMPLATE_NAME]: {
163+
values: {
164+
tuples: [[{
165+
stringValue: uri.pathWithinSourceArchive
166+
}]]
167+
}
168+
}
169+
};
170+
const results = await compileAndRunQueryAgainstDatabase(cli, qs, db, false, vscode.Uri.file(query), templates);
171+
if (results.result.resultType == messages.QueryResultType.SUCCESS) {
172+
links.push(...await getLinksFromResults(results, cli, db, filter));
173+
}
174+
}
175+
return links;
176+
} else {
177+
return [];
178+
}
179+
}
180+
178181
function fileRangeFromURI(uri: UrlValue, db: DatabaseItem): FileRange | undefined {
179182
if (typeof uri === "string") {
180183
return undefined;

‎extensions/ql-vscode/src/extension.ts‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ import { CodeQLCliServer } from './cli';
66
import { DistributionConfigListener, QueryHistoryConfigListener, QueryServerConfigListener } from './config';
77
import { DatabaseManager } from './databases';
88
import { DatabaseUI } from './databases-ui';
9-
import { createDefinitionsHandler, createReferencesHander } from './definitions';
9+
import { TemplateQueryDefinitionProvider, TemplateQueryReferenceProvider } from './definitions';
1010
import { DEFAULT_DISTRIBUTION_VERSION_CONSTRAINT, DistributionManager, DistributionUpdateCheckResultKind, FindDistributionResult, FindDistributionResultKind, GithubApiError, GithubRateLimitedError } from './distribution';
1111
import * as helpers from './helpers';
1212
import { assertNever } from './helpers-pure';
@@ -338,11 +338,11 @@ async function activateWithInstalledDistribution(ctx: ExtensionContext, distribu
338338

339339
languages.registerDefinitionProvider(
340340
{ scheme: archiveFilesystemProvider.zipArchiveScheme },
341-
await createDefinitionsHandler(cliServer, qs, dbm)
341+
new TemplateQueryDefinitionProvider(cliServer, qs, dbm)
342342
);
343343
languages.registerReferenceProvider(
344344
{ scheme: archiveFilesystemProvider.zipArchiveScheme },
345-
await createReferencesHander(cliServer, qs, dbm)
345+
new TemplateQueryReferenceProvider(cliServer, qs, dbm)
346346
);
347347
}
348348

0 commit comments

Comments
 (0)