Improve resolve handling

- Also set DocumentSelector when preprocessing in order to support same
handler type instances supporting different languages.

Author: bjorkstromm

src/Server/HandlerCollection.cs

@@ -82,11 +82,13 @@ private HandlerDescriptor GetDescriptor(string method, Type implementedType, IJs
 
             var key = "default";
             // This protects against the case where class implements many, possibly conflicting, interfaces.
-            if (registration != null &&
-                typeof(TextDocumentRegistrationOptions).GetTypeInfo().IsAssignableFrom(registration) &&
+            if ((registration != null && typeof(TextDocumentRegistrationOptions).GetTypeInfo().IsAssignableFrom(registration) ||
+                registration == null && implementedType.GetTypeInfo().ImplementedInterfaces.Any(x => x.GetTypeInfo().IsGenericType && x.GetTypeInfo().GetGenericTypeDefinition() == typeof(ICanBeResolvedHandler<>))) &&
                 handler is IRegistration<TextDocumentRegistrationOptions> handlerRegistration)
             {
-                key = handlerRegistration.GetRegistrationOptions()?.DocumentSelector ?? key;
+                key = string.IsNullOrEmpty(handlerRegistration?.GetRegistrationOptions()?.DocumentSelector)
+                    ? key
+                    : handlerRegistration?.GetRegistrationOptions()?.DocumentSelector;
             }
 
             return new HandlerDescriptor(

src/Server/Matchers/ResolveCommandMatcher.cs

@@ -13,6 +13,7 @@ public class ResolveCommandMatcher : IHandlerMatcher, IHandlerPreProcessorMatche
     {
         private readonly ILogger _logger;
         internal static string PrivateHandlerTypeName = "$$___handlerType___$$";
+        internal static string PrivateHandlerKey = "$$___handlerKey___$$";
 
         public ResolveCommandMatcher(ILogger logger)
         {
@@ -30,10 +31,15 @@ public IEnumerable<ILspHandlerDescriptor> FindHandler(object parameters, IEnumer
             if (parameters is ICanBeResolved canBeResolved)
             {
                 string handlerType = null;
+                string handlerKey = null;
                 if (canBeResolved.Data != null && canBeResolved.Data.Type == JTokenType.Object)
+                {
                     handlerType = canBeResolved.Data?[PrivateHandlerTypeName]?.ToString();
+                    handlerKey = canBeResolved.Data?[PrivateHandlerKey]?.ToString();
+                }
 
-                if (string.IsNullOrWhiteSpace(handlerType))
+                if (string.IsNullOrWhiteSpace(handlerType) &&
+                    string.IsNullOrWhiteSpace(handlerKey))
                 {
                     foreach (var descriptor in descriptors)
                     {
@@ -65,7 +71,8 @@ public IEnumerable<ILspHandlerDescriptor> FindHandler(object parameters, IEnumer
                     _logger.LogTrace("Checking handler {Method}:{Handler}",
                         descriptor.Method,
                         descriptor.Handler.GetType().FullName);
-                    if (descriptor.Handler.GetType().FullName == handlerType || descriptor.HandlerType.FullName == handlerType)
+                    if ((descriptor.Handler.GetType().FullName == handlerType || descriptor.HandlerType.FullName == handlerType) &&
+                        ((descriptor is HandlerDescriptor handlerDescriptor) && handlerDescriptor.Key == handlerKey))
                     {
                         yield return descriptor;
                     }
@@ -120,8 +127,11 @@ public object Process(ILspHandlerDescriptor descriptor, object parameters)
 
         public object Process(ILspHandlerDescriptor descriptor, object parameters, object response)
         {
+            var registrationOptions = descriptor.Registration.RegisterOptions as TextDocumentRegistrationOptions;
+
             // Only pin the handler type, if we know the source handler (codelens) is also the resolver.
-            if (response is IEnumerable<ICanBeResolved> canBeResolveds &&
+            if (registrationOptions?.DocumentSelector != null &&
+                response is IEnumerable<ICanBeResolved> canBeResolveds &&
                 descriptor?.CanBeResolvedHandlerType?.GetTypeInfo().IsAssignableFrom(descriptor.Handler.GetType()) == true)
             {
                 _logger.LogTrace("Updating Resolve items with wrapped data for {Method}:{Handler}",
@@ -135,6 +145,7 @@ public object Process(ILspHandlerDescriptor descriptor, object parameters, objec
                     var data = new JObject();
                     data["data"] = item.Data;
                     data[PrivateHandlerTypeName] = descriptor.Handler.GetType().FullName;
+                    data[PrivateHandlerKey] = registrationOptions.DocumentSelector.ToString();
                     item.Data = data;
                 }
             }

test/Lsp.Tests/HandlerResolverTests.cs

@@ -192,6 +192,22 @@ public void Should_DealWithClassesThatImplementMultipleHandlers_WithoutConflicti
             descriptor.Key.Should().Be("default");
         }
 
+        [Fact]
+        public void Should_DealWithClassesThatImplementMultipleHandlers_BySettingKeyAccordingly()
+        {
+            var codeLensHandler = Substitute.For(new Type[] { typeof(ICodeLensHandler), typeof(ICodeLensResolveHandler) }, new object[0]);
+            ((ICodeLensHandler)codeLensHandler).GetRegistrationOptions()
+                .Returns(new CodeLensRegistrationOptions() {
+                    DocumentSelector = new DocumentSelector(DocumentFilter.ForLanguage("foo"))
+                });
+
+            var handler = new HandlerCollection();
+            handler.Add(codeLensHandler as IJsonRpcHandler);
+
+            var descriptor = handler._handlers.Select(x => x.Key);
+            descriptor.ShouldAllBeEquivalentTo(new [] { "[foo]", "[foo]" });
+        }
+
         public static IEnumerable<object[]> Should_DealWithClassesThatImplementMultipleHandlers_WithoutConflictingRegistrations_Data()
         {
             var codeLensHandler = Substitute.For(new Type[] { typeof(ICodeLensHandler), typeof(ICodeLensResolveHandler) }, new object[0]);