Merge pull request #15140 from CyrusNajmabadi/addUsingSnippet

Insert a Using into the right locatoin when a snippet adds an import.

Fixes #4457

src/Test/Utilities/Shared/MarkedSource/MarkupTestFile.cs

@@ -230,11 +230,11 @@ public static void GetPositionAndSpans(string input, out string output, out int
             cursorPosition = pos.Value;
         }
 
+        public static void GetPosition(string input, out string output, out int? cursorPosition)
+            => GetPositionAndSpans(input, out output, out cursorPosition, out IList<TextSpan> spans);
+
         public static void GetPosition(string input, out string output, out int cursorPosition)
-        {
-            IList<TextSpan> spans;
-            GetPositionAndSpans(input, out output, out cursorPosition, out spans);
-        }
+            => GetPositionAndSpans(input, out output, out cursorPosition, out var spans);
 
         public static void GetPositionAndSpan(string input, out string output, out int? cursorPosition, out TextSpan? textSpan)
         {

src/VisualStudio/CSharp/Impl/Snippets/SnippetExpansionClient.cs

@@ -11,6 +11,7 @@
 using Microsoft.CodeAnalysis.CSharp.Syntax;
 using Microsoft.CodeAnalysis.Editor.Shared.Extensions;
 using Microsoft.CodeAnalysis.Formatting;
+using Microsoft.CodeAnalysis.Shared.Extensions;
 using Microsoft.VisualStudio.Editor;
 using Microsoft.VisualStudio.LanguageServices.CSharp.Snippets.SnippetFunctions;
 using Microsoft.VisualStudio.LanguageServices.Implementation.Snippets;
@@ -81,7 +82,9 @@ public override int GetExpansionFunction(IXMLDOMNode xmlFunctionNode, string bst
             }
         }
 
-        internal override Document AddImports(Document document, XElement snippetNode, bool placeSystemNamespaceFirst, CancellationToken cancellationToken)
+        internal override Document AddImports(
+            Document document, int position, XElement snippetNode,
+            bool placeSystemNamespaceFirst, CancellationToken cancellationToken)
         {
             var importsNode = snippetNode.Element(XName.Get("Imports", snippetNode.Name.NamespaceName));
             if (importsNode == null ||
@@ -104,8 +107,14 @@ internal override Document AddImports(Document document, XElement snippetNode, b
             }
 
             var root = document.GetSyntaxRootSynchronously(cancellationToken);
+            var node = root.FindToken(position).Parent;
+            var container = node.GetInnermostNamespaceDeclarationWithUsings() ?? (SyntaxNode)node.GetAncestorOrThis<CompilationUnitSyntax>();
 
-            var newRoot = ((CompilationUnitSyntax)root).AddUsingDirectives(newUsingDirectives, placeSystemNamespaceFirst);
+            var newContainer = container is NamespaceDeclarationSyntax n
+                ? (SyntaxNode)n.AddUsingDirectives(newUsingDirectives, placeSystemNamespaceFirst, Formatter.Annotation)
+                : ((CompilationUnitSyntax)container).AddUsingDirectives(newUsingDirectives, placeSystemNamespaceFirst);
+
+            var newRoot = root.ReplaceNode(container, newContainer);
             var newDocument = document.WithSyntaxRoot(newRoot);
 
             var formattedDocument = Formatter.FormatAsync(newDocument, Formatter.Annotation, cancellationToken: cancellationToken).WaitAndGetResult(cancellationToken);

src/VisualStudio/Core/Def/Implementation/Snippets/AbstractSnippetExpansionClient.cs

@@ -56,7 +56,7 @@ public AbstractSnippetExpansionClient(Guid languageServiceGuid, ITextView textVi
 
         public abstract int GetExpansionFunction(IXMLDOMNode xmlFunctionNode, string bstrFieldName, out IVsExpansionFunction pFunc);
         protected abstract ITrackingSpan InsertEmptyCommentAndGetEndPositionTrackingSpan();
-        internal abstract Document AddImports(Document document, XElement snippetNode, bool placeSystemNamespaceFirst, CancellationToken cancellationToken);
+        internal abstract Document AddImports(Document document, int position, XElement snippetNode, bool placeSystemNamespaceFirst, CancellationToken cancellationToken);
 
         public int FormatSpan(IVsTextLines pBuffer, VsTextSpan[] tsInSurfaceBuffer)
         {
@@ -107,6 +107,7 @@ public int FormatSpan(IVsTextLines pBuffer, VsTextSpan[] tsInSurfaceBuffer)
             var endPositionTrackingSpan = isFullSnippetFormat ? InsertEmptyCommentAndGetEndPositionTrackingSpan() : null;
 
             var formattingSpan = CommonFormattingHelpers.GetFormattingSpan(SubjectBuffer.CurrentSnapshot, snippetTrackingSpan.GetSpan(SubjectBuffer.CurrentSnapshot));
+            
             SubjectBuffer.CurrentSnapshot.FormatAndApplyToBuffer(formattingSpan, CancellationToken.None);
 
             if (isFullSnippetFormat)
@@ -117,7 +118,9 @@ public int FormatSpan(IVsTextLines pBuffer, VsTextSpan[] tsInSurfaceBuffer)
                 // specified in the snippet xml. In OnBeforeInsertion we have no guarantee that the
                 // snippet xml will be available, and changing the buffer during OnAfterInsertion can
                 // cause the underlying tracking spans to get out of sync.
-                AddReferencesAndImports(ExpansionSession, cancellationToken);
+                var currentStartPosition = snippetTrackingSpan.GetStartPoint(SubjectBuffer.CurrentSnapshot).Position;
+                AddReferencesAndImports(
+                    ExpansionSession, currentStartPosition, cancellationToken);
 
                 SetNewEndPosition(endPositionTrackingSpan);
             }
@@ -512,7 +515,10 @@ private void GetCaretPositionInSurfaceBuffer(out int caretLine, out int caretCol
             }
         }
 
-        private void AddReferencesAndImports(IVsExpansionSession pSession, CancellationToken cancellationToken)
+        private void AddReferencesAndImports(
+            IVsExpansionSession pSession,
+            int position,
+            CancellationToken cancellationToken)
         {
             XElement snippetNode;
             if (!TryGetSnippetNode(pSession, out snippetNode))
@@ -528,7 +534,8 @@ private void AddReferencesAndImports(IVsExpansionSession pSession, CancellationT
 
             var documentOptions = documentWithImports.GetOptionsAsync(cancellationToken).WaitAndGetResult(cancellationToken);
             var placeSystemNamespaceFirst = documentOptions.GetOption(GenerationOptions.PlaceSystemNamespaceFirst);
-            documentWithImports = AddImports(documentWithImports, snippetNode, placeSystemNamespaceFirst, cancellationToken);
+
+            documentWithImports = AddImports(documentWithImports, position, snippetNode, placeSystemNamespaceFirst, cancellationToken);
             AddReferences(documentWithImports.Project, snippetNode);
         }
 
@@ -657,4 +664,4 @@ internal bool TryGetSpanOnHigherBuffer(SnapshotSpan snapshotSpan, ITextBuffer ta
             return false;
         }
     }
-}
+}
\ No newline at end of file

src/VisualStudio/Core/Test/Snippets/CSharpSnippetExpansionClientTests.vb

@@ -6,8 +6,6 @@ Imports Microsoft.CodeAnalysis
 Imports Microsoft.CodeAnalysis.Editor.UnitTests.Extensions
 Imports Microsoft.CodeAnalysis.Editor.UnitTests.Workspaces
 Imports Microsoft.CodeAnalysis.Formatting
-Imports Microsoft.CodeAnalysis.Options
-Imports Microsoft.VisualStudio.LanguageServices
 Imports Microsoft.VisualStudio.LanguageServices.CSharp.Snippets
 Imports Microsoft.VisualStudio.Text.Projection
 Imports Roslyn.Test.Utilities
@@ -17,125 +15,159 @@ Namespace Microsoft.VisualStudio.LanguageServices.UnitTests.Snippets
 
         <WpfFact, Trait(Traits.Feature, Traits.Features.Snippets)>
         Public Async Function TestAddImport_EmptyDocument() As Task
-            Dim originalCode = <![CDATA[]]>.Value
+            Dim originalCode = ""
             Dim namespacesToAdd = {"System"}
-            Dim expectedUpdatedCode = <![CDATA[using System;
-]]>.Value
+            Dim expectedUpdatedCode = "using System;
+"
+
             Await TestSnippetAddImportsAsync(originalCode, namespacesToAdd, placeSystemNamespaceFirst:=True, expectedUpdatedCode:=expectedUpdatedCode)
         End Function
 
         <WpfFact, Trait(Traits.Feature, Traits.Features.Snippets)>
         Public Async Function TestAddImport_EmptyDocument_SystemAtTop() As Task
-            Dim originalCode = <![CDATA[]]>.Value
+            Dim originalCode = ""
             Dim namespacesToAdd = {"First.Alphabetically", "System.Bar"}
-            Dim expectedUpdatedCode = <![CDATA[using System.Bar;
+            Dim expectedUpdatedCode = "using System.Bar;
 using First.Alphabetically;
-]]>.Value
+"
             Await TestSnippetAddImportsAsync(originalCode, namespacesToAdd, placeSystemNamespaceFirst:=True, expectedUpdatedCode:=expectedUpdatedCode)
         End Function
 
         <WpfFact, Trait(Traits.Feature, Traits.Features.Snippets)>
         Public Async Function TestAddImport_EmptyDocument_SystemNotSortedToTop() As Task
-            Dim originalCode = <![CDATA[]]>.Value
+            Dim originalCode = ""
             Dim namespacesToAdd = {"First.Alphabetically", "System.Bar"}
-            Dim expectedUpdatedCode = <![CDATA[using First.Alphabetically;
+            Dim expectedUpdatedCode = "using First.Alphabetically;
 using System.Bar;
-]]>.Value
+"
+
             Await TestSnippetAddImportsAsync(originalCode, namespacesToAdd, placeSystemNamespaceFirst:=False, expectedUpdatedCode:=expectedUpdatedCode)
         End Function
 
         <WpfFact, Trait(Traits.Feature, Traits.Features.Snippets)>
         Public Async Function TestAddImport_AddsOnlyNewNamespaces() As Task
-            Dim originalCode = <![CDATA[using A.B.C;
+            Dim originalCode = "using A.B.C;
 using D.E.F;
-]]>.Value
+"
             Dim namespacesToAdd = {"D.E.F", "G.H.I"}
-            Dim expectedUpdatedCode = <![CDATA[using A.B.C;
+            Dim expectedUpdatedCode = "using A.B.C;
 using D.E.F;
 using G.H.I;
-]]>.Value
+"
+            Await TestSnippetAddImportsAsync(originalCode, namespacesToAdd, placeSystemNamespaceFirst:=True, expectedUpdatedCode:=expectedUpdatedCode)
+        End Function
+
+        <WorkItem(4457, "https://github.com/dotnet/roslyn/issues/4457")>
+        <WpfFact, Trait(Traits.Feature, Traits.Features.Snippets)>
+        Public Async Function TestAddImport_InsideNamespace() As Task
+            Dim originalCode = "
+using A;
+
+namespace N
+{
+    using B;
+
+    class C
+    {
+        $$
+    }
+}"
+            Dim namespacesToAdd = {"D"}
+            Dim expectedUpdatedCode = "
+using A;
+
+namespace N
+{
+    using B;
+    using D;
+
+    class C
+    {
+        
+    }
+}"
             Await TestSnippetAddImportsAsync(originalCode, namespacesToAdd, placeSystemNamespaceFirst:=True, expectedUpdatedCode:=expectedUpdatedCode)
         End Function
 
         <WpfFact, Trait(Traits.Feature, Traits.Features.Snippets)>
         Public Async Function TestAddImport_AddsOnlyNewAliasAndNamespacePairs() As Task
-            Dim originalCode = <![CDATA[using A = B.C;
+            Dim originalCode = "using A = B.C;
 using D = E.F;
 using G = H.I;
-]]>.Value
+"
             Dim namespacesToAdd = {"A = E.F", "D = B.C", "G = H.I", "J = K.L"}
-            Dim expectedUpdatedCode = <![CDATA[using A = B.C;
+            Dim expectedUpdatedCode = "using A = B.C;
 using A = E.F;
 using D = B.C;
 using D = E.F;
 using G = H.I;
 using J = K.L;
-]]>.Value
+"
             Await TestSnippetAddImportsAsync(originalCode, namespacesToAdd, placeSystemNamespaceFirst:=True, expectedUpdatedCode:=expectedUpdatedCode)
         End Function
 
         <WpfFact, Trait(Traits.Feature, Traits.Features.Snippets)>
         Public Async Function TestAddImport_DuplicateNamespaceDetectionDoesNotIgnoreCase() As Task
-            Dim originalCode = <![CDATA[using A.b.C;
-]]>.Value
+            Dim originalCode = "using A.b.C;
+"
             Dim namespacesToAdd = {"a.B.C", "A.B.c"}
-            Dim expectedUpdatedCode = <![CDATA[using a.B.C;
+            Dim expectedUpdatedCode = "using a.B.C;
 using A.b.C;
 using A.B.c;
-]]>.Value
+"
             Await TestSnippetAddImportsAsync(originalCode, namespacesToAdd, placeSystemNamespaceFirst:=True, expectedUpdatedCode:=expectedUpdatedCode)
         End Function
 
         <WpfFact, Trait(Traits.Feature, Traits.Features.Snippets)>
         Public Async Function TestAddImport_DuplicateAliasNamespacePairDetectionIgnoresWhitespace1() As Task
-            Dim originalCode = <![CDATA[using A = B.C;
-]]>.Value
+            Dim originalCode = "using A = B.C;
+"
             Dim namespacesToAdd = {"A  =        B.C"}
-            Dim expectedUpdatedCode = <![CDATA[using A = B.C;
-]]>.Value
+            Dim expectedUpdatedCode = "using A = B.C;
+"
             Await TestSnippetAddImportsAsync(originalCode, namespacesToAdd, placeSystemNamespaceFirst:=True, expectedUpdatedCode:=expectedUpdatedCode)
         End Function
 
         <WpfFact, Trait(Traits.Feature, Traits.Features.Snippets)>
         Public Async Function TestAddImport_DuplicateAliasNamespacePairDetectionIgnoresWhitespace2() As Task
-            Dim originalCode = <![CDATA[using A     =  B.C;
-]]>.Value
+            Dim originalCode = "using A     =  B.C;
+"
             Dim namespacesToAdd = {"A=B.C"}
-            Dim expectedUpdatedCode = <![CDATA[using A     =  B.C;
-]]>.Value
+            Dim expectedUpdatedCode = "using A     =  B.C;
+"
             Await TestSnippetAddImportsAsync(originalCode, namespacesToAdd, placeSystemNamespaceFirst:=True, expectedUpdatedCode:=expectedUpdatedCode)
         End Function
 
         <WpfFact, Trait(Traits.Feature, Traits.Features.Snippets)>
         Public Async Function TestAddImport_DuplicateAliasNamespacePairDetectionDoesNotIgnoreCase() As Task
-            Dim originalCode = <![CDATA[using A = B.C;
-]]>.Value
+            Dim originalCode = "using A = B.C;
+"
             Dim namespacesToAdd = {"a = b.C"}
-            Dim expectedUpdatedCode = <![CDATA[using a = b.C;
+            Dim expectedUpdatedCode = "using a = b.C;
 using A = B.C;
-]]>.Value
+"
             Await TestSnippetAddImportsAsync(originalCode, namespacesToAdd, placeSystemNamespaceFirst:=True, expectedUpdatedCode:=expectedUpdatedCode)
         End Function
 
         <WpfFact, Trait(Traits.Feature, Traits.Features.Snippets)>
         Public Async Function TestAddImport_OnlyFormatNewImports() As Task
-            Dim originalCode = <![CDATA[using A     =  B.C;
+            Dim originalCode = "using A     =  B.C;
 using G=   H.I;
-]]>.Value
+"
             Dim namespacesToAdd = {"D        =E.F"}
-            Dim expectedUpdatedCode = <![CDATA[using A     =  B.C;
+            Dim expectedUpdatedCode = "using A     =  B.C;
 using D = E.F;
 using G=   H.I;
-]]>.Value
+"
             Await TestSnippetAddImportsAsync(originalCode, namespacesToAdd, placeSystemNamespaceFirst:=True, expectedUpdatedCode:=expectedUpdatedCode)
         End Function
 
         <WpfFact, Trait(Traits.Feature, Traits.Features.Snippets)>
         Public Async Function TestAddImport_BadNamespaceGetsAdded() As Task
-            Dim originalCode = <![CDATA[]]>.Value
+            Dim originalCode = ""
             Dim namespacesToAdd = {"$system"}
-            Dim expectedUpdatedCode = <![CDATA[using $system;
-]]>.Value
+            Dim expectedUpdatedCode = "using $system;
+"
             Await TestSnippetAddImportsAsync(originalCode, namespacesToAdd, placeSystemNamespaceFirst:=True, expectedUpdatedCode:=expectedUpdatedCode)
         End Function
 
@@ -311,7 +343,16 @@ using G=   H.I;
             End Using
         End Function
 
-        Private Async Function TestSnippetAddImportsAsync(originalCode As String, namespacesToAdd As String(), placeSystemNamespaceFirst As Boolean, expectedUpdatedCode As String) As Tasks.Task
+        Private Async Function TestSnippetAddImportsAsync(
+                markupCode As String,
+                namespacesToAdd As String(),
+                placeSystemNamespaceFirst As Boolean,
+                expectedUpdatedCode As String) As Tasks.Task
+
+            Dim originalCode As String = Nothing
+            Dim position As Integer?
+            MarkupTestFile.GetPosition(markupCode, originalCode, position)
+
             Dim workspaceXml = <Workspace>
                                    <Project Language=<%= LanguageNames.CSharp %> CommonReferences="true">
                                        <Document><%= originalCode %></Document>
@@ -329,7 +370,7 @@ using G=   H.I;
                                                    </Import>)
             Next
 
-            Using workspace = Await TestWorkspace.CreateAsync(workspaceXml)
+            Using workspace = Await TestWorkspace.CreateCSharpAsync(originalCode)
                 Dim expansionClient = New SnippetExpansionClient(
                     Guids.VisualBasicDebuggerLanguageId,
                     workspace.Documents.Single().GetTextView(),
@@ -338,13 +379,13 @@ using G=   H.I;
 
                 Dim updatedDocument = expansionClient.AddImports(
                     workspace.CurrentSolution.Projects.Single().Documents.Single(),
+                    If(position, 0),
                     snippetNode,
                     placeSystemNamespaceFirst, CancellationToken.None)
 
-                Assert.Equal(expectedUpdatedCode.Replace(vbLf, vbCrLf),
+                Assert.Equal(expectedUpdatedCode,
                              (Await updatedDocument.GetTextAsync()).ToString())
             End Using
         End Function
-
     End Class
-End Namespace
+End Namespace
\ No newline at end of file

src/VisualStudio/Core/Test/Snippets/SnippetTestState.vb

@@ -215,7 +215,7 @@ Namespace Microsoft.VisualStudio.LanguageServices.UnitTests.Snippets
                 Throw New NotImplementedException()
             End Function
 
-            Friend Overrides Function AddImports(document As Document, snippetNode As XElement, placeSystemNamespaceFirst As Boolean, cancellationToken As CancellationToken) As Document
+            Friend Overrides Function AddImports(document As Document, position As Integer, snippetNode As XElement, placeSystemNamespaceFirst As Boolean, cancellationToken As CancellationToken) As Document
                 Return document
             End Function
         End Class

src/VisualStudio/Core/Test/Snippets/VisualBasicSnippetExpansionClientTests.vb

@@ -370,7 +370,16 @@ End Class</Test>
             End Using
         End Function
 
-        Private Async Function TestSnippetAddImportsAsync(originalCode As String, namespacesToAdd As String(), placeSystemNamespaceFirst As Boolean, expectedUpdatedCode As String) As Tasks.Task
+        Private Async Function TestSnippetAddImportsAsync(
+                markupCode As String,
+                namespacesToAdd As String(),
+                placeSystemNamespaceFirst As Boolean,
+                expectedUpdatedCode As String) As Tasks.Task
+
+            Dim originalCode As String = Nothing
+            Dim position As Integer?
+            MarkupTestFile.GetPosition(markupCode, originalCode, position)
+
             Dim workspaceXml = <Workspace>
                                    <Project Language=<%= LanguageNames.VisualBasic %> CommonReferences="true">
                                        <CompilationOptions/>
@@ -398,6 +407,7 @@ End Class</Test>
 
                 Dim updatedDocument = expansionClient.AddImports(
                     workspace.CurrentSolution.Projects.Single().Documents.Single(),
+                    If(position, 0),
                     snippetNode,
                     placeSystemNamespaceFirst, CancellationToken.None)
 

src/VisualStudio/VisualBasic/Impl/Snippets/SnippetExpansionClient.vb

@@ -87,7 +87,11 @@ Namespace Microsoft.VisualStudio.LanguageServices.VisualBasic.Snippets
             End Select
         End Function
 
-        Friend Overrides Function AddImports(document As Document, snippetNode As XElement, placeSystemNamespaceFirst As Boolean, cancellationToken As CancellationToken) As Document
+        Friend Overrides Function AddImports(document As Document,
+                                             position As Integer,
+                                             snippetNode As XElement,
+                                             placeSystemNamespaceFirst As Boolean,
+                                             cancellationToken As CancellationToken) As Document
             Dim importsNode = snippetNode.Element(XName.Get("Imports", snippetNode.Name.NamespaceName))
             If importsNode Is Nothing OrElse
                Not importsNode.HasElements() Then