Better handle completion when the display text is not in the final result

src/OmniSharp.Abstractions/Models/v1/Completion/CompletionItem.cs

@@ -1,6 +1,6 @@
 #nullable enable
 
-using System.Collections.Immutable;
+using System.Collections.Generic;
 
 namespace OmniSharp.Models.v1.Completion
 {
@@ -23,7 +23,7 @@ public class CompletionItem
         /// <summary>
         /// Tags for this completion item
         /// </summary>
-        public ImmutableArray<CompletionItemTag>? Tags { get; set; }
+        public IReadOnlyList<CompletionItemTag>? Tags { get; set; }
 
         /// <summary>
         /// A human-readable string with additional information
@@ -69,7 +69,7 @@ public class CompletionItem
         /// An optional set of characters that when pressed while this completion is active will accept it first and
         /// then type that character.
         /// </summary>
-        public ImmutableArray<char>? CommitCharacters { get; set; }
+        public IReadOnlyList<char>? CommitCharacters { get; set; }
 
         /// <summary>
         /// An optional array of additional text edits that are applied when
@@ -80,7 +80,7 @@ public class CompletionItem
         /// (for example adding an import statement at the top of the file if the completion item will
         /// insert an unqualified type).
         /// </summary>
-        public ImmutableArray<LinePositionSpanTextChange>? AdditionalTextEdits { get; set; }
+        public IReadOnlyList<LinePositionSpanTextChange>? AdditionalTextEdits { get; set; }
 
         /// <summary>
         /// Index in the completions list that this completion occurred.

src/OmniSharp.Abstractions/Models/v1/Completion/CompletionResponse.cs

@@ -1,6 +1,6 @@
 #nullable enable
 
-using System.Collections.Immutable;
+using System.Collections.Generic;
 
 namespace OmniSharp.Models.v1.Completion
 {
@@ -14,6 +14,6 @@ public class CompletionResponse
         /// <summary>
         /// The completion items.
         /// </summary>
-        public ImmutableArray<CompletionItem> Items { get; set; }
+        public IReadOnlyList<CompletionItem> Items { get; set; } = null!;
     }
 }

src/OmniSharp.Roslyn.CSharp/Services/Completion/CompletionService.cs

@@ -7,9 +7,10 @@
 using System.Linq;
 using System.Text;
 using System.Threading.Tasks;
-using System.Xml.Resolvers;
 using Microsoft.CodeAnalysis;
 using Microsoft.CodeAnalysis.Completion;
+using Microsoft.CodeAnalysis.CSharp;
+using Microsoft.CodeAnalysis.CSharp.Syntax;
 using Microsoft.CodeAnalysis.Tags;
 using Microsoft.CodeAnalysis.Text;
 using Microsoft.Extensions.Logging;
@@ -159,17 +160,19 @@ public async Task<CompletionResponse> Handle(CompletionRequest request)
             // the completion as done.
             bool seenUnimportedCompletions = false;
             bool expectingImportedItems = expandedItemsAvailable && _workspace.Options.GetOption(CompletionItemExtensions.ShowItemsFromUnimportedNamespaces, LanguageNames.CSharp) == true;
+            var syntax = await document.GetSyntaxTreeAsync();
 
             for (int i = 0; i < completions.Items.Length; i++)
             {
                 var completion = completions.Items[i];
                 var insertTextFormat = InsertTextFormat.PlainText;
-                ImmutableArray<LinePositionSpanTextChange>? additionalTextEdits = null;
+                IReadOnlyList<LinePositionSpanTextChange>? additionalTextEdits = null;
                 char sortTextPrepend = '0';
 
                 if (!completion.TryGetInsertionText(out string insertText))
                 {
-                    switch (completion.GetProviderName())
+                    string providerName = completion.GetProviderName();
+                    switch (providerName)
                     {
                         case CompletionItemExtensions.InternalsVisibleToCompletionProvider:
                             // The IVT completer doesn't add extra things before the completion
@@ -240,7 +243,7 @@ public async Task<CompletionResponse> Handle(CompletionRequest request)
                                 }
 
                                 int additionalEditEndOffset;
-                                (additionalTextEdits, additionalEditEndOffset) = GetAdditionalTextEdits(change, sourceText, typedSpan, completion.DisplayText, isImportCompletion: false);
+                                (additionalTextEdits, additionalEditEndOffset) = await GetAdditionalTextEdits(change, sourceText, (CSharpParseOptions)syntax!.Options, typedSpan, completion.DisplayText, providerName);
 
                                 // Now that we have the additional edit, adjust the rest of the new text
                                 (insertText, insertTextFormat) = getAdjustedInsertTextWithPosition(change, position, additionalEditEndOffset);
@@ -435,14 +438,16 @@ public async Task<CompletionResolveResponse> Handle(CompletionResolveRequest req
 
             request.Item.Documentation = textBuilder.ToString();
 
-            switch (lastCompletionItem.GetProviderName())
+            string providerName = lastCompletionItem.GetProviderName();
+            switch (providerName)
             {
                 case CompletionItemExtensions.ExtensionMethodImportCompletionProvider:
                 case CompletionItemExtensions.TypeImportCompletionProvider:
+                    var syntax = await document.GetSyntaxTreeAsync();
                     var sourceText = await document.GetTextAsync();
                     var typedSpan = completionService.GetDefaultCompletionListSpan(sourceText, position);
                     var change = await completionService.GetChangeAsync(document, lastCompletionItem, typedSpan);
-                    (request.Item.AdditionalTextEdits, _) = GetAdditionalTextEdits(change, sourceText, typedSpan, lastCompletionItem.DisplayText, isImportCompletion: true);
+                    (request.Item.AdditionalTextEdits, _) = await GetAdditionalTextEdits(change, sourceText, (CSharpParseOptions)syntax!.Options, typedSpan, lastCompletionItem.DisplayText, providerName);
                     break;
             }
 
@@ -452,19 +457,20 @@ public async Task<CompletionResolveResponse> Handle(CompletionResolveRequest req
             };
         }
 
-        private (ImmutableArray<LinePositionSpanTextChange> edits, int endOffset) GetAdditionalTextEdits(CompletionChange change, SourceText sourceText, TextSpan typedSpan, string completionDisplayText, bool isImportCompletion)
+        private async ValueTask<(IReadOnlyList<LinePositionSpanTextChange> edits, int endOffset)> GetAdditionalTextEdits(
+            CompletionChange change,
+            SourceText sourceText,
+            CSharpParseOptions parseOptions,
+            TextSpan typedSpan,
+            string completionDisplayText,
+            string providerName)
         {
             // We know the span starts before the text we're keying off of. So, break that
             // out into a separate edit. We need to cut out the space before the current word,
             // as the additional edit is not allowed to overlap with the insertion point.
             var additionalEditStartPosition = sourceText.Lines.GetLinePosition(change.TextChange.Span.Start);
             var additionalEditEndPosition = sourceText.Lines.GetLinePosition(typedSpan.Start - 1);
-            int additionalEditEndOffset = isImportCompletion
-                // Import completion will put the displaytext at the end of the line, override completion will
-                // put it at the front.
-                ? change.TextChange.NewText!.LastIndexOf(completionDisplayText)
-                : change.TextChange.NewText!.IndexOf(completionDisplayText);
-
+            int additionalEditEndOffset = await getAdditionalTextEditEndOffset(change, sourceText, parseOptions, typedSpan, completionDisplayText, providerName);
             if (additionalEditEndOffset < 1)
             {
                 // The first index of this was either 0 and the edit span was wrong,
@@ -484,6 +490,44 @@ public async Task<CompletionResolveResponse> Handle(CompletionResolveRequest req
                 EndLine = additionalEditEndPosition.Line,
                 EndColumn = additionalEditEndPosition.Character,
             }), additionalEditEndOffset);
+
+            static async ValueTask<int> getAdditionalTextEditEndOffset(CompletionChange change, SourceText sourceText, CSharpParseOptions parseOptions, TextSpan typedSpan, string completionDisplayText, string providerName)
+            {
+                // For many simple cases, we can just find the first or last index of the completionDisplayText and that's good
+                // enough
+                int endOffset = (providerName == CompletionItemExtensions.ExtensionMethodImportCompletionProvider ||
+                                 providerName == CompletionItemExtensions.TypeImportCompletionProvider)
+                    // Import completion will put the displaytext at the end of the line, override completion will
+                    // put it at the front.
+                    ? change.TextChange.NewText!.LastIndexOf(completionDisplayText)
+                    : change.TextChange.NewText!.IndexOf(completionDisplayText);
+
+                if (endOffset > -1)
+                {
+                    return endOffset;
+                }
+
+                // The DisplayText wasn't in the final string. This can happen in a few cases:
+                //  * The override or partial method completion is involving types that need
+                //    to have a using added in the final version, and won't be fully qualified
+                //    as they were in the DisplayText
+                //  * Nullable context differences, such as if the thing you're overriding is
+                //    annotated but the final context being generated into does not have
+                //    annotations enabled.
+                // For these cases, we currently should only be seeing override or partial
+                // completions, as import completions don't have nullable annotations or
+                // fully-qualified types in their DisplayTexts. If that ever changes, we'll have
+                // to adjust the API here.
+                //
+                // In order to find the correct location here, we parse the change. The location
+                // of the name of the last method is the correct location in the string.
+                Debug.Assert(providerName == CompletionItemExtensions.OverrideCompletionProvider ||
+                             providerName == CompletionItemExtensions.PartialMethodCompletionProvider);
+
+                var parsedTree = CSharpSyntaxTree.ParseText(change.TextChange.NewText, parseOptions);
+                var lastMethodDecl = (await parsedTree.GetRootAsync()).DescendantNodes().OfType<MethodDeclarationSyntax>().Last();
+                return lastMethodDecl.Identifier.SpanStart;
+            }
         }
     }
 }