feat: add method to get locations with provider name

add TextDocumentLocationProviderRegistry.getLocationsWithProviderName method

Author: beyang

src/client/client.test.ts

@@ -84,10 +84,15 @@ describe('Client', () => {
 
                 // Request registration.
                 client.handleRegistrationRequest({
-                    registrations: [{ id: 'a', method: 'm', registerOptions: { a: 1 }, overwriteExisting: true }],
+                    registrations: [
+                        { id: 'a', method: 'm', registerOptions: { a: 1, providerName: '' }, overwriteExisting: true },
+                    ],
                 })
                 assert.deepStrictEqual(registerCalls, [
-                    { message: { method: 'm' }, data: { id: 'a', registerOptions: { a: 1 }, overwriteExisting: true } },
+                    {
+                        message: { method: 'm' },
+                        data: { id: 'a', registerOptions: { a: 1, providerName: '' }, overwriteExisting: true },
+                    },
                 ] as typeof registerCalls)
 
                 // Request unregistration.

src/client/client.ts

@@ -468,6 +468,7 @@ export class Client implements Unsubscribable {
             if (!options.documentSelector && this.options.documentSelector) {
                 options.documentSelector = this.options.documentSelector
             }
+            options.providerName = this.name
             const data: RegistrationData<any> = {
                 id: registration.id,
                 registerOptions: options,

src/client/features/common.test.ts

@@ -28,7 +28,7 @@ class TextDocumentFeature extends AbstractTextDocumentFeature<TextDocumentRegist
     }
 }
 
-const FIXTURE_REGISTER_OPTIONS: TextDocumentRegistrationOptions = { documentSelector: ['*'] }
+const FIXTURE_REGISTER_OPTIONS: TextDocumentRegistrationOptions = { documentSelector: ['*'], providerName: 'test' }
 
 describe('TextDocumentFeature', () => {
     describe('dynamic registration', () => {

src/client/features/hover.ts

@@ -42,7 +42,10 @@ export class TextDocumentHoverFeature extends Feature<TextDocumentRegistrationOp
         }
         this.register(this.messages, {
             id: uuidv4(),
-            registerOptions: { documentSelector },
+            registerOptions: {
+                documentSelector,
+                providerName: this.client.name,
+            },
         })
     }
 

src/client/features/location.test.ts

@@ -25,7 +25,7 @@ const create = <P extends TextDocumentPositionParams, F extends TextDocumentLoca
     registry: TextDocumentLocationProviderRegistry<P>
     feature: F
 } => {
-    const client = { options: { middleware: {} } } as Client
+    const client = { options: { middleware: {} }, name: 'test' } as Client
     const registry = new RegistryClass()
     const feature = new FeatureClass(client, registry)
     return { client, registry, feature }

src/client/features/location.ts

@@ -58,7 +58,10 @@ export abstract class TextDocumentLocationFeature<
         }
         this.register(this.messages, {
             id: uuidv4(),
-            registerOptions: { documentSelector },
+            registerOptions: {
+                documentSelector,
+                providerName: this.client.name,
+            },
         })
     }
 

src/client/features/textDocument.test.ts

@@ -46,7 +46,7 @@ describe('TextDocumentNotificationFeature', () => {
         }
     }
 
-    const FIXTURE_REGISTER_OPTIONS: TextDocumentRegistrationOptions = { documentSelector: ['*'] }
+    const FIXTURE_REGISTER_OPTIONS: TextDocumentRegistrationOptions = { documentSelector: ['*'], providerName: 'test' }
 
     describe('registration', () => {
         it('supports dynamic registration and unregistration', () => {

src/client/features/textDocument.ts

@@ -124,7 +124,10 @@ export class TextDocumentDidOpenFeature extends TextDocumentNotificationFeature<
         if (documentSelector && textDocumentSyncOptions && textDocumentSyncOptions.openClose) {
             this.register(this.messages, {
                 id: uuidv4(),
-                registerOptions: { documentSelector },
+                registerOptions: {
+                    documentSelector,
+                    providerName: this.client.name,
+                },
             })
         }
     }
@@ -171,7 +174,10 @@ export class TextDocumentDidCloseFeature extends TextDocumentNotificationFeature
         if (documentSelector && textDocumentSyncOptions && textDocumentSyncOptions.openClose) {
             this.register(this.messages, {
                 id: uuidv4(),
-                registerOptions: { documentSelector },
+                registerOptions: {
+                    documentSelector,
+                    providerName: this.client.name,
+                },
             })
         }
     }

src/environment/providers/location.ts

@@ -1,9 +1,9 @@
-import { combineLatest, from, Observable } from 'rxjs'
+import { combineLatest, from, Observable, Unsubscribable } from 'rxjs'
 import { map, switchMap } from 'rxjs/operators'
 import { Location } from 'vscode-languageserver-types'
 import { ReferenceParams, TextDocumentPositionParams, TextDocumentRegistrationOptions } from '../../protocol'
 import { FeatureProviderRegistry } from './registry'
-import { flattenAndCompact } from './util'
+import { flattenAndCompact, flattenAndCompactNonNull } from './util'
 
 /**
  * Function signature for retrieving related locations given a location (e.g., definition, implementation, and type
@@ -22,6 +22,33 @@ export class TextDocumentLocationProviderRegistry<
     public getLocation(params: P): Observable<L | L[] | null> {
         return getLocation<P, L>(this.providers, params)
     }
+
+    public getLocationsWithProviderName(params: P): Observable<{ providerName: string; location: L }[] | null> {
+        return getLocationsWithProviderName<P, L>(this.providersWithName, params)
+    }
+
+    public readonly providersWithName: Observable<
+        { name: string; provider: ProvideTextDocumentLocationSignature<P, L> }[]
+    > = this.entries.pipe(
+        map(providers =>
+            providers.map(({ provider, registrationOptions }) => ({
+                name: registrationOptions.providerName,
+                provider,
+            }))
+        )
+    )
+
+    public registerProvider(
+        registrationOptions: TextDocumentRegistrationOptions,
+        provider: ProvideTextDocumentLocationSignature<P, L>
+    ): Unsubscribable {
+        if (!registrationOptions.providerName) {
+            // Input validation is necessary here, because registrationOptions is deserialized from
+            // the registration request and this is the first place where where we have the provider type.
+            throw new Error('registered provider must have name')
+        }
+        return super.registerProvider(registrationOptions, provider)
+    }
 }
 
 /**
@@ -62,6 +89,35 @@ export function getLocations<
     )
 }
 
+export function getLocationsWithProviderName<
+    P extends TextDocumentPositionParams = TextDocumentPositionParams,
+    L extends Location = Location
+>(
+    providersWithName: Observable<{ name: string; provider: ProvideTextDocumentLocationSignature<P, L> }[]>,
+    params: P
+): Observable<{ providerName: string; location: L }[]> {
+    return providersWithName.pipe(
+        switchMap(async providersWithName => {
+            const resultsByProvider = await Promise.all(
+                providersWithName.map(async ({ provider, name }) => {
+                    const providerRes = await provider(params).toPromise()
+                    return flattenAndCompactNonNull([providerRes]).map(loc => ({
+                        providerName: name,
+                        location: loc,
+                    }))
+                })
+            )
+            const flattenedResults: { providerName: string; location: L }[] = []
+            for (const providerResults of resultsByProvider) {
+                for (const res of providerResults) {
+                    flattenedResults.push(res)
+                }
+            }
+            return flattenedResults
+        })
+    )
+}
+
 /**
  * Provides reference results from all extensions.
  *

src/environment/providers/registry.ts

@@ -9,7 +9,7 @@ interface Entry<O, P> {
 
 /** Base class for provider registries for features. */
 export abstract class FeatureProviderRegistry<O, P> {
-    private entries = new BehaviorSubject<Entry<O, P>[]>([])
+    protected entries = new BehaviorSubject<Entry<O, P>[]>([])
 
     public constructor(initialEntries?: Entry<O, P>[]) {
         if (initialEntries) {

src/environment/providers/util.ts

@@ -8,3 +8,8 @@ export function flattenAndCompact<T>(value: (T | T[] | null)[] | null): T[] | nu
     const merged = flatten(compact(value))
     return merged.length === 0 ? null : merged
 }
+
+/** Flattens and compacts the argument. If it is null or if the result is empty, it returns null. */
+export function flattenAndCompactNonNull<T>(value: (T | T[] | null)[] | null): T[] {
+    return value ? flatten(compact(value)) : []
+}

src/protocol/textDocument.ts

@@ -358,6 +358,11 @@ export interface TextDocumentRegistrationOptions {
      * the document selector provided on the client side will be used.
      */
     documentSelector: DocumentSelector | null
+
+    /**
+     * User-visible name of the of the provider being registered.
+     */
+    providerName: string
 }
 
 /**