Merge pull request #1924 from lf-lang/do-not-drop-comments

Fix for edge case in which comments are dropped
lf-lang · Aug 1, 2023 · 081351e · 081351e
2 parents 371eb53 + 0c0171c
commit 081351e
Showing 1 changed file with 32 additions and 11 deletions.
diff --git a/core/src/main/java/org/lflang/ast/ToLf.java b/core/src/main/java/org/lflang/ast/ToLf.java
@@ -2,7 +2,6 @@
 
 import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.Collection;
 import java.util.Comparator;
 import java.util.HashSet;
 import java.util.List;
@@ -117,14 +116,23 @@ public MalleableString doSwitch(EObject eObject) {
     Predicate<INode> doesNotBelongToPrevious =
         doesNotBelongToAncestor.and(
             previous == null ? n -> true : ASTUtils.sameLine(previous).negate());
-    Stream<String> precedingComments =
-        ASTUtils.getPrecedingComments(node, doesNotBelongToPrevious).map(String::strip);
-    Collection<String> allComments = new ArrayList<>();
-    precedingComments.forEachOrdered(allComments::add);
-    getContainedComments(node).stream()
-        .filter(doesNotBelongToAncestor)
-        .map(INode::getText)
-        .forEachOrdered(allComments::add);
+    var allComments = new ArrayList<String>();
+    if (eObject.eContents().isEmpty() && !(eObject instanceof Code)) {
+      getContainedComments(node).stream()
+          .filter(doesNotBelongToAncestor)
+          .filter(doesNotBelongToPrevious)
+          .map(INode::getText)
+          .forEach(allComments::add);
+    } else {
+      Stream<String> precedingComments =
+          ASTUtils.getPrecedingComments(node, doesNotBelongToPrevious.and(doesNotBelongToAncestor))
+              .map(String::strip);
+      precedingComments.forEachOrdered(allComments::add);
+      getContainedCodeComments(node).stream()
+          .filter(doesNotBelongToAncestor)
+          .map(INode::getText)
+          .forEach(allComments::add);
+    }
     allComments.addAll(followingComments);
     if (allComments.stream().anyMatch(s -> KEEP_FORMAT_COMMENT.matcher(s).matches())) {
       return MalleableString.anyOf(StringUtil.trimCodeBlock(node.getText(), 0))
@@ -140,8 +148,9 @@ static Set<INode> getAncestorComments(INode node) {
     for (ICompositeNode ancestor = node.getParent();
         ancestor != null;
         ancestor = ancestor.getParent()) {
-      ancestorComments.addAll(getContainedComments(ancestor));
+      ancestorComments.addAll(getContainedCodeComments(ancestor));
       ASTUtils.getPrecedingCommentNodes(ancestor, u -> true).forEachOrdered(ancestorComments::add);
+      ancestorComments.addAll(getContainedCodeComments(ancestor));
     }
     return ancestorComments;
   }
@@ -192,7 +201,7 @@ private static Stream<String> getFollowingComments(
    * Return comments contained by {@code node} that logically belong to this node (and not to any of
    * its children).
    */
-  private static List<INode> getContainedComments(INode node) {
+  private static List<INode> getContainedCodeComments(INode node) {
     ArrayList<INode> ret = new ArrayList<>();
     boolean inSemanticallyInsignificantLeadingRubbish = true;
     for (INode child : node.getAsTreeIterable()) {
@@ -212,6 +221,18 @@ private static List<INode> getContainedComments(INode node) {
     return ret;
   }
 
+  /**
+   * Return all comments that are part of {@code node}, regardless of where they appear relative to
+   * the main content of the node.
+   */
+  private static List<INode> getContainedComments(INode node) {
+    var ret = new ArrayList<INode>();
+    for (INode child : node.getAsTreeIterable()) {
+      if (ASTUtils.isComment(child)) ret.add(child);
+    }
+    return ret;
+  }
+
   @Override
   public MalleableString caseCodeExpr(CodeExpr object) {
     return caseCode(object.getCode());