|
9 | 9 | using Rubberduck.Refactorings.EncapsulateField;
|
10 | 10 | using System.Collections.Generic;
|
11 | 11 | using Rubberduck.SmartIndenter;
|
| 12 | +using Antlr4.Runtime; |
| 13 | +using Rubberduck.Parsing.Grammar; |
| 14 | +using Rubberduck.Parsing; |
12 | 15 |
|
13 | 16 | namespace Rubberduck.Refactorings.EncapsulateFieldInsertNewCode
|
14 | 17 | {
|
@@ -106,17 +109,33 @@ private void InsertBlocks(EncapsulateFieldInsertNewCodeModel model, IRewriteSess
|
106 | 109 |
|
107 | 110 | var rewriter = rewriteSession.CheckOutModuleRewriter(model.QualifiedModuleName);
|
108 | 111 |
|
109 |
| - var codeSectionStartIndex = _declarationFinderProvider.DeclarationFinder |
110 |
| - .Members(model.QualifiedModuleName).Where(m => m.IsMember()) |
| 112 | + var firstMember = _declarationFinderProvider.DeclarationFinder |
| 113 | + .Members(model.QualifiedModuleName) |
| 114 | + .Where(m => m.IsMember()) |
111 | 115 | .OrderBy(c => c.Selection)
|
112 |
| - .FirstOrDefault()?.Context.Start.TokenIndex; |
| 116 | + .FirstOrDefault(); |
113 | 117 |
|
114 |
| - if (codeSectionStartIndex.HasValue) |
| 118 | + var codeSectionStartIndex = firstMember?.Context.Start.TokenIndex; |
| 119 | + |
| 120 | + var firstAnnotationContext = firstMember?.Annotations? |
| 121 | + .OrderBy(c => c.QualifiedSelection) |
| 122 | + .FirstOrDefault()?.Context; |
| 123 | + |
| 124 | + if (firstAnnotationContext != null) |
| 125 | + { |
| 126 | + codeSectionStartIndex = firstAnnotationContext |
| 127 | + .GetAncestor<VBAParser.CommentOrAnnotationContext>().start.TokenIndex; |
| 128 | + } |
| 129 | + |
| 130 | + if (codeSectionStartIndex != null) |
115 | 131 | {
|
116 | 132 | rewriter.InsertBefore(codeSectionStartIndex.Value, $"{formattedContent}{NewLines.DOUBLE_SPACE}");
|
117 |
| - return; |
118 | 133 | }
|
119 |
| - rewriter.InsertBefore(rewriter.TokenStream.Size - 1, $"{NewLines.DOUBLE_SPACE}{formattedContent}"); |
| 134 | + else |
| 135 | + { |
| 136 | + var insertIndex = rewriter.TokenStream.Size - 1; |
| 137 | + rewriter.InsertBefore(insertIndex, $"{NewLines.DOUBLE_SPACE}{formattedContent}"); |
| 138 | + } |
120 | 139 | }
|
121 | 140 | }
|
122 | 141 | }
|
0 commit comments