Add rough first-draft of document highlight provider (#409)

shared/activate.ts

@@ -240,6 +240,17 @@ function activateWithoutLSP(
     ctx.subscriptions.add(
         sourcegraph.languages.registerHoverProvider(selector, wrapper.hover())
     )
+
+    // Do not try to register this provider on pre-3.18 instances as it
+    // didn't exist.
+    if (sourcegraph.languages.registerDocumentHighlightProvider) {
+        ctx.subscriptions.add(
+            sourcegraph.languages.registerDocumentHighlightProvider(
+                selector,
+                wrapper.documentHighlights()
+            )
+        )
+    }
 }
 
 /**

shared/lsif/providers.test.ts

@@ -13,7 +13,7 @@ import {
 } from './providers'
 
 const doc = createStubTextDocument({
-    uri: 'https://sourcegraph.test/repo@rev/-/raw/foo.ts',
+    uri: 'git://repo@ref#/foo.ts',
     languageId: 'typescript',
     text: undefined,
 })
@@ -28,7 +28,11 @@ const pos = new sourcegraph.Position(5, 10)
 const range1 = new sourcegraph.Range(1, 2, 3, 4)
 const range2 = new sourcegraph.Range(2, 3, 4, 5)
 const range3 = new sourcegraph.Range(3, 4, 5, 6)
+const range4 = new sourcegraph.Range(4, 5, 6, 7)
+const range5 = new sourcegraph.Range(5, 6, 7, 8)
+const range6 = new sourcegraph.Range(6, 7, 8, 9)
 
+const resource0 = makeResource('repo', 'rev', '/foo.ts')
 const resource1 = makeResource('repo1', 'deadbeef1', '/a.ts')
 const resource2 = makeResource('repo2', 'deadbeef2', '/b.ts')
 const resource3 = makeResource('repo3', 'deadbeef3', '/c.ts')
@@ -309,6 +313,54 @@ describe('graphql providers', () => {
             )
         })
     })
+
+    describe('document highlights provider', () => {
+        it('should correctly parse result', async () => {
+            const queryGraphQLFn = sinon.spy<
+                QueryGraphQLFn<GenericLSIFResponse<ReferencesResponse | null>>
+            >(() =>
+                makeEnvelope({
+                    references: {
+                        nodes: [
+                            { resource: resource0, range: range1 },
+                            { resource: resource1, range: range2 },
+                            { resource: resource0, range: range3 },
+                            { resource: resource2, range: range4 },
+                            { resource: resource0, range: range5 },
+                            { resource: resource3, range: range6 },
+                        ],
+                        pageInfo: {},
+                    },
+                })
+            )
+
+            console.log(
+                await gatherValues(
+                    createProviders(queryGraphQLFn).documentHighlights(doc, pos)
+                )
+            )
+
+            assert.deepEqual(
+                await gatherValues(
+                    createProviders(queryGraphQLFn).documentHighlights(doc, pos)
+                ),
+                [[{ range: range1 }, { range: range3 }, { range: range5 }]]
+            )
+        })
+
+        it('should deal with empty payload', async () => {
+            const queryGraphQLFn = sinon.spy<
+                QueryGraphQLFn<GenericLSIFResponse<ReferencesResponse | null>>
+            >(() => makeEnvelope())
+
+            assert.deepEqual(
+                await gatherValues(
+                    createProviders(queryGraphQLFn).documentHighlights(doc, pos)
+                ),
+                [null]
+            )
+        })
+    })
 })
 
 async function gatherValues<T>(g: AsyncGenerator<T>): Promise<T[]> {

shared/lsif/providers.ts

@@ -6,6 +6,7 @@ import { asyncGeneratorFromPromise, concat } from '../util/ix'
 import { parseGitURI } from '../util/uri'
 import { LocationConnectionNode, nodeToLocation } from './conversion'
 import { Logger } from '../logging'
+import { isDefined } from '../util/helpers'
 
 /**
  * The maximum number of chained GraphQL requests to make for a single
@@ -57,6 +58,9 @@ export function createGraphQLProviders(
         definition: asyncGeneratorFromPromise(definition(queryGraphQL)),
         references: references(queryGraphQL),
         hover: asyncGeneratorFromPromise(hover(queryGraphQL)),
+        documentHighlights: asyncGeneratorFromPromise(
+            documentHighlights(queryGraphQL)
+        ),
     }
 }
 
@@ -287,6 +291,79 @@ function hover(
     }
 }
 
+/** Retrieve references ranges of the current hover position to highlight. */
+export function documentHighlights(
+    queryGraphQL: QueryGraphQLFn<GenericLSIFResponse<ReferencesResponse | null>>
+): (
+    doc: sourcegraph.TextDocument,
+    position: sourcegraph.Position
+) => Promise<sourcegraph.DocumentHighlight[] | null> {
+    return async (
+        doc: sourcegraph.TextDocument,
+        position: sourcegraph.Position
+    ): Promise<sourcegraph.DocumentHighlight[] | null> => {
+        const query = `
+            query ReferencesForDocumentHighlights($repository: String!, $commit: String!, $path: String!, $line: Int!, $character: Int!) {
+                repository(name: $repository) {
+                    commit(rev: $commit) {
+                        blob(path: $path) {
+                            lsif {
+                                references(line: $line, character: $character) {
+                                    nodes {
+                                        resource {
+                                            path
+                                            repository {
+                                                name
+                                            }
+                                            commit {
+                                                oid
+                                            }
+                                        }
+                                        range {
+                                            start {
+                                                line
+                                                character
+                                            }
+                                            end {
+                                                line
+                                                character
+                                            }
+                                        }
+                                    }
+                                }
+                            }
+                        }
+                    }
+                }
+            }
+        `
+
+        // Make the request for the page starting at the after cursor
+        const lsifObj: ReferencesResponse | null = await queryLSIF(
+            {
+                doc,
+                position,
+                query,
+            },
+            queryGraphQL
+        )
+        if (!lsifObj) {
+            return null
+        }
+
+        const { path: targetPath } = parseGitURI(new URL(doc.uri))
+
+        const {
+            references: { nodes },
+        } = lsifObj
+
+        return nodes
+            .filter(({ resource: { path } }) => path === targetPath)
+            .map(({ range }) => range && { range })
+            .filter(isDefined)
+    }
+}
+
 /**
  * Perform an LSIF request to the GraphQL API.
  *

shared/providers.test.ts

@@ -7,6 +7,7 @@ import {
     createDefinitionProvider,
     createHoverProvider,
     createReferencesProvider,
+    createDocumentHighlightProvider,
 } from './providers'
 
 const doc = createStubTextDocument({
@@ -211,6 +212,25 @@ describe('createHoverProvider', () => {
     })
 })
 
+describe('createDocumentHighlightProvider', () => {
+    it('uses LSIF document highlights', async () => {
+        const result = createDocumentHighlightProvider(
+            () =>
+                asyncGeneratorFromValues([
+                    [{ range: range1 }, { range: range2 }],
+                ]),
+            () => asyncGeneratorFromValues([]),
+            () => asyncGeneratorFromValues([])
+        ).provideDocumentHighlights(doc, pos) as Observable<
+            sourcegraph.DocumentHighlight[]
+        >
+
+        assert.deepStrictEqual(await gatherValues(result), [
+            [{ range: range1 }, { range: range2 }],
+        ])
+    })
+})
+
 async function* asyncGeneratorFromValues<P>(
     source: P[]
 ): AsyncGenerator<P, void, undefined> {

shared/providers.ts

@@ -14,12 +14,14 @@ export interface Providers {
     definition: DefinitionProvider
     references: ReferencesProvider
     hover: HoverProvider
+    documentHighlights: DocumentHighlightProvider
 }
 
 export interface SourcegraphProviders {
     definition: sourcegraph.DefinitionProvider
     references: sourcegraph.ReferenceProvider
     hover: sourcegraph.HoverProvider
+    documentHighlights: sourcegraph.DocumentHighlightProvider
 }
 
 export type DefinitionProvider = (
@@ -38,16 +40,23 @@ export type HoverProvider = (
     pos: sourcegraph.Position
 ) => AsyncGenerator<sourcegraph.Hover | null, void, undefined>
 
+export type DocumentHighlightProvider = (
+    doc: sourcegraph.TextDocument,
+    pos: sourcegraph.Position
+) => AsyncGenerator<sourcegraph.DocumentHighlight[] | null, void, undefined>
+
 export const noopProviders = {
     definition: noopAsyncGenerator,
     references: noopAsyncGenerator,
     hover: noopAsyncGenerator,
+    documentHighlights: noopAsyncGenerator,
 }
 
 export interface ProviderWrapper {
     definition: DefinitionWrapper
     references: ReferencesWrapper
     hover: HoverWrapper
+    documentHighlights: DocumentHighlightWrapper
 }
 
 export type DefinitionWrapper = (
@@ -62,6 +71,10 @@ export type HoverWrapper = (
     provider?: HoverProvider
 ) => sourcegraph.HoverProvider
 
+export type DocumentHighlightWrapper = (
+    provider?: DocumentHighlightProvider
+) => sourcegraph.DocumentHighlightProvider
+
 export class NoopProviderWrapper implements ProviderWrapper {
     public definition = (
         provider?: DefinitionProvider
@@ -97,6 +110,18 @@ export class NoopProviderWrapper implements ProviderWrapper {
                 ? observableFromAsyncIterator(() => provider(doc, pos))
                 : NEVER,
     })
+
+    public documentHighlights = (
+        provider?: DocumentHighlightProvider
+    ): sourcegraph.DocumentHighlightProvider => ({
+        provideDocumentHighlights: (
+            doc: sourcegraph.TextDocument,
+            pos: sourcegraph.Position
+        ) =>
+            provider
+                ? observableFromAsyncIterator(() => provider(doc, pos))
+                : NEVER,
+    })
 }
 
 /**
@@ -146,6 +171,16 @@ export function createProviderWrapper(
             wrapped.hover = provider
             return provider
         },
+
+        documentHighlights: (lspProvider?: DocumentHighlightProvider) => {
+            const provider = createDocumentHighlightProvider(
+                lsifProviders.documentHighlights,
+                searchProviders.documentHighlights,
+                lspProvider
+            )
+            wrapped.documentHighlights = provider
+            return provider
+        },
     }
 }
 
@@ -355,6 +390,39 @@ export function createHoverProvider(
     }
 }
 
+/**
+ * Creates a document highlight provider.
+ *
+ * @param lsifProvider The LSIF-based document highlight provider.
+ * @param searchProvider The search-based document highlight provider.
+ * @param lspProvider An optional LSP-based document highlight provider.
+ */
+export function createDocumentHighlightProvider(
+    lsifProvider: DocumentHighlightProvider,
+    searchProvider: DocumentHighlightProvider,
+    lspProvider?: DocumentHighlightProvider
+): sourcegraph.DocumentHighlightProvider {
+    return {
+        provideDocumentHighlights: wrapProvider(async function*(
+            doc: sourcegraph.TextDocument,
+            pos: sourcegraph.Position
+        ): AsyncGenerator<
+            sourcegraph.DocumentHighlight[] | null | undefined,
+            void,
+            undefined
+        > {
+            const emitter = new TelemetryEmitter()
+
+            for await (const lsifResult of lsifProvider(doc, pos)) {
+                if (lsifResult) {
+                    await emitter.emitOnce('lsifDocumentHighlight')
+                    yield lsifResult
+                }
+            }
+        }),
+    }
+}
+
 /**
  * Add a badge property to a single value or to a list of values. Returns the
  * modified result in the same shape as the input.

shared/search/providers.ts

@@ -293,10 +293,16 @@ export function createProviders(
         }
     }
 
+    const documentHighlights = (
+        doc: sourcegraph.TextDocument,
+        pos: sourcegraph.Position
+    ): Promise<sourcegraph.DocumentHighlight[] | null> => Promise.resolve(null)
+
     return {
         definition: asyncGeneratorFromPromise(definition),
         references: asyncGeneratorFromPromise(references),
         hover: asyncGeneratorFromPromise(hover),
+        documentHighlights: asyncGeneratorFromPromise(documentHighlights),
     }
 }