Ignore comment in wrap_str

2 files changed, 34 insertions, 1 deletions

diff --git a/src/comment.rs b/src/comment.rs
index ed83a3925b0..2e7402ed178 100644
--- a/src/comment.rs
+++ b/src/comment.rs
@@ -1139,6 +1139,21 @@ pub fn recover_comment_removed(
     }
 }
 
+pub fn filter_normal_code(code: &str) -> String {
+    let mut buffer = String::with_capacity(code.len());
+    LineClasses::new(code).for_each(|(kind, line)| match kind {
+        FullCodeCharKind::Normal | FullCodeCharKind::InString => {
+            buffer.push_str(&line);
+            buffer.push('\n');
+        }
+        _ => (),
+    });
+    if !code.ends_with("\n") && buffer.ends_with("\n") {
+        buffer.pop();
+    }
+    buffer
+}
+
 /// Return true if the two strings of code have the same payload of comments.
 /// The payload of comments is everything in the string except:
 ///     - actual code (not comments)
@@ -1392,4 +1407,21 @@ mod test {
         let s = format!("    r#\"\n        test\n    \"#");
         assert_eq!(remove_trailing_white_spaces(&s), s);
     }
+
+    #[test]
+    fn test_filter_normal_code() {
+        let s = r#"
+fn main() {
+    println!("hello, world");
+}
+"#;
+        assert_eq!(s, filter_normal_code(s));
+        let s_with_comment = r#"
+fn main() {
+    // hello, world
+    println!("hello, world");
+}
+"#;
+        assert_eq!(s, filter_normal_code(s_with_comment));
+    }
 }
diff --git a/src/utils.rs b/src/utils.rs
index 5f92255e79c..9eac3fd4f0d 100644
--- a/src/utils.rs
+++ b/src/utils.rs
@@ -18,6 +18,7 @@ use syntax::ast::{
 use syntax::codemap::{BytePos, Span, NO_EXPANSION};
 use syntax::ptr;
 
+use comment::filter_normal_code;
 use rewrite::RewriteContext;
 use shape::Shape;
 
@@ -350,7 +351,7 @@ macro_rules! skip_out_of_file_lines_range_visitor {
 // Wraps String in an Option. Returns Some when the string adheres to the
 // Rewrite constraints defined for the Rewrite trait and None otherwise.
 pub fn wrap_str(s: String, max_width: usize, shape: Shape) -> Option<String> {
-    if is_valid_str(&s, max_width, shape) {
+    if is_valid_str(&filter_normal_code(&s), max_width, shape) {
         Some(s)
     } else {
         None