Use trim_tries to extract post comment over simple trim_matches

[rust.git] / src / chains.rs

diff --git a/src/chains.rs b/src/chains.rs
index 3f5ff75a722880916ed7b91a7d8841099a46b930..f73f57a19a6ab6bb833e4de7dc430132621467e8 100644 (file)
--- a/src/chains.rs
+++ b/src/chains.rs
@@ -65,7 +65,7 @@
 //!            .qux
 //! ```
 
-use comment::rewrite_comment;
+use comment::{rewrite_comment, CharClasses, FullCodeCharKind, RichChar};
 use config::IndentStyle;
 use expr::rewrite_call;
 use lists::{extract_post_comment, extract_pre_comment, get_comment_end};
@@ -288,9 +288,9 @@ fn handle_post_comment(
                 return;
             }
             // HACK: Treat `?`s as separators.
-            let trimmed_snippet = post_comment_snippet.trim_matches('?');
-            let comment_end = get_comment_end(trimmed_snippet, "?", "", false);
-            let maybe_post_comment = extract_post_comment(trimmed_snippet, comment_end, "?")
+            let trimmed_snippet = trim_tries(post_comment_snippet);
+            let comment_end = get_comment_end(&trimmed_snippet, "?", "", false);
+            let maybe_post_comment = extract_post_comment(&trimmed_snippet, comment_end, "?")
                 .and_then(|comment| {
                     if comment.is_empty() {
                         None
@@ -334,9 +334,8 @@ fn handle_post_comment(
             // Pre-comment
             if handle_comment {
                 let pre_comment_span = mk_sp(prev_span_end, chain_item.span.lo());
-                let pre_comment_snippet = context.snippet(pre_comment_span);
-                let pre_comment_snippet = pre_comment_snippet.trim().trim_matches('?');
-                let (pre_comment, _) = extract_pre_comment(pre_comment_snippet);
+                let pre_comment_snippet = trim_tries(context.snippet(pre_comment_span));
+                let (pre_comment, _) = extract_pre_comment(&pre_comment_snippet);
                 match pre_comment {
                     Some(ref comment) if !comment.is_empty() => {
                         children.push(ChainItem::comment(
@@ -870,3 +869,28 @@ fn is_block_expr(context: &RewriteContext, expr: &ast::Expr, repr: &str) -> bool
         _ => false,
     }
 }
+
+/// Remove try operators (`?`s) that appear in the given string. If removing
+/// them leaves an empty line, remove that line as well unless it is the first
+/// line (we need the first newline for detecting pre/post comment).
+fn trim_tries(s: &str) -> String {
+    let mut result = String::with_capacity(s.len());
+    let mut line_buffer = String::with_capacity(s.len());
+    for (kind, rich_char) in CharClasses::new(s.chars()) {
+        match rich_char.get_char() {
+            '\n' => {
+                if result.is_empty() || !line_buffer.trim().is_empty() {
+                    result.push_str(&line_buffer);
+                    result.push('\n')
+                }
+                line_buffer.clear();
+            }
+            '?' if kind == FullCodeCharKind::Normal => continue,
+            c => line_buffer.push(c),
+        }
+    }
+    if !line_buffer.trim().is_empty() {
+        result.push_str(&line_buffer);
+    }
+    result
+}