1717using Microsoft . CodeAnalysis . Formatting ;
1818using Microsoft . CodeAnalysis . PooledObjects ;
1919using Microsoft . CodeAnalysis . RemoveUnnecessaryImports ;
20+ using Microsoft . CodeAnalysis . Shared . Collections ;
2021using Microsoft . CodeAnalysis . Shared . Extensions ;
2122using Roslyn . Utilities ;
2223
@@ -33,7 +34,7 @@ public static async Task<Document> ConvertToTopLevelStatementsAsync(
3334 Contract . ThrowIfNull ( typeDeclaration ) ; // checked by analyzer
3435
3536 var generator = document . GetRequiredLanguageService < SyntaxGenerator > ( ) ;
36- var root = await document . GetRequiredSyntaxRootAsync ( cancellationToken ) . ConfigureAwait ( false ) ;
37+ var root = ( CompilationUnitSyntax ) await document . GetRequiredSyntaxRootAsync ( cancellationToken ) . ConfigureAwait ( false ) ;
3738
3839 var semanticModel = await document . GetRequiredSemanticModelAsync ( cancellationToken ) . ConfigureAwait ( false ) ;
3940 var rootWithGlobalStatements = GetRootWithGlobalStatements (
@@ -97,7 +98,7 @@ private static void AddUsingDirectives(NameSyntax name, SyntaxAnnotation annotat
9798 private static SyntaxNode GetRootWithGlobalStatements (
9899 SemanticModel semanticModel ,
99100 SyntaxGenerator generator ,
100- SyntaxNode root ,
101+ CompilationUnitSyntax root ,
101102 TypeDeclarationSyntax typeDeclaration ,
102103 MethodDeclarationSyntax methodDeclaration ,
103104 CancellationToken cancellationToken )
@@ -107,19 +108,17 @@ private static SyntaxNode GetRootWithGlobalStatements(
107108 semanticModel , typeDeclaration , methodDeclaration , cancellationToken ) ;
108109
109110 var namespaceDeclaration = typeDeclaration . Parent as BaseNamespaceDeclarationSyntax ;
111+
110112 if ( namespaceDeclaration != null &&
111113 namespaceDeclaration . Members . Count >= 2 )
112114 {
113115 // Our parent namespace has another symbol in it. Keep the namespace declaration around, removing only
114116 // the existing Program type from it.
115117 editor . RemoveNode ( typeDeclaration ) ;
116- editor . ReplaceNode (
117- root ,
118- ( current , _ ) =>
119- {
120- var currentRoot = ( CompilationUnitSyntax ) current ;
121- return currentRoot . WithMembers ( currentRoot . Members . InsertRange ( 0 , globalStatements ) ) ;
122- } ) ;
118+ editor . InsertBefore ( namespaceDeclaration , globalStatements ) ;
119+
120+ // We want to place the trailing directive on the namespace declaration we're preceding.
121+ AddDirectivesToNextMemberOrEndOfFile ( root . Members . IndexOf ( namespaceDeclaration ) ) ;
123122 }
124123 else if ( namespaceDeclaration != null )
125124 {
@@ -136,18 +135,60 @@ private static SyntaxNode GetRootWithGlobalStatements(
136135 globalStatements [ 0 ] . WithPrependedLeadingTrivia ( fileBanner ) ) ;
137136 }
138137
139- editor . ReplaceNode (
140- root ,
141- root . ReplaceNode ( namespaceDeclaration , globalStatements ) ) ;
138+ editor . ReplaceNode ( namespaceDeclaration , ( _ , _ ) => globalStatements ) ;
139+
140+ // We're removing the namespace itself. So we want to place the trailing directive on the element that follows that.
141+ AddDirectivesToNextMemberOrEndOfFile ( root . Members . IndexOf ( namespaceDeclaration ) + 1 ) ;
142142 }
143143 else
144144 {
145145 // type wasn't in a namespace. just remove the type and replace it with the new global statements.
146- editor . ReplaceNode (
147- root , root . ReplaceNode ( typeDeclaration , globalStatements ) ) ;
146+ editor . ReplaceNode ( typeDeclaration , ( _ , _ ) => globalStatements ) ;
147+
148+ // We're removing the namespace itself. So we want to place the trailing directive on the element that follows that.
149+ AddDirectivesToNextMemberOrEndOfFile ( root . Members . IndexOf ( typeDeclaration ) + 1 ) ;
148150 }
149151
150152 return editor . GetChangedRoot ( ) ;
153+
154+ void AddDirectivesToNextMemberOrEndOfFile ( int memberIndexToPlaceTrailingDirectivesOn )
155+ {
156+ // If the method has trailing directive on the close brace, move them to whatever will come after the
157+ // final global statement in the new file. That could be the next namespace/type member declaration. Or
158+ // it could be the end of file token if there are no more members in the file.
159+ if ( methodDeclaration . Body is not BlockSyntax block )
160+ return ;
161+
162+ var leadingCloseBraceTrivia = block . CloseBraceToken . LeadingTrivia ;
163+ if ( ! leadingCloseBraceTrivia . Any ( t => t . IsDirective ) )
164+ return ;
165+
166+ if ( memberIndexToPlaceTrailingDirectivesOn < root . Members . Count )
167+ {
168+ editor . ReplaceNode (
169+ root . Members [ memberIndexToPlaceTrailingDirectivesOn ] ,
170+ ( current , _ ) =>
171+ {
172+ var updated = current . WithPrependedLeadingTrivia ( leadingCloseBraceTrivia ) ;
173+ updated = updated . ReplaceToken (
174+ updated . GetFirstToken ( ) ,
175+ updated . GetFirstToken ( ) . WithAdditionalAnnotations ( Formatter . Annotation ) ) ;
176+ return updated ;
177+ } ) ;
178+ }
179+ else
180+ {
181+ editor . ReplaceNode (
182+ root ,
183+ ( current , _ ) =>
184+ {
185+ var currentRoot = ( CompilationUnitSyntax ) current ;
186+ return currentRoot . WithEndOfFileToken ( currentRoot . EndOfFileToken
187+ . WithPrependedLeadingTrivia ( leadingCloseBraceTrivia )
188+ . WithAdditionalAnnotations ( Formatter . Annotation ) ) ;
189+ } ) ;
190+ }
191+ }
151192 }
152193
153194 private static ImmutableArray < GlobalStatementSyntax > GetGlobalStatements (
0 commit comments