Add handling of comment indentation to yaml AutoFormat

rewrite-yaml/src/main/java/org/openrewrite/yaml/MergeYamlVisitor.java

@@ -129,7 +129,7 @@ private Yaml.Mapping mergeMapping(Yaml.Mapping m1, Yaml.Mapping m2, P p, Cursor
                 }
             }
             if (shouldAutoFormat) {
-                return autoFormat(incomingEntry, p, cursor);
+                incomingEntry = autoFormat(incomingEntry, p, cursor);
             }
             return incomingEntry;
         }));

rewrite-yaml/src/main/java/org/openrewrite/yaml/format/AutoFormatVisitor.java

@@ -36,41 +36,40 @@ public AutoFormatVisitor(@Nullable Tree stopAfter) {
     }
 
     @Override
-    public @Nullable Yaml preVisit(Yaml tree, P p) {
+    public Yaml preVisit(Yaml tree, P p) {
         stopAfterPreVisit();
         Yaml.Documents docs = getCursor().firstEnclosingOrThrow(Yaml.Documents.class);
         Cursor cursor = getCursor().getParentOrThrow();
 
-        Yaml y = new NormalizeFormatVisitor<>(stopAfter).visit(tree, p, cursor.fork());
+        Yaml y = new NormalizeFormatVisitor<>(stopAfter).visitNonNull(tree, p, cursor.fork());
 
-        y = new MinimumViableSpacingVisitor<>(stopAfter).visit(y, p, cursor.fork());
+        y = new MinimumViableSpacingVisitor<>(stopAfter).visitNonNull(y, p, cursor.fork());
 
         y = new IndentsVisitor<>(Optional.ofNullable(docs.getStyle(IndentsStyle.class))
                 .orElse(Autodetect.tabsAndIndents(docs, YamlDefaultStyles.indents())), stopAfter)
-                .visit(y, p, cursor.fork());
+                .visitNonNull(y, p, cursor.fork());
 
         y = new NormalizeLineBreaksVisitor<>(Optional.ofNullable(docs.getStyle(GeneralFormatStyle.class))
                 .orElse(Autodetect.generalFormat(docs)), stopAfter)
-                .visit(y, p, cursor.fork());
+                .visitNonNull(y, p, cursor.fork());
 
         return y;
     }
 
     @Override
     public Yaml.Documents visitDocuments(Yaml.Documents documents, P p) {
-        Yaml.Documents y = (Yaml.Documents) new NormalizeFormatVisitor<>(stopAfter).visit(documents, p);
+        Yaml.Documents y = (Yaml.Documents) new NormalizeFormatVisitor<>(stopAfter).visitNonNull(documents, p);
 
-        y = (Yaml.Documents) new MinimumViableSpacingVisitor<>(stopAfter).visit(y, p);
+        y = (Yaml.Documents) new MinimumViableSpacingVisitor<>(stopAfter).visitNonNull(y, p);
 
         y = (Yaml.Documents) new IndentsVisitor<>(Optional.ofNullable(documents.getStyle(IndentsStyle.class))
                 .orElse(Autodetect.tabsAndIndents(y, YamlDefaultStyles.indents())), stopAfter)
-                .visit(documents, p);
+                .visitNonNull(documents, p);
 
         y = (Yaml.Documents) new NormalizeLineBreaksVisitor<>(Optional.ofNullable(documents.getStyle(GeneralFormatStyle.class))
                 .orElse(Autodetect.generalFormat(y)), stopAfter)
-                .visit(documents, p);
+                .visitNonNull(documents, p);
 
-        assert y != null;
         return y;
     }
 

rewrite-yaml/src/main/java/org/openrewrite/yaml/format/IndentsVisitor.java

@@ -18,6 +18,7 @@
 import org.jspecify.annotations.Nullable;
 import org.openrewrite.Cursor;
 import org.openrewrite.Tree;
+import org.openrewrite.internal.StringUtils;
 import org.openrewrite.yaml.YamlIsoVisitor;
 import org.openrewrite.yaml.style.IndentsStyle;
 import org.openrewrite.yaml.tree.Yaml;
@@ -73,19 +74,24 @@ public IndentsVisitor(IndentsStyle style, @Nullable Tree stopAfter) {
 
                 getCursor().getParentOrThrow().putMessage("sequenceEntryIndent", indent);
                 // the +1 is for the '-' character
-                getCursor().getParentOrThrow().putMessage("lastIndent", indent +
-                        firstIndent(((Yaml.Sequence.Entry) y).getBlock()).length() + 1);
+                getCursor().getParentOrThrow().putMessage("lastIndent",
+                        indent + firstIndent(((Yaml.Sequence.Entry) y).getBlock()).length() + 1);
             } else if (y instanceof Yaml.Mapping.Entry) {
                 y = y.withPrefix(indentTo(y.getPrefix(), indent + style.getIndentSize()));
                 getCursor().putMessage("lastIndent", indent + style.getIndentSize());
+            } else if (y instanceof Yaml.Document) {
+                y = y.withPrefix(indentComments(y.getPrefix(), 0));
+            }
+        } else if (y instanceof Yaml.Mapping.Entry) {
+            if (getCursor().getParentOrThrow(2).getValue() instanceof Yaml.Sequence.Entry) {
+                // this is a mapping entry that begins a sequence entry and anything below it should be indented further to the right now, e.g.:
+                //
+                // - key:
+                //     value
+                getCursor().putMessage("lastIndent", indent + style.getIndentSize());
+            } else {
+                y = y.withPrefix(indentComments(y.getPrefix(), indent));
             }
-        } else if (y instanceof Yaml.Mapping.Entry &&
-                getCursor().getParentOrThrow(2).getValue() instanceof Yaml.Sequence.Entry) {
-            // this is a mapping entry that begins a sequence entry and anything below it should be indented further to the right now, e.g.:
-            //
-            // - key:
-            //     value
-            getCursor().putMessage("lastIndent", indent + style.getIndentSize());
         }
         return y;
     }
@@ -117,7 +123,7 @@ private String indentTo(String prefix, int column) {
         }
 
         int indent = findIndent(prefix);
-
+        prefix = indentComments(prefix, indent);
         if (indent != column) {
             int shift = column - indent;
             prefix = indent(prefix, shift);
@@ -126,6 +132,21 @@ private String indentTo(String prefix, int column) {
         return prefix;
     }
 
+    private String indentComments(String prefix, int indent) {
+        // If the prefix contains a newline followed by a comment ensure the comment begins at the indentation column
+        if (prefix.contains("#")) {
+            String reindentedComments = prefix.replaceAll("\n\\s*#", "\n" + StringUtils.repeat(" ", indent) + "#");
+            // If a document begins with a comment it might not have a newline before it
+            if (getCursor().getValue() instanceof Yaml.Document) {
+                reindentedComments = prefix.replaceFirst("^\\s*#", "#");
+            }
+            if (!reindentedComments.equals(prefix)) {
+                prefix = reindentedComments;
+            }
+        }
+        return prefix;
+    }
+
     private String indent(String whitespace, int shift) {
         StringBuilder newWhitespace = new StringBuilder(whitespace);
         shift(newWhitespace, shift);
@@ -154,12 +175,13 @@ private int findIndent(String prefix) {
     }
 
     private String firstIndent(Yaml yaml) {
-        AtomicReference<String> indent = new AtomicReference<>();
+        AtomicReference<@Nullable String> indent = new AtomicReference<>();
 
         new YamlIsoVisitor<AtomicReference<String>>() {
             @Override
             public @Nullable Yaml visit(@Nullable Tree tree, AtomicReference<String> indent) {
                 Yaml y = (Yaml) tree;
+                //noinspection ConstantValue
                 if (indent.get() != null) {
                     return y;
                 }

rewrite-yaml/src/test/java/org/openrewrite/yaml/MergeYamlTest.java

@@ -1145,7 +1145,7 @@ void comment() {
                 //language=yaml
                 """
                   
-                  # new stuff
+                    # new stuff
                   new-property: value
                   """,
                 false,

rewrite-yaml/src/test/java/org/openrewrite/yaml/format/IndentsTest.java

@@ -25,6 +25,7 @@
 import static org.openrewrite.test.RewriteTest.toRecipe;
 import static org.openrewrite.yaml.Assertions.yaml;
 
+@SuppressWarnings("KubernetesUnknownResourcesInspection")
 class IndentsTest implements RewriteTest {
 
     @Override
@@ -115,4 +116,67 @@ void maintainIndentSpacingOnMixedTypeSequences() {
           )
         );
     }
+
+    @Test
+    void indentSequenceComments() {
+        rewriteRun(
+          yaml("""
+              key:
+              # under-indented
+                  # over-indented
+                - a
+              """,
+            """
+              key:
+                # under-indented
+                # over-indented
+                - a
+              """
+          )
+        );
+    }
+
+    @Test
+    void indentMappingComments() {
+        rewriteRun(
+          yaml("""
+              key: # no change
+              # under-indented
+                  # over-indented
+                a : # no change
+                # under-indented
+                    # over-indented
+                  b : c
+              """,
+            """
+              key: # no change
+                # under-indented
+                # over-indented
+                a : # no change
+                  # under-indented
+                  # over-indented
+                  b : c
+              """
+          )
+        );
+    }
+
+    @Test
+    void indentRootComments() {
+        rewriteRun(
+          yaml("""
+                # over-indented 1
+              key: value # no change
+                # over-indented 2
+              key2: value2
+              """,
+            """
+              # over-indented 1
+              key: value # no change
+              # over-indented 2
+              key2: value2
+              """
+          )
+        );
+    }
 }