Merge pull request #14919 from CyrusNajmabadi/moveTypeImports

Only remove imports when moving a type, if the imports moved to the new file.

Fixes #12730

src/Compilers/Core/Portable/InternalUtilities/EnumerableExtensions.cs

@@ -307,11 +307,6 @@ private static class Comparisons<T> where T : IComparable<T>
             public static readonly IComparer<T> Comparer = Comparer<T>.Create(CompareTo);
         }
 
-        private static class Functions<T>
-        {
-            public static readonly Func<T, T> Identity = t => t;
-        }
-
         public static bool IsSorted<T>(this IEnumerable<T> enumerable, IComparer<T> comparer)
         {
             using (var e = enumerable.GetEnumerator())
@@ -393,4 +388,14 @@ public static IComparer<T> ToComparer<T>(this Comparison<T> comparison)
             return Comparer<T>.Create(comparison);
         }
     }
+
+    /// <summary>
+    /// Cached versions of commonly used delegates.
+    /// </summary>
+    /// <typeparam name="T"></typeparam>
+    internal static class Functions<T>
+    {
+        public static readonly Func<T, T> Identity = t => t;
+        public static readonly Func<T, bool> True = t => true;
+    }
 }
\ No newline at end of file

src/EditorFeatures/CSharpTest/CodeActions/MoveType/MoveTypeTests.MoveToNewFile.cs

@@ -130,6 +130,7 @@ class Class2 { }";
 
             var codeAfterMove =
 @"// Banner Text
+using System;
 
 class Class2 { }";
 
@@ -708,5 +709,43 @@ class InnerClass3
 }";
             await TestMoveTypeToNewFileAsync(code, codeAfterMove, expectedDocumentName, destinationDocumentText);
         }
+
+        [WpfFact, Trait(Traits.Feature, Traits.Features.CodeActionsMoveType)]
+        public async Task MoveTypeUsings1()
+        {
+            var code =
+@"
+// Only used by inner type.
+using System;
+
+// Unused by both types.
+using System.Collections;
+
+class Outer { 
+    [||]class Inner {
+        DateTime d;
+    }
+}";
+            var codeAfterMove = @"
+// Unused by both types.
+using System.Collections;
+
+partial class Outer {
+}";
+
+            var expectedDocumentName = "Inner.cs";
+            var destinationDocumentText =
+@"
+// Only used by inner type.
+using System;
+
+partial class Outer {
+    class Inner { 
+        DateTime d;
+    }
+}";
+
+            await TestMoveTypeToNewFileAsync(code, codeAfterMove, expectedDocumentName, destinationDocumentText);
+        }
     }
 }
\ No newline at end of file

src/EditorFeatures/VisualBasicTest/CodeActions/MoveType/MoveTypeTests.MoveToNewFile.vb

@@ -7,10 +7,10 @@ Namespace Microsoft.CodeAnalysis.Editor.VisualBasic.UnitTests.CodeRefactorings.M
         <WpfFact, Trait(Traits.Feature, Traits.Features.CodeActionsMoveType)>
         Public Async Function TestMissing_OnMatchingFileName() As Task
             Dim code =
-<File>
+"
 [||]Class test1
 End Class
-</File>.ConvertTestSourceTag()
+"
 
             Await TestMissingAsync(code)
         End Function
@@ -18,12 +18,12 @@ End Class
         <WpfFact, Trait(Traits.Feature, Traits.Features.CodeActionsMoveType)>
         Public Async Function TestMissing_Nested_OnMatchingFileName_Simple() As Task
             Dim code =
-<File>
+"
 Class Outer
     [||]Class test1
     End Class
 End Class
-</File>.ConvertTestSourceTag()
+"
 
             Await TestMissingAsync(code)
     End Function
@@ -31,79 +31,78 @@ End Class
     <WpfFact, Trait(Traits.Feature, Traits.Features.CodeActionsMoveType)>
     Public Async Function MultipleTypesInFileWithNoContainerNamespace() As Task
         Dim code =
-<File>
+"
 [||]Class Class1
 End Class
 
 Class Class2
 End Class
-</File>
+"
             Dim codeAfterMove =
-<File>
+"
 Class Class2
 End Class
-</File>
+"
             Dim expectedDocumentName = "Class1.vb"
 
             Dim destinationDocumentText =
-<File>
+"
 Class Class1
 End Class
-</File>
+"
             Await TestMoveTypeToNewFileAsync(code, codeAfterMove, expectedDocumentName, destinationDocumentText)
         End Function
 
         <WpfFact, Trait(Traits.Feature, Traits.Features.CodeActionsMoveType)>
         Public Async Function MoveNestedTypeToNewFile_Simple() As Task
             Dim code =
-<File>
+"
 Public Class Class1
     Class Class2[||]
     End Class
 End Class
-</File>
+"
             Dim codeAfterMove =
-<File>
+"
 Public Partial Class Class1
 End Class
-</File>
+"
             Dim expectedDocumentName = "Class2.vb"
 
             Dim destinationDocumentText =
-<File>
+"
 Public Partial Class Class1
     Class Class2
-
     End Class
 End Class
-</File>
+"
             Await TestMoveTypeToNewFileAsync(code, codeAfterMove, expectedDocumentName, destinationDocumentText)
         End Function
 
         <WpfFact, Trait(Traits.Feature, Traits.Features.CodeActionsMoveType)>
         Public Async Function MoveNestedTypeToNewFile_Simple_DottedName() As Task
             Dim code =
-<File>
+"
 Public Class Class1
     Class Class2[||]
     End Class
 End Class
-</File>
+"
             Dim codeAfterMove =
-<File>
+"
 Public Partial Class Class1
 End Class
-</File>
+"
             Dim expectedDocumentName = "Class1.Class2.vb"
 
             Dim destinationDocumentText =
-<File>
+"
 Public Partial Class Class1
     Class Class2
 
     End Class
 End Class
-</File>
+"
             Await TestMoveTypeToNewFileAsync(code, codeAfterMove, expectedDocumentName, destinationDocumentText, index:=1)
         End Function
 
@@ -111,33 +110,75 @@ End Class
         <WpfFact, Trait(Traits.Feature, Traits.Features.CodeActionsMoveType)>
         Public Async Function MoveNestedTypeToNewFile_RemoveComments() As Task
             Dim code =
-<File>
+"
 ''' Outer comment
 Public Class Class1
     ''' Inner comment
     Class Class2[||]
     End Class
 End Class
-</File>
+"
             Dim codeAfterMove =
-<File>
+"
 ''' Outer comment
 Public Partial Class Class1
 End Class
-</File>
+"
             Dim expectedDocumentName = "Class1.Class2.vb"
 
             Dim destinationDocumentText =
-<File>
+"
 Public Partial Class Class1
     ''' Inner comment
     Class Class2
     End Class
 End Class
-</File>
+"
             Await TestMoveTypeToNewFileAsync(
                 code, codeAfterMove, expectedDocumentName, destinationDocumentText,
                 index:=1, compareTokens:=False)
         End Function
+
+        <WpfFact, Trait(Traits.Feature, Traits.Features.CodeActionsMoveType)>
+        Public Async Function TestImports() As Task
+            Dim code =
+"
+' Used only by inner
+Imports System
+
+' Not used
+Imports System.Collections
+
+Class Outer
+    [||]Class Inner
+        Sub M(d as DateTime)
+        End Sub
+    End Class
+End Class
+"
+            Dim codeAfterMove =
+"
+' Not used
+Imports System.Collections
+
+Partial Class Outer
+End Class
+"
+            Dim expectedDocumentName = "Inner.vb"
+
+            Dim destinationDocumentText =
+"
+' Used only by inner
+Imports System
+
+Partial Class Outer
+    Class Inner
+        Sub M(d as DateTime)
+        End Sub
+    End Class
+End Class
+"
+            Await TestMoveTypeToNewFileAsync(code, codeAfterMove, expectedDocumentName, destinationDocumentText)
+        End Function
     End Class
 End Namespace
\ No newline at end of file

src/EditorFeatures/VisualBasicTest/Diagnostics/RemoveUnnecessaryImports/RemoveUnnecessaryImportsTests.vb

 ' Copyright (c) Microsoft.  All Rights Reserved.  Licensed under the Apache License, Version 2.0.  See License.txt in the project root for license information.
 
-Option Strict Off
 Imports Microsoft.CodeAnalysis.CodeFixes
 Imports Microsoft.CodeAnalysis.Diagnostics
 Imports Microsoft.CodeAnalysis.Editor.UnitTests.Extensions

src/Features/CSharp/Portable/Diagnostics/Analyzers/CSharpRemoveUnnecessaryImportsDiagnosticAnalyzer.cs

@@ -28,7 +28,8 @@ protected override LocalizableString GetTitleAndMessageFormatForClassificationId
         protected override IEnumerable<SyntaxNode> GetUnnecessaryImports(
             SemanticModel semanticModel, SyntaxNode root, CancellationToken cancellationToken)
         {
-            return CSharpRemoveUnnecessaryImportsService.GetUnnecessaryImports(semanticModel, root, cancellationToken);
+            return CSharpRemoveUnnecessaryImportsService.GetUnnecessaryImportsShared(
+                semanticModel, root, predicate: null, cancellationToken: cancellationToken);
         }
 
         protected override IEnumerable<TextSpan> GetFixableDiagnosticSpans(

src/Features/CSharp/Portable/RemoveUnnecessaryImports/CSharpRemoveUnnecessaryImportsService.cs

 // Copyright (c) Microsoft.  All Rights Reserved.  Licensed under the Apache License, Version 2.0.  See License.txt in the project root for license information.
 
+using System;
 using System.Collections.Generic;
+using System.Collections.Immutable;
 using System.Composition;
 using System.Linq;
 using System.Threading;
@@ -9,23 +11,26 @@
 using Microsoft.CodeAnalysis.Formatting;
 using Microsoft.CodeAnalysis.Host.Mef;
 using Microsoft.CodeAnalysis.Internal.Log;
-using Microsoft.CodeAnalysis.Options;
 using Microsoft.CodeAnalysis.RemoveUnnecessaryImports;
 using Microsoft.CodeAnalysis.Shared.Extensions;
 using Microsoft.CodeAnalysis.Text;
+using Roslyn.Utilities;
 
 namespace Microsoft.CodeAnalysis.CSharp.RemoveUnnecessaryImports
 {
     [ExportLanguageService(typeof(IRemoveUnnecessaryImportsService), LanguageNames.CSharp), Shared]
     internal partial class CSharpRemoveUnnecessaryImportsService :
-        AbstractRemoveUnnecessaryImportsService<UsingDirectiveSyntax>, IRemoveUnnecessaryImportsService
+        AbstractRemoveUnnecessaryImportsService<UsingDirectiveSyntax>
     {
-        public static IEnumerable<SyntaxNode> GetUnnecessaryImports(SemanticModel semanticModel, SyntaxNode root, CancellationToken cancellationToken)
+        public static ImmutableArray<UsingDirectiveSyntax> GetUnnecessaryImportsShared(
+            SemanticModel semanticModel, SyntaxNode root, 
+            Func<SyntaxNode, bool> predicate, CancellationToken cancellationToken)
         {
+            predicate = predicate ?? Functions<SyntaxNode>.True;
             var diagnostics = semanticModel.GetDiagnostics(cancellationToken: cancellationToken);
             if (!diagnostics.Any())
             {
-                return null;
+                return ImmutableArray<UsingDirectiveSyntax>.Empty;
             }
 
             var unnecessaryImports = new HashSet<UsingDirectiveSyntax>();
@@ -36,26 +41,26 @@ public static IEnumerable<SyntaxNode> GetUnnecessaryImports(SemanticModel semant
                 {
                     var node = root.FindNode(diagnostic.Location.SourceSpan) as UsingDirectiveSyntax;
 
-                    if (node != null)
+                    if (node != null && predicate(node))
                     {
                         unnecessaryImports.Add(node);
                     }
                 }
             }
 
-            if (cancellationToken.IsCancellationRequested || !unnecessaryImports.Any())
-            {
-                return null;
-            }
-
-            return unnecessaryImports;
+            return unnecessaryImports.ToImmutableArray();
         }
 
-        public async Task<Document> RemoveUnnecessaryImportsAsync(Document document, CancellationToken cancellationToken)
+        public override async Task<Document> RemoveUnnecessaryImportsAsync(
+            Document document, 
+            Func<SyntaxNode, bool> predicate,
+            CancellationToken cancellationToken)
         {
+            predicate = predicate ?? Functions<SyntaxNode>.True;
             using (Logger.LogBlock(FunctionId.Refactoring_RemoveUnnecessaryImports_CSharp, cancellationToken))
             {
-                var unnecessaryImports = await GetCommonUnnecessaryImportsOfAllContextAsync(document, cancellationToken).ConfigureAwait(false);
+                var unnecessaryImports = await GetCommonUnnecessaryImportsOfAllContextAsync(
+                    document, predicate, cancellationToken).ConfigureAwait(false);
                 if (unnecessaryImports == null || unnecessaryImports.Any(import => import.OverlapsHiddenPosition(cancellationToken)))
                 {
                     return document;
@@ -66,18 +71,16 @@ public async Task<Document> RemoveUnnecessaryImportsAsync(Document document, Can
                 var oldRoot = (CompilationUnitSyntax)root;
                 var newRoot = (CompilationUnitSyntax)new Rewriter(unnecessaryImports, cancellationToken).Visit(oldRoot);
 
-                if (cancellationToken.IsCancellationRequested)
-                {
-                    return null;
-                }
-
+                cancellationToken.ThrowIfCancellationRequested();
                 return document.WithSyntaxRoot(await FormatResultAsync(document, newRoot, cancellationToken).ConfigureAwait(false));
             }
         }
 
-        protected override IEnumerable<UsingDirectiveSyntax> GetUnusedUsings(SemanticModel model, SyntaxNode root, CancellationToken cancellationToken)
+        protected override ImmutableArray<UsingDirectiveSyntax> GetUnnecessaryImports(
+            SemanticModel model, SyntaxNode root,
+            Func<SyntaxNode, bool> predicate, CancellationToken cancellationToken)
         {
-            return GetUnnecessaryImports(model, root, cancellationToken) as IEnumerable<UsingDirectiveSyntax>;
+            return GetUnnecessaryImportsShared(model, root, predicate, cancellationToken);
         }
 
         private async Task<SyntaxNode> FormatResultAsync(Document document, CompilationUnitSyntax newRoot, CancellationToken cancellationToken)

src/Features/Core/Portable/CodeRefactorings/MoveType/AbstractMoveTypeService.MoveTypeEditor.cs

@@ -50,7 +50,9 @@ internal override async Task<ImmutableArray<CodeActionOperation>> GetOperationsA
                 var projectToBeUpdated = SemanticDocument.Document.Project;
                 var newDocumentId = DocumentId.CreateNewId(projectToBeUpdated.Id, FileName);
 
-                var solutionWithNewDocument = await AddNewDocumentWithSingleTypeDeclarationAndImportsAsync(newDocumentId).ConfigureAwait(false);
+                var documentWithMovedType = await AddNewDocumentWithSingleTypeDeclarationAndImportsAsync(newDocumentId).ConfigureAwait(false);
+
+                var solutionWithNewDocument = documentWithMovedType.Project.Solution;
 
                 // Get the original source document again, from the latest forked solution.
                 var sourceDocument = solutionWithNewDocument.GetDocument(SemanticDocument.Document.Id);
@@ -58,7 +60,7 @@ internal override async Task<ImmutableArray<CodeActionOperation>> GetOperationsA
                 // update source document to add partial modifiers to type chain
                 // and/or remove type declaration from original source document.
                 var solutionWithBothDocumentsUpdated = await RemoveTypeFromSourceDocumentAsync(
-                      sourceDocument).ConfigureAwait(false);
+                      sourceDocument, documentWithMovedType).ConfigureAwait(false);
 
                 return ImmutableArray.Create<CodeActionOperation>(new ApplyChangesOperation(solutionWithBothDocumentsUpdated));
             }
@@ -69,7 +71,7 @@ internal override async Task<ImmutableArray<CodeActionOperation>> GetOperationsA
             /// </summary>
             /// <param name="newDocumentId">id for the new document to be added</param>
             /// <returns>the new solution which contains a new document with the type being moved</returns>
-            private async Task<Solution> AddNewDocumentWithSingleTypeDeclarationAndImportsAsync(
+            private async Task<Document> AddNewDocumentWithSingleTypeDeclarationAndImportsAsync(
                 DocumentId newDocumentId)
             {
                 var document = SemanticDocument.Document;
@@ -101,20 +103,22 @@ internal override async Task<ImmutableArray<CodeActionOperation>> GetOperationsA
                 // update the text for the new document
                 solutionWithNewDocument = solutionWithNewDocument.WithDocumentSyntaxRoot(newDocumentId, modifiedRoot, PreservationMode.PreserveIdentity);
 
-                // get the updated document, perform clean up like remove unused usings.
+                // get the updated document, give it the minimal set of imports that the type
+                // inside it needs.
+                var service = document.GetLanguageService<IRemoveUnnecessaryImportsService>();
                 var newDocument = solutionWithNewDocument.GetDocument(newDocumentId);
-                newDocument = await CleanUpDocumentAsync(newDocument).ConfigureAwait(false);
+                newDocument = await service.RemoveUnnecessaryImportsAsync(newDocument, CancellationToken).ConfigureAwait(false);
 
-                return newDocument.Project.Solution;
+                return newDocument;
             }
 
             /// <summary>
             /// update the original document and remove the type that was moved.
             /// perform other fix ups as necessary.
             /// </summary>
-            /// <param name="sourceDocument">original document</param>
             /// <returns>an updated solution with the original document fixed up as appropriate.</returns>
-            private async Task<Solution> RemoveTypeFromSourceDocumentAsync(Document sourceDocument)
+            private async Task<Solution> RemoveTypeFromSourceDocumentAsync(
+                Document sourceDocument, Document documentWithMovedType)
             {
                 var documentEditor = await DocumentEditor.CreateAsync(sourceDocument, CancellationToken).ConfigureAwait(false);
 
@@ -126,7 +130,24 @@ private async Task<Solution> RemoveTypeFromSourceDocumentAsync(Document sourceDo
 
                 var updatedDocument = documentEditor.GetChangedDocument();
 
-                updatedDocument = await CleanUpDocumentAsync(updatedDocument).ConfigureAwait(false);
+                // Now, remove any imports that we no longer need *if* they were used in the new
+                // file with the moved type.  Essentially, those imports were here just to serve
+                // that new type and we shoudl remove them.  If we have *other* imports that that
+                // other file does not use *and* we do not use, we'll still keep those around.
+                // Those may be important to the user for code they're about to write, and we 
+                // don't want to interfere with them by removing them.
+                var service = sourceDocument.GetLanguageService<IRemoveUnnecessaryImportsService>();
+
+                var syntaxFacts = sourceDocument.GetLanguageService<ISyntaxFactsService>();
+
+                var rootWithMovedType = await documentWithMovedType.GetSyntaxRootAsync(CancellationToken).ConfigureAwait(false);
+                var movedImports = rootWithMovedType.DescendantNodes()
+                                                    .Where(syntaxFacts.IsUsingOrExternOrImport)
+                                                    .ToImmutableArray();
+
+                Func<SyntaxNode, bool> predicate = n => movedImports.Contains(i => i.IsEquivalentTo(n));
+                updatedDocument = await service.RemoveUnnecessaryImportsAsync(
+                    updatedDocument, predicate, CancellationToken).ConfigureAwait(false);
 
                 return updatedDocument.Project.Solution;
             }
@@ -206,16 +227,6 @@ private static bool FilterToTopLevelMembers(SyntaxNode node, SyntaxNode typeNode
                     }
                 }
             }
-
-            /// <summary>
-            /// Perform clean ups on a given document.
-            /// </summary>
-            private Task<Document> CleanUpDocumentAsync(Document document)
-            {
-                return document
-                    .GetLanguageService<IRemoveUnnecessaryImportsService>()
-                    .RemoveUnnecessaryImportsAsync(document, CancellationToken);
-            }
         }
     }
 }
\ No newline at end of file

src/Features/Core/Portable/RemoveUnnecessaryImports/AbstractRemoveUnnecessaryImportsService.cs

 // Copyright (c) Microsoft.  All Rights Reserved.  Licensed under the Apache License, Version 2.0.  See License.txt in the project root for license information.
 
+using System;
 using System.Collections.Generic;
-using System.Linq;
+using System.Collections.Immutable;
 using System.Threading;
 using System.Threading.Tasks;
-using Microsoft.CodeAnalysis.Host;
 using Microsoft.CodeAnalysis.Shared.Extensions;
 using Roslyn.Utilities;
 
 namespace Microsoft.CodeAnalysis.RemoveUnnecessaryImports
 {
-    internal abstract class AbstractRemoveUnnecessaryImportsService<T> : ILanguageService, IEqualityComparer<T> where T : SyntaxNode
+    internal abstract class AbstractRemoveUnnecessaryImportsService<T> : 
+        IRemoveUnnecessaryImportsService, IEqualityComparer<T> where T : SyntaxNode
     {
-        protected abstract IEnumerable<T> GetUnusedUsings(SemanticModel model, SyntaxNode root, CancellationToken cancellationToken);
+        public Task<Document> RemoveUnnecessaryImportsAsync(Document document, CancellationToken cancellationToken)
+            => RemoveUnnecessaryImportsAsync(document, predicate: null, cancellationToken: cancellationToken);
 
-        protected async Task<HashSet<T>> GetCommonUnnecessaryImportsOfAllContextAsync(Document document, CancellationToken cancellationToken)
+        public abstract Task<Document> RemoveUnnecessaryImportsAsync(Document fromDocument, Func<SyntaxNode, bool> predicate, CancellationToken cancellationToken);
+
+        protected abstract ImmutableArray<T> GetUnnecessaryImports(
+            SemanticModel model, SyntaxNode root, 
+            Func<SyntaxNode, bool> predicate, CancellationToken cancellationToken);
+
+        protected async Task<HashSet<T>> GetCommonUnnecessaryImportsOfAllContextAsync(
+            Document document, Func<SyntaxNode, bool> predicate, CancellationToken cancellationToken)
         {
             var model = await document.GetSemanticModelAsync(cancellationToken).ConfigureAwait(false);
             var root = await document.GetSyntaxRootAsync(cancellationToken).ConfigureAwait(false);
 
-            var unnecessaryImports = new HashSet<T>(GetUnnecessaryImportsOrEmpty(model, root, cancellationToken), this);
+            var unnecessaryImports = new HashSet<T>(this);
+            unnecessaryImports.AddRange(GetUnnecessaryImports(
+                model, root, predicate, cancellationToken));
             foreach (var current in document.GetLinkedDocuments())
             {
                 var currentModel = await current.GetSemanticModelAsync(cancellationToken).ConfigureAwait(false);
                 var currentRoot = await current.GetSyntaxRootAsync(cancellationToken).ConfigureAwait(false);
 
-                unnecessaryImports.IntersectWith(GetUnnecessaryImportsOrEmpty(currentModel, currentRoot, cancellationToken));
+                unnecessaryImports.IntersectWith(GetUnnecessaryImports(
+                    currentModel, currentRoot, predicate, cancellationToken));
             }
 
             return unnecessaryImports;
         }
 
-        private IEnumerable<T> GetUnnecessaryImportsOrEmpty(SemanticModel model, SyntaxNode root, CancellationToken cancellationToken)
-        {
-            var imports = GetUnusedUsings(model, root, cancellationToken) ?? SpecializedCollections.EmptyEnumerable<T>();
-            return imports.Cast<T>();
-        }
-
         bool IEqualityComparer<T>.Equals(T x, T y)
         {
             return x.Span == y.Span;

src/Features/Core/Portable/RemoveUnnecessaryImports/IRemoveUnnecessaryImportsService.cs

 // Copyright (c) Microsoft.  All Rights Reserved.  Licensed under the Apache License, Version 2.0.  See License.txt in the project root for license information.
 
+using System;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.Host;
@@ -8,8 +9,7 @@ namespace Microsoft.CodeAnalysis.RemoveUnnecessaryImports
 {
     internal interface IRemoveUnnecessaryImportsService : ILanguageService
     {
-        /// <returns>Returns the rewritten document, or the document passed in if no changes were made. If cancellation
-        /// was observed, it returns null.</returns>
         Task<Document> RemoveUnnecessaryImportsAsync(Document document, CancellationToken cancellationToken);
+        Task<Document> RemoveUnnecessaryImportsAsync(Document fromDocument, Func<SyntaxNode, bool> predicate, CancellationToken cancellationToken);
     }
 }
\ No newline at end of file

src/Features/VisualBasic/Portable/Diagnostics/Analyzers/VisualBasicRemoveUnnecessaryImportsDiagnosticAnalyzer.vb

@@ -20,13 +20,14 @@ Namespace Microsoft.CodeAnalysis.VisualBasic.Diagnostics.RemoveUnnecessaryImport
 
         Protected Overrides Function GetUnnecessaryImports(
                 semanticModel As SemanticModel, root As SyntaxNode, cancellationToken As CancellationToken) As IEnumerable(Of SyntaxNode)
-            Return VisualBasicRemoveUnnecessaryImportsService.GetUnnecessaryImports(semanticModel, root, cancellationToken)
+            Return VisualBasicRemoveUnnecessaryImportsService.GetUnnecessaryImportsShared(
+                semanticModel, root, predicate:=Nothing, cancellationToken:=cancellationToken)
         End Function
 
         Protected Overrides Function GetFixableDiagnosticSpans(
                 nodes As IEnumerable(Of SyntaxNode), tree As SyntaxTree, cancellationToken As CancellationToken) As IEnumerable(Of TextSpan)
             ' Create one fixable diagnostic that contains the entire Imports list.
-            Return SpecializedCollections.SingletonEnumerable(Of TextSpan)(tree.GetCompilationUnitRoot().Imports.GetContainedSpan())
+            Return SpecializedCollections.SingletonEnumerable(tree.GetCompilationUnitRoot().Imports.GetContainedSpan())
         End Function
 
         Protected Overrides Function GetLastTokenDelegateForContiguousSpans() As Func(Of SyntaxNode, SyntaxToken)

src/Features/VisualBasic/Portable/RemoveUnnecessaryImports/VisualBasicRemoveUnnecessaryImportsService.vb

 ' Copyright (c) Microsoft.  All Rights Reserved.  Licensed under the Apache License, Version 2.0.  See License.txt in the project root for license information.
 
+Imports System.Collections.Immutable
 Imports System.Composition
 Imports System.Threading
 Imports Microsoft.CodeAnalysis
@@ -13,12 +14,16 @@ Namespace Microsoft.CodeAnalysis.VisualBasic.RemoveUnnecessaryImports
     <ExportLanguageService(GetType(IRemoveUnnecessaryImportsService), LanguageNames.VisualBasic), [Shared]>
     Partial Friend Class VisualBasicRemoveUnnecessaryImportsService
         Inherits AbstractRemoveUnnecessaryImportsService(Of ImportsClauseSyntax)
-        Implements IRemoveUnnecessaryImportsService
 
-        Public Shared Function GetUnnecessaryImports(model As SemanticModel, root As SyntaxNode, cancellationToken As CancellationToken) As IEnumerable(Of SyntaxNode)
-            Dim unnecessaryImports = DirectCast(GetIndividualUnnecessaryImports(model, root, cancellationToken), ISet(Of ImportsClauseSyntax))
-            If unnecessaryImports Is Nothing Then
-                Return Nothing
+        Public Shared Function GetUnnecessaryImportsShared(
+                model As SemanticModel,
+                root As SyntaxNode,
+                predicate As Func(Of SyntaxNode, Boolean),
+                cancellationToken As CancellationToken) As ImmutableArray(Of SyntaxNode)
+            predicate = If(predicate, Functions(Of SyntaxNode).True)
+            Dim unnecessaryImports = GetIndividualUnnecessaryImportsShared(model, root, predicate, cancellationToken)
+            If Not unnecessaryImports.Any() Then
+                Return ImmutableArray(Of SyntaxNode).Empty
             End If
 
             Return unnecessaryImports.Select(
@@ -29,14 +34,20 @@ Namespace Microsoft.CodeAnalysis.VisualBasic.RemoveUnnecessaryImports
                         Else
                             Return i
                         End If
-                    End Function).ToSet()
+                    End Function).ToImmutableArray()
         End Function
 
-        Public Async Function RemoveUnnecessaryImportsAsync(document As Document, cancellationToken As CancellationToken) As Task(Of Document) Implements IRemoveUnnecessaryImportsService.RemoveUnnecessaryImportsAsync
+        Public Overrides Async Function RemoveUnnecessaryImportsAsync(
+                document As Document,
+                predicate As Func(Of SyntaxNode, Boolean),
+                cancellationToken As CancellationToken) As Task(Of Document)
+
+            predicate = If(predicate, Functions(Of SyntaxNode).True)
             Using Logger.LogBlock(FunctionId.Refactoring_RemoveUnnecessaryImports_VisualBasic, cancellationToken)
 
-                Dim unnecessaryImports = Await GetCommonUnnecessaryImportsOfAllContextAsync(document, cancellationToken).ConfigureAwait(False)
-                If unnecessaryImports Is Nothing OrElse unnecessaryImports.Any(Function(import) import.OverlapsHiddenPosition(cancellationToken)) Then
+                Dim unnecessaryImports = Await GetCommonUnnecessaryImportsOfAllContextAsync(
+                    document, predicate, cancellationToken).ConfigureAwait(False)
+                If unnecessaryImports.Any(Function(import) import.OverlapsHiddenPosition(cancellationToken)) Then
                     Return document
                 End If
 
@@ -46,56 +57,53 @@ Namespace Microsoft.CodeAnalysis.VisualBasic.RemoveUnnecessaryImports
                 Dim newRoot = New Rewriter(unnecessaryImports, cancellationToken).Visit(oldRoot)
                 newRoot = newRoot.WithAdditionalAnnotations(Formatter.Annotation)
 
-                If cancellationToken.IsCancellationRequested Then
-                    Return Nothing
-                End If
-
+                cancellationToken.ThrowIfCancellationRequested()
                 Return document.WithSyntaxRoot(newRoot)
             End Using
         End Function
 
-        Protected Overrides Function GetUnusedUsings(model As SemanticModel, root As SyntaxNode, cancellationToken As CancellationToken) As IEnumerable(Of ImportsClauseSyntax)
-            Return DirectCast(GetIndividualUnnecessaryImports(model, root, cancellationToken), IEnumerable(Of ImportsClauseSyntax))
+        Protected Overrides Function GetUnnecessaryImports(
+                model As SemanticModel, root As SyntaxNode,
+                predicate As Func(Of SyntaxNode, Boolean),
+                cancellationToken As CancellationToken) As ImmutableArray(Of ImportsClauseSyntax)
+            Return GetIndividualUnnecessaryImportsShared(model, root, predicate, cancellationToken)
         End Function
 
-        Private Shared Function GetIndividualUnnecessaryImports(semanticModel As SemanticModel, root As SyntaxNode, cancellationToken As CancellationToken) As IEnumerable(Of SyntaxNode)
-            Dim diagnostics = semanticModel.GetDiagnostics(cancellationToken:=cancellationToken)
+        Private Shared Function GetIndividualUnnecessaryImportsShared(
+                semanticModel As SemanticModel, root As SyntaxNode,
+                predicate As Func(Of SyntaxNode, Boolean),
+                CancellationToken As CancellationToken) As ImmutableArray(Of ImportsClauseSyntax)
+            Dim diagnostics = semanticModel.GetDiagnostics(cancellationToken:=CancellationToken)
 
             Dim unnecessaryImports = New HashSet(Of ImportsClauseSyntax)
 
             For Each diagnostic In diagnostics
                 If diagnostic.Id = "BC50000" Then
                     Dim node = root.FindNode(diagnostic.Location.SourceSpan)
-                    If node IsNot Nothing Then
+                    If node IsNot Nothing AndAlso predicate(node) Then
                         unnecessaryImports.Add(DirectCast(node, ImportsClauseSyntax))
                     End If
                 End If
 
                 If diagnostic.Id = "BC50001" Then
                     Dim node = TryCast(root.FindNode(diagnostic.Location.SourceSpan), ImportsStatementSyntax)
-                    If node IsNot Nothing Then
+                    If node IsNot Nothing AndAlso predicate(node) Then
                         unnecessaryImports.AddRange(node.ImportsClauses)
                     End If
                 End If
             Next
 
-            If cancellationToken.IsCancellationRequested Then
-                Return Nothing
-            End If
-
             Dim oldRoot = DirectCast(root, CompilationUnitSyntax)
-            AddRedundantImports(oldRoot, semanticModel, unnecessaryImports, cancellationToken)
-
-            If unnecessaryImports.Count = 0 Then
-                Return Nothing
-            End If
+            AddRedundantImports(oldRoot, semanticModel, unnecessaryImports, predicate, CancellationToken)
 
-            Return unnecessaryImports
+            Return unnecessaryImports.ToImmutableArray()
         End Function
 
-        Private Shared Sub AddRedundantImports(compilationUnit As CompilationUnitSyntax,
+        Private Shared Sub AddRedundantImports(
+                compilationUnit As CompilationUnitSyntax,
                 semanticModel As SemanticModel,
                 unnecessaryImports As HashSet(Of ImportsClauseSyntax),
+                predicate As Func(Of SyntaxNode, Boolean),
                 cancellationToken As CancellationToken)
             ' Now that we've visited the tree, add any imports that bound to project level
             ' imports.  We definitely can remove them.
@@ -106,18 +114,20 @@ Namespace Microsoft.CodeAnalysis.VisualBasic.RemoveUnnecessaryImports
                     Dim simpleImportsClause = TryCast(clause, SimpleImportsClauseSyntax)
                     If simpleImportsClause IsNot Nothing Then
                         If simpleImportsClause.Alias Is Nothing Then
-                            AddRedundantMemberImportsClause(simpleImportsClause, semanticModel, unnecessaryImports, cancellationToken)
+                            AddRedundantMemberImportsClause(simpleImportsClause, semanticModel, unnecessaryImports, predicate, cancellationToken)
                         Else
-                            AddRedundantAliasImportsClause(simpleImportsClause, semanticModel, unnecessaryImports, cancellationToken)
+                            AddRedundantAliasImportsClause(simpleImportsClause, semanticModel, unnecessaryImports, predicate, cancellationToken)
                         End If
                     End If
                 Next
             Next
         End Sub
 
-        Private Shared Sub AddRedundantAliasImportsClause(clause As SimpleImportsClauseSyntax,
+        Private Shared Sub AddRedundantAliasImportsClause(
+                clause As SimpleImportsClauseSyntax,
                 semanticModel As SemanticModel,
                 unnecessaryImports As HashSet(Of ImportsClauseSyntax),
+                predicate As Func(Of SyntaxNode, Boolean),
                 cancellationToken As CancellationToken)
 
             Dim semanticInfo = semanticModel.GetSymbolInfo(clause.Name, cancellationToken)
@@ -129,14 +139,18 @@ Namespace Microsoft.CodeAnalysis.VisualBasic.RemoveUnnecessaryImports
 
             Dim compilation = semanticModel.Compilation
             Dim aliasSymbol = compilation.AliasImports.FirstOrDefault(Function(a) a.Name = clause.Alias.Identifier.ValueText)
-            If aliasSymbol IsNot Nothing AndAlso aliasSymbol.Target.Equals(semanticInfo.Symbol) Then
+            If aliasSymbol IsNot Nothing AndAlso
+               aliasSymbol.Target.Equals(semanticInfo.Symbol) AndAlso
+               predicate(clause) Then
                 unnecessaryImports.Add(clause)
             End If
         End Sub
 
-        Private Shared Sub AddRedundantMemberImportsClause(clause As SimpleImportsClauseSyntax,
+        Private Shared Sub AddRedundantMemberImportsClause(
+                clause As SimpleImportsClauseSyntax,
                 semanticModel As SemanticModel,
                 unnecessaryImports As HashSet(Of ImportsClauseSyntax),
+                predicate As Func(Of SyntaxNode, Boolean),
                 cancellationToken As CancellationToken)
 
             Dim semanticInfo = semanticModel.GetSymbolInfo(clause.Name, cancellationToken)
@@ -147,7 +161,8 @@ Namespace Microsoft.CodeAnalysis.VisualBasic.RemoveUnnecessaryImports
             End If
 
             Dim compilation = semanticModel.Compilation
-            If compilation.MemberImports.Contains(namespaceOrType) Then
+            If compilation.MemberImports.Contains(namespaceOrType) AndAlso
+               predicate(clause) Then
                 unnecessaryImports.Add(clause)
             End If
         End Sub

src/VisualStudio/Xaml/Impl/Features/OrganizeImports/XamlRemoveUnnecessaryImportsService.cs

 // Copyright (c) Microsoft.  All Rights Reserved.  Licensed under the Apache License, Version 2.0.  See License.txt in the project root for license information.
 
+using System;
 using System.Composition;
 using System.Threading;
 using System.Threading.Tasks;
@@ -21,6 +22,10 @@ public XamlRemoveUnnecessaryImportsService(IXamlRemoveUnnecessaryNamespacesServi
         }
 
         public Task<Document> RemoveUnnecessaryImportsAsync(Document document, CancellationToken cancellationToken)
+            => RemoveUnnecessaryImportsAsync(document, predicate: null, cancellationToken: cancellationToken);
+
+        public Task<Document> RemoveUnnecessaryImportsAsync(
+            Document document, Func<SyntaxNode, bool> predicate, CancellationToken cancellationToken)
         {
             return _removeService.RemoveUnnecessaryNamespacesAsync(document, cancellationToken) ?? Task.FromResult(document);
         }