fix: idempotent issue for continue/break/new with comments #114

src/data_model.rs

@@ -2758,12 +2758,17 @@ impl ReturnStatement {
 
 impl<'a> DocBuild<'a> for ReturnStatement {
     fn build_inner(&self, b: &'a DocBuilder<'a>, result: &mut Vec<DocRef<'a>>) {
+        if self.exp.is_none() {
+            build_with_comments_and_punc_attached(b, &self.node_context, result, |b, result| {
+                result.push(b.txt("return"));
+            });
+            return;
+        }
+
         build_with_comments_and_punc(b, &self.node_context, result, |b, result| {
             result.push(b.txt("return"));
-            if let Some(ref exp) = self.exp {
-                result.push(b.txt(" "));
-                result.push(exp.build(b));
-            }
+            result.push(b.txt(" "));
+            result.push(self.exp.as_ref().unwrap().build(b));
         });
     }
 }
@@ -3324,22 +3329,33 @@ impl BreakStatement {
     pub fn new(node: Node) -> Self {
         assert_check(node, "break_statement");
 
+        let identifier = node.try_c_by_k("identifier").map(|n| ValueNode::new(n));
+        let node_context = if identifier.is_none() {
+            NodeContext::with_inner_punctuation(&node)
+        } else {
+            NodeContext::with_punctuation(&node)
+        };
+
         Self {
-            identifier: node.try_c_by_k("identifier").map(|n| ValueNode::new(n)),
-            node_context: NodeContext::with_inner_punctuation(&node),
+            identifier,
+            node_context,
         }
     }
 }
 
 impl<'a> DocBuild<'a> for BreakStatement {
     fn build_inner(&self, b: &'a DocBuilder<'a>, result: &mut Vec<DocRef<'a>>) {
+        if self.identifier.is_none() {
+            build_with_comments_and_punc_attached(b, &self.node_context, result, |b, result| {
+                result.push(b.txt("break"));
+            });
+            return;
+        }
+
         build_with_comments_and_punc(b, &self.node_context, result, |b, result| {
             result.push(b.txt("break"));
-
-            if let Some(ref n) = self.identifier {
-                result.push(b.txt(" "));
-                result.push(n.build(b));
-            }
+            result.push(b.txt(" "));
+            result.push(self.identifier.as_ref().unwrap().build(b));
         });
     }
 }
@@ -3370,13 +3386,17 @@ impl ContinueStatement {
 
 impl<'a> DocBuild<'a> for ContinueStatement {
     fn build_inner(&self, b: &'a DocBuilder<'a>, result: &mut Vec<DocRef<'a>>) {
+        if self.identifier.is_none() {
+            build_with_comments_and_punc_attached(b, &self.node_context, result, |b, result| {
+                result.push(b.txt("continue"));
+            });
+            return;
+        }
+
         build_with_comments_and_punc(b, &self.node_context, result, |b, result| {
             result.push(b.txt("continue"));
-
-            if let Some(ref n) = self.identifier {
-                result.push(b.txt(" "));
-                result.push(n.build(b));
-            }
+            result.push(b.txt(" "));
+            result.push(self.identifier.as_ref().unwrap().build(b));
         });
     }
 }

src/formatter.rs

@@ -13,6 +13,9 @@ use std::thread;
 use std::{fs, path::Path};
 use tree_sitter::{Node, Parser, Tree};
 
+#[allow(unused_imports)]
+use crate::utility::print_comment_map;
+
 #[derive(Clone, Debug, Deserialize)]
 pub struct Config {
     #[serde(default = "default_max_width")]
@@ -154,7 +157,8 @@ impl Formatter {
 
         let result = pretty_print(doc_ref, config.max_width);
 
-        //print_comment_map(&ast_tree);
+        // debugging tool: use this to print named node value + comments in bucket
+        // print_comment_map(&ast_tree);
 
         assert_no_missing_comments();
 

src/utility.rs

@@ -284,6 +284,24 @@ pub fn build_with_comments<'a, F>(
     handle_post_comments(b, bucket, result);
 }
 
+pub fn build_with_comments_core<'a, F>(
+    b: &'a DocBuilder<'a>,
+    node_context: &NodeContext,
+    result: &mut Vec<DocRef<'a>>,
+    handle_members: F,
+) where
+    F: FnOnce(&'a DocBuilder<'a>, &mut Vec<DocRef<'a>>),
+{
+    let bucket = get_comment_bucket(&node_context.id);
+    handle_pre_comments(b, bucket, result);
+
+    if bucket.dangling_comments.is_empty() {
+        handle_members(b, result);
+    } else {
+        result.push(b.concat(handle_dangling_comments(b, bucket)));
+    }
+}
+
 pub fn build_with_comments_and_punc<'a, F>(
     b: &'a DocBuilder<'a>,
     node_context: &NodeContext,
@@ -292,13 +310,40 @@ pub fn build_with_comments_and_punc<'a, F>(
 ) where
     F: FnOnce(&'a DocBuilder<'a>, &mut Vec<DocRef<'a>>),
 {
-    build_with_comments(b, node_context, result, handle_members);
+    build_with_comments_core(b, node_context, result, handle_members);
+
+    let bucket = get_comment_bucket(&node_context.id);
+    if bucket.dangling_comments.is_empty() {
+        handle_post_comments(b, bucket, result);
+    }
 
     if let Some(ref n) = node_context.punc {
         result.push(n.build(b));
     }
 }
 
+// fix: https://github.com/xixiaofinland/afmt/issues/114
+pub fn build_with_comments_and_punc_attached<'a, F>(
+    b: &'a DocBuilder<'a>,
+    node_context: &NodeContext,
+    result: &mut Vec<DocRef<'a>>,
+    handle_members: F,
+) where
+    F: FnOnce(&'a DocBuilder<'a>, &mut Vec<DocRef<'a>>),
+{
+    build_with_comments_core(b, node_context, result, handle_members);
+
+    let bucket = get_comment_bucket(&node_context.id);
+
+    if let Some(ref n) = node_context.punc {
+        result.push(n.build(b));
+    }
+
+    if bucket.dangling_comments.is_empty() {
+        handle_post_comments(b, bucket, result);
+    }
+}
+
 pub fn handle_dangling_comments_in_bracket_surround<'a>(
     b: &'a DocBuilder<'a>,
     bucket: &CommentBucket,

tests/static/punc_and_comment_flip.cls

@@ -0,0 +1,77 @@
+// fix: https://github.com/xixiaofinland/afmt/issues/114
+class TestClass {
+  {
+    continue; /*t1*/
+
+    continue; //t2
+
+    continue;
+    // t3
+
+    continue;
+    // t4
+
+    continue true; //t5
+
+    continue;
+    // t4 /*t*/
+
+    continue
+    /*t*/;
+    // t4
+
+    continue
+    /*t*/;
+    // t4
+  }
+
+  {
+    break; /*t1*/
+
+    break; //t2
+
+    break;
+    // t3
+
+    break;
+    // t4
+
+    break true; //t5
+
+    break;
+    // t4 /*t*/
+
+    break
+    /*t*/;
+    // t4
+
+    break
+    /*t*/;
+    // t4
+  }
+
+  {
+    return; /*t1*/
+
+    return; //t2
+
+    return;
+    // t3
+
+    return;
+    // t4
+
+    return true; //t5
+
+    return;
+    // t4 /*t*/
+
+    return
+    /*t*/;
+    // t4
+
+    return
+    /*t*/;
+    // t4
+  }
+}

tests/static/punc_and_comment_flip.in

@@ -0,0 +1,89 @@
+// fix: https://github.com/xixiaofinland/afmt/issues/114
+class TestClass {
+  {
+    continue; /*t1*/
+
+    continue; //t2
+
+    continue // t3
+    ;
+
+    continue
+    // t4
+    ;
+
+    continue true; //t5
+
+    continue
+    // t4 /*t*/
+    ;
+
+    continue
+    // t4
+    /*t*/
+    ;
+
+    continue
+    /*t*/
+    // t4
+    ;
+  }
+
+  {
+    break; /*t1*/
+
+    break; //t2
+
+    break // t3
+    ;
+
+    break
+    // t4
+    ;
+
+    break true; //t5
+
+    break
+    // t4 /*t*/
+    ;
+
+    break
+    // t4
+    /*t*/
+    ;
+
+    break
+    /*t*/
+    // t4
+    ;
+  }
+
+  {
+    return; /*t1*/
+
+    return; //t2
+
+    return // t3
+    ;
+
+    return
+    // t4
+    ;
+
+    return true; //t5
+
+    return
+    // t4 /*t*/
+    ;
+
+    return
+    // t4
+    /*t*/
+    ;
+
+    return
+    /*t*/
+    // t4
+    ;
+  }
+}