When merging yaml files comments on the last element of existing block are placed at the wrong line

rewrite-core/src/main/java/org/openrewrite/Cursor.java

@@ -57,7 +57,7 @@ public Cursor getRoot() {
     /**
      * @return true if this cursor is the root of the tree, false otherwise
      */
-    final boolean isRoot() {
+    final public boolean isRoot() {
         return ROOT_VALUE.equals(value);
     }
 

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

@@ -84,6 +84,7 @@ public String getDescription() {
     }
 
     final static String FOUND_MATCHING_ELEMENT = "FOUND_MATCHING_ELEMENT";
+    final static String REMOVE_PREFIX = "REMOVE_PREFIX";
 
     @Override
     public TreeVisitor<?, ExecutionContext> getVisitor() {
@@ -96,7 +97,7 @@ public TreeVisitor<?, ExecutionContext> getVisitor() {
                     .map(docs -> {
                         // Any comments will have been put on the parent Document node, preserve by copying to the mapping
                         Yaml.Document doc = docs.getDocuments().get(0);
-                        if(doc.getBlock() instanceof Yaml.Mapping) {
+                        if (doc.getBlock() instanceof Yaml.Mapping) {
                             Yaml.Mapping m = (Yaml.Mapping) doc.getBlock();
                             return m.withEntries(ListUtils.mapFirst(m.getEntries(), entry -> entry.withPrefix(doc.getPrefix())));
                         } else if (doc.getBlock() instanceof Yaml.Sequence) {
@@ -110,9 +111,11 @@ public TreeVisitor<?, ExecutionContext> getVisitor() {
             @Override
             public Yaml.Document visitDocument(Yaml.Document document, ExecutionContext ctx) {
                 if ("$".equals(key)) {
-                    return document.withBlock((Yaml.Block) new MergeYamlVisitor<>(document.getBlock(), yaml,
-                            Boolean.TRUE.equals(acceptTheirs), objectIdentifyingProperty).visitNonNull(document.getBlock(),
-                            ctx, getCursor()));
+                    Yaml.Document d = document.withBlock((Yaml.Block)
+                            new MergeYamlVisitor<>(document.getBlock(), yaml, Boolean.TRUE.equals(acceptTheirs), objectIdentifyingProperty)
+                                    .visitNonNull(document.getBlock(), ctx, getCursor())
+                    );
+                    return getCursor().getMessage(REMOVE_PREFIX, false) ? d.withEnd(d.getEnd().withPrefix("")) : d;
                 }
                 Yaml.Document d = super.visitDocument(document, ctx);
                 if (d == document && !getCursor().getMessage(FOUND_MATCHING_ELEMENT, false)) {

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

@@ -22,16 +22,28 @@
 import org.openrewrite.Cursor;
 import org.openrewrite.internal.ListUtils;
 import org.openrewrite.yaml.tree.Yaml;
+import org.openrewrite.yaml.tree.Yaml.Document;
+import org.openrewrite.yaml.tree.Yaml.Mapping;
+import org.openrewrite.yaml.tree.Yaml.Mapping.Entry;
+import org.openrewrite.yaml.tree.Yaml.Scalar;
 
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Optional;
+import java.util.regex.Pattern;
 import java.util.stream.Collectors;
 
+import static org.openrewrite.Cursor.ROOT_VALUE;
+import static org.openrewrite.internal.ListUtils.*;
+import static org.openrewrite.yaml.MergeYaml.REMOVE_PREFIX;
+
 @AllArgsConstructor
 @RequiredArgsConstructor
 public class MergeYamlVisitor<P> extends YamlVisitor<P> {
-    private final Yaml scope;
+
+    private static final Pattern LINE_BREAK = Pattern.compile("\\R");
+
+    private final Yaml existing;
     private final Yaml incoming;
     private final boolean acceptTheirs;
 
@@ -47,9 +59,9 @@ public MergeYamlVisitor(Yaml scope, @Language("yml") String yamlString, boolean
                         .map(Yaml.Documents.class::cast)
                         .map(docs -> {
                             // Any comments will have been put on the parent Document node, preserve by copying to the mapping
-                            Yaml.Document doc = docs.getDocuments().get(0);
-                            if(doc.getBlock() instanceof Yaml.Mapping) {
-                                Yaml.Mapping m = (Yaml.Mapping) doc.getBlock();
+                            Document doc = docs.getDocuments().get(0);
+                            if (doc.getBlock() instanceof Mapping) {
+                                Mapping m = (Mapping) doc.getBlock();
                                 return m.withEntries(ListUtils.mapFirst(m.getEntries(), entry -> entry.withPrefix(doc.getPrefix())));
                             } else if (doc.getBlock() instanceof Yaml.Sequence) {
                                 Yaml.Sequence s = (Yaml.Sequence) doc.getBlock();
@@ -63,24 +75,24 @@ public MergeYamlVisitor(Yaml scope, @Language("yml") String yamlString, boolean
     }
 
     @Override
-    public Yaml visitScalar(Yaml.Scalar existingScalar, P p) {
-        if (scope.isScope(existingScalar) && incoming instanceof Yaml.Scalar) {
-            return mergeScalar(existingScalar, (Yaml.Scalar) incoming);
+    public Yaml visitScalar(Scalar existingScalar, P p) {
+        if (existing.isScope(existingScalar) && incoming instanceof Scalar) {
+            return mergeScalar(existingScalar, (Scalar) incoming);
         }
         return super.visitScalar(existingScalar, p);
     }
 
     @Override
     public Yaml visitSequence(Yaml.Sequence existingSeq, P p) {
-        if (scope.isScope(existingSeq)) {
-            if (incoming instanceof Yaml.Mapping) {
+        if (existing.isScope(existingSeq)) {
+            if (incoming instanceof Mapping) {
                 // Distribute the incoming mapping to each entry in the sequence
-                return existingSeq.withEntries(ListUtils.map(existingSeq.getEntries(), (i, existingSeqEntry) -> {
-                    Yaml.Block b = (Yaml.Block) new MergeYamlVisitor<>(existingSeqEntry.getBlock(), incoming,
-                            acceptTheirs, objectIdentifyingProperty, shouldAutoFormat)
-                            .visitNonNull(existingSeqEntry.getBlock(), p, new Cursor(getCursor(), existingSeqEntry));
-                    return existingSeqEntry.withBlock(b);
-                }));
+                return existingSeq.withEntries(map(existingSeq.getEntries(), (i, existingSeqEntry) ->
+                        existingSeqEntry.withBlock((Yaml.Block)
+                                new MergeYamlVisitor<>(existingSeqEntry.getBlock(), incoming, acceptTheirs, objectIdentifyingProperty, shouldAutoFormat)
+                                        .visitNonNull(existingSeqEntry.getBlock(), p, new Cursor(getCursor(), existingSeqEntry))
+                        )
+                ));
             } else if (incoming instanceof Yaml.Sequence) {
                 return mergeSequence(existingSeq, (Yaml.Sequence) incoming, p, getCursor());
             }
@@ -89,96 +101,160 @@ public Yaml visitSequence(Yaml.Sequence existingSeq, P p) {
     }
 
     @Override
-    public Yaml visitMapping(Yaml.Mapping existingMapping, P p) {
-        if (scope.isScope(existingMapping) && incoming instanceof Yaml.Mapping) {
-            return mergeMapping(existingMapping, (Yaml.Mapping) incoming, p, getCursor());
+    public Yaml visitMapping(Mapping existingMapping, P p) {
+        if (existing.isScope(existingMapping) && incoming instanceof Mapping) {
+            Mapping mapping = mergeMapping(existingMapping, (Mapping) incoming, p, getCursor());
+
+            if (getCursor().getMessage(REMOVE_PREFIX, false)) {
+                List<Entry> entries = ((Mapping) getCursor().getValue()).getEntries();
+                return mapping.withEntries(mapLast(mapping.getEntries(), it -> it.withPrefix("\n" + grabPartLineBreak(entries.get(entries.size() - 1), 1))));
+            }
+
+            return mapping;
         }
         return super.visitMapping(existingMapping, p);
     }
 
-    private static boolean keyMatches(Yaml.Mapping.Entry e1, Yaml.Mapping.Entry e2) {
-        return e1.getKey().getValue().equals(e2.getKey().getValue());
+    private static boolean keyMatches(@Nullable Entry e1, @Nullable Entry e2) {
+        return e1 != null && e2 != null && e1.getKey().getValue().equals(e2.getKey().getValue());
     }
 
-    private boolean keyMatches(Yaml.Mapping m1, Yaml.Mapping m2) {
+    private boolean keyMatches(Mapping m1, Mapping m2) {
         Optional<String> nameToAdd = m2.getEntries().stream()
                 .filter(e -> objectIdentifyingProperty != null && objectIdentifyingProperty.equals(e.getKey().getValue()))
-                .map(e -> ((Yaml.Scalar) e.getValue()).getValue())
+                .map(e -> ((Scalar) e.getValue()).getValue())
                 .findAny();
 
         return nameToAdd.map(nameToAddValue -> m1.getEntries().stream()
                         .filter(e -> objectIdentifyingProperty.equals(e.getKey().getValue()))
-                        .map(e -> ((Yaml.Scalar) e.getValue()).getValue())
+                        .map(e -> ((Scalar) e.getValue()).getValue())
                         .anyMatch(existingName -> existingName.equals(nameToAddValue)))
                 .orElse(false);
     }
 
-    private Yaml.Mapping mergeMapping(Yaml.Mapping m1, Yaml.Mapping m2, P p, Cursor cursor) {
-        List<Yaml.Mapping.Entry> mutatedEntries = ListUtils.map(m1.getEntries(), existingEntry -> {
-            for (Yaml.Mapping.Entry incomingEntry : m2.getEntries()) {
+    private Mapping mergeMapping(Mapping m1, Mapping m2, P p, Cursor cursor) {
+        List<Entry> mergedEntries = map(m1.getEntries(), existingEntry -> {
+            for (Entry incomingEntry : m2.getEntries()) {
                 if (keyMatches(existingEntry, incomingEntry)) {
-                    return existingEntry.withValue((Yaml.Block) new MergeYamlVisitor<>(existingEntry.getValue(),
-                            incomingEntry.getValue(), acceptTheirs, objectIdentifyingProperty, shouldAutoFormat)
-                            .visitNonNull(existingEntry.getValue(), p, new Cursor(cursor, existingEntry)));
+                    return existingEntry.withValue((Yaml.Block)
+                            new MergeYamlVisitor<>(existingEntry.getValue(), incomingEntry.getValue(), acceptTheirs, objectIdentifyingProperty, shouldAutoFormat)
+                                    .visitNonNull(existingEntry.getValue(), p, new Cursor(cursor, existingEntry)));
                 }
             }
             return existingEntry;
         });
 
-        mutatedEntries = ListUtils.concatAll(mutatedEntries, ListUtils.map(m2.getEntries(), incomingEntry -> {
-            for (Yaml.Mapping.Entry existingEntry : m1.getEntries()) {
-                if (keyMatches(existingEntry, incomingEntry)) {
+        List<Entry> mutatedEntries = concatAll(mergedEntries, map(m2.getEntries(), it -> {
+            for (Entry existingEntry : m1.getEntries()) {
+                if (keyMatches(existingEntry, it)) {
                     return null;
                 }
             }
-            if (shouldAutoFormat) {
-                incomingEntry = autoFormat(incomingEntry, p, cursor);
+            // workaround: autoFormat put sometimes extra spaces before elements
+            if (!mergedEntries.isEmpty() && it.getValue() instanceof Scalar && hasLineBreak(mergedEntries.get(0), 2)) {
+                return it.withPrefix("\n" + grabPartLineBreak(mergedEntries.get(0), 1));
             }
-            return incomingEntry;
+            return shouldAutoFormat ? autoFormat(it, p, cursor) : it;
         }));
 
+        if (m1.getEntries().size() < mutatedEntries.size() && !getCursor().isRoot()) {
+            Cursor c = getCursor().dropParentUntil(it -> {
+                if (ROOT_VALUE.equals(it) || it instanceof Document) {
+                    return true;
+                }
+
+                if (it instanceof Mapping) {
+                    List<Entry> entries = ((Mapping) it).getEntries();
+                    // last member should search further upwards until two entries are found
+                    if (entries.get(entries.size() - 1).equals(getCursor().getParentOrThrow().getValue())) {
+                        return false;
+                    }
+                    return entries.size() > 1;
+                }
+
+                return false;
+            });
+
+            if (c.getValue() instanceof Document || c.getValue() instanceof Mapping) {
+                String comment = null;
+
+                if (c.getValue() instanceof Document) {
+                    comment = ((Document) c.getValue()).getEnd().getPrefix();
+                } else {
+                    List<Entry> entries = ((Mapping) c.getValue()).getEntries();
+
+                    // get comment from next element in same mapping block
+                    for (int i = 0; i < entries.size() - 1; i++) {
+                        if (entries.get(i).getValue().equals(getCursor().getValue())) {
+                            comment = grabPartLineBreak(entries.get(i + 1), 0);
+                            break;
+                        }
+                    }
+                    // or retrieve it for last item from next element (could potentially be much higher in the tree)
+                    if (comment == null && hasLineBreak(entries.get(entries.size() - 1), 1)) {
+                        comment = grabPartLineBreak(entries.get(entries.size() - 1), 0);
+                    }
+                }
+
+                if (comment != null) {
+                    Entry last = mutatedEntries.get(mutatedEntries.size() - 1);
+                    mutatedEntries.set(mutatedEntries.size() - 1, last.withPrefix(comment + last.getPrefix()));
+                    c.putMessage(REMOVE_PREFIX, true);
+                }
+            }
+        }
+
+        removePrefixForDirectChildren(m1.getEntries(), mutatedEntries);
+
         return m1.withEntries(mutatedEntries);
     }
 
+    private void removePrefixForDirectChildren(List<Entry> m1Entries, List<Entry> mutatedEntries) {
+        for (int i = 0; i < m1Entries.size() - 1; i++) {
+            if (m1Entries.get(i).getValue() instanceof Mapping && mutatedEntries.get(i).getValue() instanceof Mapping &&
+                    ((Mapping) m1Entries.get(i).getValue()).getEntries().size() < ((Mapping) mutatedEntries.get(i).getValue()).getEntries().size()) {
+                mutatedEntries.set(i + 1, mutatedEntries.get(i + 1).withPrefix("\n" + grabPartLineBreak(mutatedEntries.get(i + 1), 1)));
+            }
+        }
+    }
+
     private Yaml.Sequence mergeSequence(Yaml.Sequence s1, Yaml.Sequence s2, P p, Cursor cursor) {
         if (acceptTheirs) {
             return s1;
         }
 
-        boolean isSequenceOfScalars = s2.getEntries().stream().allMatch(entry -> entry.getBlock() instanceof Yaml.Scalar);
+        boolean isSequenceOfScalars = s2.getEntries().stream().allMatch(entry -> entry.getBlock() instanceof Scalar);
 
         if (isSequenceOfScalars) {
             List<Yaml.Sequence.Entry> incomingEntries = new ArrayList<>(s2.getEntries());
 
             nextEntry:
             for (Yaml.Sequence.Entry entry : s1.getEntries()) {
-                if (entry.getBlock() instanceof Yaml.Scalar) {
-                    String existingScalar = ((Yaml.Scalar) entry.getBlock()).getValue();
+                if (entry.getBlock() instanceof Scalar) {
+                    String existingScalar = ((Scalar) entry.getBlock()).getValue();
                     for (Yaml.Sequence.Entry incomingEntry : incomingEntries) {
-                        if (((Yaml.Scalar) incomingEntry.getBlock()).getValue().equals(existingScalar)) {
+                        if (((Scalar) incomingEntry.getBlock()).getValue().equals(existingScalar)) {
                             incomingEntries.remove(incomingEntry);
                             continue nextEntry;
                         }
                     }
                 }
             }
 
-            return s1.withEntries(ListUtils.concatAll(s1.getEntries(),
-                    ListUtils.map(incomingEntries, incomingEntry -> autoFormat(incomingEntry, p, cursor))));
+            return s1.withEntries(concatAll(s1.getEntries(), map(incomingEntries, it -> autoFormat(it, p, cursor))));
         } else {
             if (objectIdentifyingProperty == null) {
                 // No identifier set to match entries on, so cannot continue
                 return s1;
             } else {
-                List<Yaml.Sequence.Entry> mutatedEntries = ListUtils.map(s2.getEntries(), entry -> {
-                    Yaml.Mapping incomingMapping = (Yaml.Mapping) entry.getBlock();
+                List<Yaml.Sequence.Entry> mutatedEntries = map(s2.getEntries(), entry -> {
+                    Mapping incomingMapping = (Mapping) entry.getBlock();
                     for (Yaml.Sequence.Entry existingEntry : s1.getEntries()) {
-                        Yaml.Mapping existingMapping = (Yaml.Mapping) existingEntry.getBlock();
+                        Mapping existingMapping = (Mapping) existingEntry.getBlock();
                         if (keyMatches(existingMapping, incomingMapping)) {
                             Yaml.Sequence.Entry e1 = existingEntry.withBlock(mergeMapping(existingMapping, incomingMapping, p, cursor));
-                            if(e1 == existingEntry) {
+                            if (e1 == existingEntry) {
                                 // Made no change, no need to consider the entry "mutated"
-                                //noinspection DataFlowIssue
                                 return null;
                             }
                             return e1;
@@ -190,10 +266,9 @@ private Yaml.Sequence mergeSequence(Yaml.Sequence s1, Yaml.Sequence s2, P p, Cur
                     return s1;
                 }
 
-                List<Yaml.Sequence.Entry> entries = ListUtils.concatAll(
-                        s1.getEntries().stream().filter(entry -> !mutatedEntries.contains(entry))
-                                .collect(Collectors.toList()),
-                        ListUtils.map(mutatedEntries, entry -> autoFormat(entry, p, cursor)));
+                List<Yaml.Sequence.Entry> entries = concatAll(
+                        s1.getEntries().stream().filter(entry -> !mutatedEntries.contains(entry)).collect(Collectors.toList()),
+                        map(mutatedEntries, entry -> autoFormat(entry, p, cursor)));
 
                 if (entries.size() != s1.getEntries().size()) {
                     return s1.withEntries(entries);
@@ -208,7 +283,17 @@ private Yaml.Sequence mergeSequence(Yaml.Sequence s1, Yaml.Sequence s2, P p, Cur
         }
     }
 
-    private Yaml.Scalar mergeScalar(Yaml.Scalar y1, Yaml.Scalar y2) {
+    private boolean hasLineBreak(Entry entry, int atLeastParts) {
+        boolean a = LINE_BREAK.matcher(entry.getPrefix()).find();
+        return a && !grabPartLineBreak(entry, atLeastParts - 1).isEmpty();
+    }
+
+    private String grabPartLineBreak(Entry entry, int index) {
+        String[] parts = LINE_BREAK.split(entry.getPrefix());
+        return parts.length > index ? parts[index] : "";
+    }
+
+    private Scalar mergeScalar(Scalar y1, Scalar y2) {
         String s1 = y1.getValue();
         String s2 = y2.getValue();
         return !s1.equals(s2) && !acceptTheirs ? y1.withValue(s2) : y1;