Merge pull request #401 from jcreedcmu/jcreed/pagination

Implement pagination for BQRS results.
2020-06-05 10:23:27 -04:00
parent 109c8755c3 c90dae89c1
commit 91bd7f5971
9 changed files with 322 additions and 55 deletions
--- a/extensions/ql-vscode/src/adapt.ts
+++ b/extensions/ql-vscode/src/adapt.ts
@@ -101,3 +101,34 @@ export function adaptBqrs(schema: AdaptedSchema, page: DecodedBqrsChunk): RawRes
    rows: page.tuples.map(adaptRow),
  };
 }
+
+/**
+ * This type has two branches; we are in the process of changing from
+ * one to the other. The old way is to parse them inside the webview,
+ * the new way is to parse them in the extension. The main motivation
+ * for this transition is to make pagination possible in such a way
+ * that only one page needs to be sent from the extension to the webview.
+ */
+export type ParsedResultSets = ExtensionParsedResultSets | WebviewParsedResultSets;
+
+/**
+ * The old method doesn't require any nontrivial information to be included here,
+ * just a tag to indicate that it is being used.
+ */
+export interface WebviewParsedResultSets {
+  t: 'WebviewParsed';
+  selectedTable?: string; // when undefined, means 'show default table'
+}
+
+/**
+ * The new method includes which bqrs page is being sent, and the
+ * actual results parsed on the extension side.
+ */
+export interface ExtensionParsedResultSets {
+  t: 'ExtensionParsed';
+  pageNumber: number;
+  numPages: number;
+  selectedTable?: string; // when undefined, means 'show default table'
+  resultSetNames: string[];
+  resultSet: RawResultSet;
+}
--- a/extensions/ql-vscode/src/interface-types.ts
+++ b/extensions/ql-vscode/src/interface-types.ts
@@ -1,6 +1,6 @@
 import * as sarif from 'sarif';
 import { ResolvableLocationValue } from 'semmle-bqrs';
-import { RawResultSet } from './adapt';
+import { ParsedResultSets } from './adapt';

 /**
 * Only ever show this many results per run in interpreted results.
@@ -12,6 +12,11 @@ export const INTERPRETED_RESULTS_PER_RUN_LIMIT = 100;
 */
 export const RAW_RESULTS_LIMIT = 10000;

+/**
+ * Show this many rows in a raw result table at a time.
+ */
+export const RAW_RESULTS_PAGE_SIZE = 100;
+
 export interface DatabaseInfo {
  name: string;
  databaseUri: string;
@@ -81,9 +86,10 @@ export interface SetStateMsg {

  /**
   * An experimental way of providing results from the extension.
-   * Should be undefined unless config.EXPERIMENTAL_BQRS_SETTING is set to true.
+   * Should be in the WebviewParsedResultSets branch of the type
+   * unless config.EXPERIMENTAL_BQRS_SETTING is set to true.
   */
-  resultSets?: RawResultSet[];
+  parsedResultSets: ParsedResultSets;
 }

 /** Advance to the next or previous path no in the path viewer */
@@ -101,7 +107,8 @@ export type FromResultsViewMsg =
  | ToggleDiagnostics
  | ChangeRawResultsSortMsg
  | ChangeInterpretedResultsSortMsg
-  | ResultViewLoaded;
+  | ResultViewLoaded
+  | ChangePage;

 interface ViewSourceFileMsg {
  t: 'viewSourceFile';
@@ -122,6 +129,12 @@ interface ResultViewLoaded {
  t: 'resultViewLoaded';
 }

+interface ChangePage {
+  t: 'changePage';
+  pageNumber: number; // 0-indexed, displayed to the user as 1-indexed
+  selectedTable: string;
+}
+
 export enum SortDirection {
  asc, desc
 }
--- a/extensions/ql-vscode/src/interface-utils.ts
+++ b/extensions/ql-vscode/src/interface-utils.ts
@@ -0,0 +1,22 @@
+import { RawResultSet } from "./adapt";
+import { ResultSetSchema } from "semmle-bqrs";
+import { Interpretation } from "./interface-types";
+
+export const SELECT_TABLE_NAME = '#select';
+export const ALERTS_TABLE_NAME = 'alerts';
+
+export type RawTableResultSet = { t: 'RawResultSet' } & RawResultSet;
+export type PathTableResultSet = { t: 'SarifResultSet'; readonly schema: ResultSetSchema; name: string } & Interpretation;
+
+export type ResultSet =
+  | RawTableResultSet
+  | PathTableResultSet;
+
+export function getDefaultResultSet(resultSets: readonly ResultSet[]): string {
+  return getDefaultResultSetName(resultSets.map(resultSet => resultSet.schema.name));
+}
+
+export function getDefaultResultSetName(resultSetNames: readonly string[]): string {
+  // Choose first available result set from the array
+  return [ALERTS_TABLE_NAME, SELECT_TABLE_NAME, resultSetNames[0]].filter(resultSetName => resultSetNames.includes(resultSetName))[0];
+}
--- a/extensions/ql-vscode/src/interface.ts
+++ b/extensions/ql-vscode/src/interface.ts
@@ -10,14 +10,15 @@ import { CodeQLCliServer } from './cli';
 import { DatabaseItem, DatabaseManager } from './databases';
 import { showAndLogErrorMessage } from './helpers';
 import { assertNever } from './helpers-pure';
-import { FromResultsViewMsg, Interpretation, INTERPRETED_RESULTS_PER_RUN_LIMIT, IntoResultsViewMsg, QueryMetadata, ResultsPaths, SortedResultSetInfo, SortedResultsMap, InterpretedResultsSortState, SortDirection } from './interface-types';
+import { FromResultsViewMsg, Interpretation, INTERPRETED_RESULTS_PER_RUN_LIMIT, IntoResultsViewMsg, QueryMetadata, ResultsPaths, SortedResultSetInfo, SortedResultsMap, InterpretedResultsSortState, SortDirection, RAW_RESULTS_PAGE_SIZE } from './interface-types';
 import { Logger } from './logging';
 import * as messages from './messages';
 import { CompletedQuery, interpretResults } from './query-results';
 import { QueryInfo, tmpDir } from './run-queries';
 import { parseSarifLocation, parseSarifPlainTextMessage } from './sarif-utils';
-import { adaptSchema, adaptBqrs, RawResultSet } from './adapt';
+import { adaptSchema, adaptBqrs, RawResultSet, ParsedResultSets } from './adapt';
 import { EXPERIMENTAL_BQRS_SETTING } from './config';
+import { getDefaultResultSetName } from './interface-utils';

 /**
 * interface.ts
@@ -115,8 +116,13 @@ function sortInterpretedResults(results: Sarif.Result[], sortState: InterpretedR
  }
 }

+function numPagesOfResultSet(resultSet: RawResultSet): number {
+  return Math.ceil(resultSet.schema.tupleCount / RAW_RESULTS_PAGE_SIZE);
+}
+
 export class InterfaceManager extends DisposableObject {
  private _displayedQuery?: CompletedQuery;
+  private _interpretation?: Interpretation;
  private _panel: vscode.WebviewPanel | undefined;
  private _panelLoaded = false;
  private _panelLoadedCallBacks: (() => void)[] = [];
@@ -288,6 +294,9 @@ export class InterfaceManager extends DisposableObject {
          query.updateInterpretedSortState(this.cliServer, msg.sortState)
        );
        break;
+      case "changePage":
+        await this.showPageOfResults(msg.selectedTable, msg.pageNumber);
+        break;
      default:
        assertNever(msg);
    }
@@ -339,6 +348,7 @@ export class InterfaceManager extends DisposableObject {
    );

    this._displayedQuery = results;
+    this._interpretation = interpretation;

    const panel = this.getPanel();
    await this.waitForPanelLoaded();
@@ -364,18 +374,37 @@ export class InterfaceManager extends DisposableObject {
      });
    }

-    let resultSets: RawResultSet[] | undefined;
+    const getParsedResultSets = async (): Promise<ParsedResultSets> => {
+      if (EXPERIMENTAL_BQRS_SETTING.getValue()) {
+        const schemas = await this.cliServer.bqrsInfo(results.query.resultsPaths.resultsPath, RAW_RESULTS_PAGE_SIZE);

-    if (EXPERIMENTAL_BQRS_SETTING.getValue()) {
-      resultSets = [];
-      const schemas = await this.cliServer.bqrsInfo(results.query.resultsPaths.resultsPath);
-      for (const schema of schemas["result-sets"]) {
-        const chunk = await this.cliServer.bqrsDecode(results.query.resultsPaths.resultsPath, schema.name);
+        const resultSetNames = schemas["result-sets"].map(resultSet => resultSet.name);
+
+        // This may not wind up being the page we actually show, if there are interpreted results,
+        // but speculatively send it anyway.
+        const selectedTable = getDefaultResultSetName(resultSetNames);
+        const schema = schemas["result-sets"].find(resultSet => resultSet.name == selectedTable)!;
+        if (schema === undefined) {
+          return { t: 'WebviewParsed' };
+        }
+
+        const chunk = await this.cliServer.bqrsDecode(results.query.resultsPaths.resultsPath, schema.name, RAW_RESULTS_PAGE_SIZE, schema.pagination?.offsets[0]);
        const adaptedSchema = adaptSchema(schema);
        const resultSet = adaptBqrs(adaptedSchema, chunk);
-        resultSets.push(resultSet);
+
+        return {
+          t: 'ExtensionParsed',
+          pageNumber: 0,
+          numPages: numPagesOfResultSet(resultSet),
+          resultSet,
+          selectedTable: undefined,
+          resultSetNames
+        };
      }
-    }
+      else {
+        return { t: 'WebviewParsed' };
+      }
+    };

    await this.postMessage({
      t: "setState",
@@ -384,7 +413,7 @@ export class InterfaceManager extends DisposableObject {
      resultsPath: this.convertPathToWebviewUri(
        results.query.resultsPaths.resultsPath
      ),
-      resultSets,
+      parsedResultSets: await getParsedResultSets(),
      sortedResultsMap,
      database: results.database,
      shouldKeepOldResultsWhileRendering,
@@ -392,6 +421,59 @@ export class InterfaceManager extends DisposableObject {
    });
  }

+  /**
+   * Show a page of raw results from the chosen table.
+   */
+  public async showPageOfResults(selectedTable: string, pageNumber: number): Promise<void> {
+    const results = this._displayedQuery;
+    if (results === undefined) {
+      throw new Error('trying to view a page of a query that is not loaded');
+    }
+
+    const sortedResultsMap: SortedResultsMap = {};
+    results.sortedResultsInfo.forEach(
+      (v, k) =>
+        (sortedResultsMap[k] = this.convertPathPropertiesToWebviewUris(
+          v
+        ))
+    );
+
+    const schemas = await this.cliServer.bqrsInfo(results.query.resultsPaths.resultsPath, RAW_RESULTS_PAGE_SIZE);
+
+    const resultSetNames = schemas["result-sets"].map(resultSet => resultSet.name);
+
+    const schema = schemas["result-sets"].find(resultSet => resultSet.name == selectedTable)!;
+    if (schema === undefined)
+      throw new Error(`Query result set '${selectedTable}' not found.`);
+
+    const chunk = await this.cliServer.bqrsDecode(results.query.resultsPaths.resultsPath, schema.name, RAW_RESULTS_PAGE_SIZE, schema.pagination?.offsets[pageNumber]);
+    const adaptedSchema = adaptSchema(schema);
+    const resultSet = adaptBqrs(adaptedSchema, chunk);
+
+    const parsedResultSets: ParsedResultSets = {
+      t: 'ExtensionParsed',
+      pageNumber,
+      resultSet,
+      numPages: numPagesOfResultSet(resultSet),
+      selectedTable: selectedTable,
+      resultSetNames
+    };
+
+    await this.postMessage({
+      t: "setState",
+      interpretation: this._interpretation,
+      origResultsPaths: results.query.resultsPaths,
+      resultsPath: this.convertPathToWebviewUri(
+        results.query.resultsPaths.resultsPath
+      ),
+      parsedResultSets,
+      sortedResultsMap,
+      database: results.database,
+      shouldKeepOldResultsWhileRendering: false,
+      metadata: results.query.metadata
+    });
+  }
+
  private async getTruncatedResults(
    metadata: QueryMetadata | undefined,
    resultsPaths: ResultsPaths,
--- a/extensions/ql-vscode/src/view/alert-table.tsx
+++ b/extensions/ql-vscode/src/view/alert-table.tsx
@@ -5,9 +5,10 @@ import * as Keys from '../result-keys';
 import { LocationStyle } from 'semmle-bqrs';
 import * as octicons from './octicons';
 import { className, renderLocation, ResultTableProps, zebraStripe, selectableZebraStripe, jumpToLocation, nextSortDirection } from './result-table-utils';
-import { PathTableResultSet, onNavigation, NavigationEvent, vscode } from './results';
+import { onNavigation, NavigationEvent, vscode } from './results';
 import { parseSarifPlainTextMessage, parseSarifLocation } from '../sarif-utils';
 import { InterpretedResultsSortColumn, SortDirection, InterpretedResultsSortState } from '../interface-types';
+import { PathTableResultSet } from '../interface-utils';

 export type PathTableProps = ResultTableProps & { resultSet: PathTableResultSet };
 export interface PathTableState {
--- a/extensions/ql-vscode/src/view/raw-results-table.tsx
+++ b/extensions/ql-vscode/src/view/raw-results-table.tsx
@@ -1,12 +1,14 @@
 import * as React from "react";
 import { renderLocation, ResultTableProps, zebraStripe, className, nextSortDirection } from "./result-table-utils";
-import { RawTableResultSet, vscode } from "./results";
+import { vscode } from "./results";
 import { ResultValue } from "../adapt";
 import { SortDirection, RAW_RESULTS_LIMIT, RawResultsSortState } from "../interface-types";
+import { RawTableResultSet } from "../interface-utils";

 export type RawTableProps = ResultTableProps & {
  resultSet: RawTableResultSet;
  sortState?: RawResultsSortState;
+  offset: number;
 };

 export class RawTable extends React.Component<RawTableProps, {}> {
@@ -28,7 +30,7 @@ export class RawTable extends React.Component<RawTableProps, {}> {
      <tr key={rowIndex} {...zebraStripe(rowIndex)}>
        {
          [
-            <td key={-1}>{rowIndex + 1}</td>,
+            <td key={-1}>{rowIndex + 1 + this.props.offset}</td>,
            ...row.map((value, columnIndex) =>
              <td key={columnIndex}>
                {
--- a/extensions/ql-vscode/src/view/result-table-utils.tsx
+++ b/extensions/ql-vscode/src/view/result-table-utils.tsx
@@ -1,8 +1,9 @@
 import * as React from 'react';
 import { LocationValue, ResolvableLocationValue, tryGetResolvableLocation } from 'semmle-bqrs';
 import { RawResultsSortState, QueryMetadata, SortDirection } from '../interface-types';
-import { ResultSet, vscode } from './results';
+import { vscode } from './results';
 import { assertNever } from '../helpers-pure';
+import { ResultSet } from '../interface-utils';

 export interface ResultTableProps {
  resultSet: ResultSet;
@@ -10,6 +11,7 @@ export interface ResultTableProps {
  metadata?: QueryMetadata;
  resultsPath: string | undefined;
  sortState?: RawResultsSortState;
+  offset: number;

  /**
   * Holds if there are any raw results. When that is the case, we
--- a/extensions/ql-vscode/src/view/result-tables.tsx
+++ b/extensions/ql-vscode/src/view/result-tables.tsx
@@ -1,14 +1,17 @@
 import * as React from 'react';
-import { DatabaseInfo, Interpretation, RawResultsSortState, QueryMetadata, ResultsPaths, InterpretedResultsSortState } from '../interface-types';
+import { DatabaseInfo, Interpretation, RawResultsSortState, QueryMetadata, ResultsPaths, InterpretedResultsSortState, RAW_RESULTS_PAGE_SIZE } from '../interface-types';
 import { PathTable } from './alert-table';
 import { RawTable } from './raw-results-table';
 import { ResultTableProps, tableSelectionHeaderClassName, toggleDiagnosticsClassName, alertExtrasClassName } from './result-table-utils';
-import { ResultSet, vscode } from './results';
+import { vscode } from './results';
+import { ParsedResultSets, ExtensionParsedResultSets } from '../adapt';
+import { ResultSet, ALERTS_TABLE_NAME, SELECT_TABLE_NAME, getDefaultResultSet } from '../interface-utils';

 /**
 * Properties for the `ResultTables` component.
 */
 export interface ResultTablesProps {
+  parsedResultSets: ParsedResultSets;
  rawResultSets: readonly ResultSet[];
  interpretation: Interpretation | undefined;
  database: DatabaseInfo;
@@ -25,10 +28,9 @@ export interface ResultTablesProps {
 */
 interface ResultTablesState {
  selectedTable: string; // name of selected result set
+  selectedPage: string; // stringified selected page
 }

-const ALERTS_TABLE_NAME = 'alerts';
-const SELECT_TABLE_NAME = '#select';
 const UPDATING_RESULTS_TEXT_CLASS_NAME = "vscode-codeql__result-tables-updating-text";

 function getResultCount(resultSet: ResultSet): number {
@@ -75,23 +77,66 @@ export class ResultTables
    return resultSets;
  }

+  private getResultSetNames(resultSets: ResultSet[]): string[] {
+    if (this.props.parsedResultSets.t === 'ExtensionParsed') {
+      return this.props.parsedResultSets.resultSetNames.concat([ALERTS_TABLE_NAME]);
+    }
+    else {
+      return resultSets.map(resultSet => resultSet.schema.name);
+    }
+  }
+
+  /**
+   * Holds if we have a result set obtained from the extension that came
+   * from the ExtensionParsed branch of ParsedResultSets. This is evidence
+   * that the user has the experimental flag turned on that allows extension-side
+   * bqrs parsing.
+   */
+  paginationAllowed(): boolean {
+    return this.props.parsedResultSets.t === 'ExtensionParsed';
+  }
+
+  /**
+   * Holds if we actually should show pagination interface right now. This is
+   * still false for the time being when we're viewing alerts.
+   */
+  paginationEnabled(): boolean {
+    return this.paginationAllowed() &&
+      this.props.parsedResultSets.selectedTable !== ALERTS_TABLE_NAME &&
+      this.state.selectedTable !== ALERTS_TABLE_NAME;
+  }
+
  constructor(props: ResultTablesProps) {
    super(props);

-    this.state = {
-      // Get the result set that should be displayed by default
-      selectedTable: ResultTables.getDefaultResultSet(this.getResultSets())
-    };
-  }
+    const selectedTable = props.parsedResultSets.selectedTable || getDefaultResultSet(this.getResultSets());

-  private static getDefaultResultSet(resultSets: readonly ResultSet[]): string {
-    const resultSetNames = resultSets.map(resultSet => resultSet.schema.name);
-    // Choose first available result set from the array
-    return [ALERTS_TABLE_NAME, SELECT_TABLE_NAME, resultSets[0].schema.name].filter(resultSetName => resultSetNames.includes(resultSetName))[0];
+    let selectedPage: string;
+    switch (props.parsedResultSets.t) {
+      case 'ExtensionParsed':
+        selectedPage = (props.parsedResultSets.pageNumber + 1) + '';
+        break;
+      case 'WebviewParsed':
+        selectedPage = '';
+        break;
+    }
+
+    this.state = { selectedTable, selectedPage };
  }

  private onTableSelectionChange = (event: React.ChangeEvent<HTMLSelectElement>): void => {
-    this.setState({ selectedTable: event.target.value });
+    const selectedTable = event.target.value;
+    const fetchPageFromExtension = this.paginationAllowed() && selectedTable !== ALERTS_TABLE_NAME;
+
+    if (fetchPageFromExtension) {
+      vscode.postMessage({
+        t: 'changePage',
+        pageNumber: 0,
+        selectedTable
+      });
+    }
+    else
+      this.setState({ selectedTable });
  }

  private alertTableExtras(): JSX.Element | undefined {
@@ -118,24 +163,81 @@ export class ResultTables
    </div>;
  }

+  getOffset(): number {
+    const { parsedResultSets } = this.props;
+    switch (parsedResultSets.t) {
+      case 'ExtensionParsed':
+        return parsedResultSets.pageNumber * RAW_RESULTS_PAGE_SIZE;
+      case 'WebviewParsed':
+        return 0;
+    }
+  }
+
+  renderPageButtons(resultSets: ExtensionParsedResultSets): JSX.Element {
+    const selectedTable = this.state.selectedTable;
+    const onChange = (e: React.ChangeEvent<HTMLInputElement>) => {
+      this.setState({ selectedPage: e.target.value });
+    };
+    const choosePage = (input: string) => {
+      const pageNumber = parseInt(input);
+      if (pageNumber !== undefined && !isNaN(pageNumber)) {
+        const actualPageNumber = Math.max(0, Math.min(pageNumber - 1, resultSets.numPages - 1));
+        vscode.postMessage({
+          t: 'changePage',
+          pageNumber: actualPageNumber,
+          selectedTable,
+        });
+      }
+    };
+
+    const prevPage = (e: React.MouseEvent<HTMLButtonElement, MouseEvent>) => {
+      vscode.postMessage({
+        t: 'changePage',
+        pageNumber: Math.max(resultSets.pageNumber - 1, 0),
+        selectedTable,
+      });
+    };
+    const nextPage = (e: React.MouseEvent<HTMLButtonElement, MouseEvent>) => {
+      vscode.postMessage({
+        t: 'changePage',
+        pageNumber: Math.min(resultSets.pageNumber + 1, resultSets.numPages - 1),
+        selectedTable,
+      });
+    };
+    return <span>
+      <button onClick={prevPage} >&lt;</button>
+      <input value={this.state.selectedPage} onChange={onChange}
+        onBlur={e => choosePage(e.target.value)}
+        onKeyDown={e => { if (e.keyCode === 13) choosePage((e.target as HTMLInputElement).value); }}
+      />
+      <button value=">" onClick={nextPage} >&gt;</button>
+    </span>;
+  }
+
+  renderButtons(): JSX.Element {
+    if (this.props.parsedResultSets.t === 'ExtensionParsed' && this.paginationEnabled())
+      return this.renderPageButtons(this.props.parsedResultSets);
+    else
+      return <span />;
+  }
+
  render(): React.ReactNode {
    const { selectedTable } = this.state;
    const resultSets = this.getResultSets();
+    const resultSetNames = this.getResultSetNames(resultSets);

    const resultSet = resultSets.find(resultSet => resultSet.schema.name == selectedTable);
    const nonemptyRawResults = resultSets.some(resultSet => resultSet.t == 'RawResultSet' && resultSet.rows.length > 0);
    const numberOfResults = resultSet && renderResultCountString(resultSet);

+    const resultSetOptions =
+      resultSetNames.map(name => <option key={name} value={name}>{name}</option>);
+
    return <div>
+      {this.renderButtons()}
      <div className={tableSelectionHeaderClassName}>
        <select value={selectedTable} onChange={this.onTableSelectionChange}>
-          {
-            resultSets.map(resultSet =>
-              <option key={resultSet.schema.name} value={resultSet.schema.name}>
-                {resultSet.schema.name}
-              </option>
-            )
-          }
+          {resultSetOptions}
        </select>
        {numberOfResults}
        {selectedTable === ALERTS_TABLE_NAME ? this.alertTableExtras() : undefined}
@@ -152,7 +254,8 @@ export class ResultTables
          resultsPath={this.props.resultsPath}
          sortState={this.props.sortStates.get(resultSet.schema.name)}
          nonemptyRawResults={nonemptyRawResults}
-          showRawResults={() => { this.setState({ selectedTable: SELECT_TABLE_NAME }); }} />
+          showRawResults={() => { this.setState({ selectedTable: SELECT_TABLE_NAME }); }}
+          offset={this.getOffset()} />
      }
    </div>;
  }
--- a/extensions/ql-vscode/src/view/results.tsx
+++ b/extensions/ql-vscode/src/view/results.tsx
@@ -1,12 +1,13 @@
 import * as React from 'react';
 import * as Rdom from 'react-dom';
 import * as bqrs from 'semmle-bqrs';
-import { ElementBase, PrimitiveColumnValue, PrimitiveTypeKind, ResultSetSchema, tryGetResolvableLocation } from 'semmle-bqrs';
+import { ElementBase, PrimitiveColumnValue, PrimitiveTypeKind, tryGetResolvableLocation } from 'semmle-bqrs';
 import { assertNever } from '../helpers-pure';
 import { DatabaseInfo, FromResultsViewMsg, Interpretation, IntoResultsViewMsg, SortedResultSetInfo, RawResultsSortState, NavigatePathMsg, QueryMetadata, ResultsPaths } from '../interface-types';
 import { EventHandlers as EventHandlerList } from './event-handler-list';
 import { ResultTables } from './result-tables';
-import { RawResultSet, ResultValue, ResultRow } from '../adapt';
+import { ResultValue, ResultRow, ParsedResultSets } from '../adapt';
+import { ResultSet } from '../interface-utils';

 /**
 * results.tsx
@@ -24,13 +25,6 @@ interface VsCodeApi {
 declare const acquireVsCodeApi: () => VsCodeApi;
 export const vscode = acquireVsCodeApi();

-export type RawTableResultSet = { t: 'RawResultSet' } & RawResultSet;
-export type PathTableResultSet = { t: 'SarifResultSet'; readonly schema: ResultSetSchema; name: string } & Interpretation;
-
-export type ResultSet =
-  | RawTableResultSet
-  | PathTableResultSet;
-
 async function* getChunkIterator(response: Response): AsyncIterableIterator<Uint8Array> {
  if (!response.ok) {
    throw new Error(`Failed to load results: (${response.status}) ${response.statusText}`);
@@ -107,8 +101,8 @@ async function parseResultSets(response: Response): Promise<readonly ResultSet[]
 }

 interface ResultsInfo {
+  parsedResultSets: ParsedResultSets;
  resultsPath: string;
-  resultSets: ResultSet[] | undefined;
  origResultsPaths: ResultsPaths;
  database: DatabaseInfo;
  interpretation: Interpretation | undefined;
@@ -169,7 +163,7 @@ class App extends React.Component<{}, ResultsViewState> {
      case 'setState':
        this.updateStateWithNewResultsInfo({
          resultsPath: msg.resultsPath,
-          resultSets: msg.resultSets?.map(x => ({ t: 'RawResultSet', ...x })),
+          parsedResultSets: msg.parsedResultSets,
          origResultsPaths: msg.origResultsPaths,
          sortedResultsMap: new Map(Object.entries(msg.sortedResultsMap)),
          database: msg.database,
@@ -221,6 +215,16 @@ class App extends React.Component<{}, ResultsViewState> {
    });
  }

+  private async getResultSets(resultsInfo: ResultsInfo): Promise<readonly ResultSet[]> {
+    const parsedResultSets = resultsInfo.parsedResultSets;
+    switch (parsedResultSets.t) {
+      case 'WebviewParsed': return await this.fetchResultSets(resultsInfo);
+      case 'ExtensionParsed': {
+        return [{ t: 'RawResultSet', ...parsedResultSets.resultSet }];
+      }
+    }
+  }
+
  private async loadResults(): Promise<void> {
    const resultsInfo = this.state.nextResultsInfo;
    if (resultsInfo === null) {
@@ -230,7 +234,7 @@ class App extends React.Component<{}, ResultsViewState> {
    let results: Results | null = null;
    let statusText = '';
    try {
-      const resultSets = resultsInfo.resultSets || await this.getResultSets(resultsInfo);
+      const resultSets = await this.getResultSets(resultsInfo);
      results = {
        resultSets,
        database: resultsInfo.database,
@@ -265,7 +269,11 @@ class App extends React.Component<{}, ResultsViewState> {
    });
  }

-  private async getResultSets(resultsInfo: ResultsInfo): Promise<readonly ResultSet[]> {
+  /**
+   * This is deprecated, because it calls `fetch`. We are moving
+   * towards doing all bqrs parsing in the extension.
+   */
+  private async fetchResultSets(resultsInfo: ResultsInfo): Promise<readonly ResultSet[]> {
    const unsortedResponse = await fetch(resultsInfo.resultsPath);
    const unsortedResultSets = await parseResultSets(unsortedResponse);
    return Promise.all(unsortedResultSets.map(async unsortedResultSet => {
@@ -291,7 +299,10 @@ class App extends React.Component<{}, ResultsViewState> {
  render(): JSX.Element {
    const displayedResults = this.state.displayedResults;
    if (displayedResults.results !== null && displayedResults.resultsInfo !== null) {
-      return <ResultTables rawResultSets={displayedResults.results.resultSets}
+      const parsedResultSets = displayedResults.resultsInfo.parsedResultSets;
+      return <ResultTables
+        parsedResultSets={parsedResultSets}
+        rawResultSets={displayedResults.results.resultSets}
        interpretation={displayedResults.resultsInfo ? displayedResults.resultsInfo.interpretation : undefined}
        database={displayedResults.results.database}
        origResultsPaths={displayedResults.resultsInfo.origResultsPaths}