Fix TreeBuilder#buildSnippet formatting

Netflix-Skunkworks · Feb 17, 2020 · d985bba · d985bba
1 parent fe09b4a
commit d985bba
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 4 deletions.
diff --git a/rewrite-core/src/main/java/com/netflix/rewrite/tree/TreeBuilder.java b/rewrite-core/src/main/java/com/netflix/rewrite/tree/TreeBuilder.java
@@ -36,6 +36,7 @@
 import static com.netflix.rewrite.internal.StringUtils.trimIndent;
 import static com.netflix.rewrite.tree.Formatting.format;
 import static com.netflix.rewrite.tree.Tr.randomId;
+import static com.netflix.rewrite.tree.visitor.refactor.Formatter.enclosingIndent;
 import static java.util.Arrays.stream;
 import static java.util.Collections.emptyList;
 import static java.util.Collections.singletonList;
@@ -121,7 +122,11 @@ public static <T extends Tree> List<T> buildSnippet(Tr.CompilationUnit containin
         Formatter formatter = new Formatter(cu);
 
         Tr.CompilationUnit formattedCu = (Tr.CompilationUnit) new TransformVisitor(block.getStatements().stream()
-                .flatMap(stat -> formatter.shiftRight(stat, insertionScope.getTree(), insertionScope.getTree()).visit(cu).stream())
+                .flatMap(stat -> {
+                    ShiftFormatRightVisitor shiftRight = new ShiftFormatRightVisitor(stat.getId(), enclosingIndent(insertionScope.getTree()) +
+                            formatter.findIndent(enclosingIndent(insertionScope.getTree()), stat).getEnclosingIndent(), formatter.isIndentedWithSpaces());
+                    return shiftRight.visit(cu).stream();
+                })
                 .collect(toList())).visit(cu);
 
         List<Tree> formattedStatements = formattedCu.getClasses().get(0).getBody().getStatements();

diff --git a/rewrite-core/src/main/java/com/netflix/rewrite/tree/visitor/refactor/Formatter.java b/rewrite-core/src/main/java/com/netflix/rewrite/tree/visitor/refactor/Formatter.java
@@ -101,16 +101,20 @@ public Formatting format(Cursor cursor) {
     public ShiftFormatRightVisitor shiftRight(Tree moving, Tree into, Tree enclosesBoth) {
         // NOTE: This isn't absolutely perfect... suppose the block moving was indented with tabs and the surrounding source was spaces.
         // Should be close enough in the vast majority of cases.
-        int shift = enclosingIndent(into) + findIndent(enclosingIndent(enclosesBoth), moving).getEnclosingIndent();
+        int shift = enclosingIndent(into) - findIndent(enclosingIndent(enclosesBoth), moving).getEnclosingIndent();
         return new ShiftFormatRightVisitor(moving.getId(), shift, wholeSourceIndent().isIndentedWithSpaces());
     }
 
-    private int enclosingIndent(Tree enclosesBoth) {
+    public static int enclosingIndent(Tree enclosesBoth) {
         return enclosesBoth instanceof Tr.Block ? ((Tr.Block<?>) enclosesBoth).getIndent() :
                 (int) enclosesBoth.getFormatting().getPrefix().chars().dropWhile(c -> c == '\n' || c == '\r')
                         .takeWhile(Character::isWhitespace).count();
     }
 
+    public boolean isIndentedWithSpaces() {
+        return wholeSourceIndent().isIndentedWithSpaces();
+    }
+
     /**
      * Discover the most common indentation level of a tree, and whether this indentation is built with spaces or tabs.
      */

diff --git a/rewrite-core/src/test/kotlin/com/netflix/rewrite/tree/TreeBuilderTest.kt b/rewrite-core/src/test/kotlin/com/netflix/rewrite/tree/TreeBuilderTest.kt
@@ -120,4 +120,4 @@ class TreeBuilderTest {
         assertEquals("a.A.*", name.printTrimmed())
         assertEquals("*", name.simpleName)
     }
-}
+}