fix to extract comments to stop internal error (#3857)

Author: topecongiro

src/comment.rs

@@ -976,6 +976,7 @@ fn left_trim_comment_line<'a>(line: &'a str, style: &CommentStyle<'_>) -> (&'a s
 
 pub(crate) trait FindUncommented {
     fn find_uncommented(&self, pat: &str) -> Option<usize>;
+    fn find_last_uncommented(&self, pat: &str) -> Option<usize>;
 }
 
 impl FindUncommented for str {
@@ -1001,6 +1002,19 @@ impl FindUncommented for str {
             None => Some(self.len() - pat.len()),
         }
     }
+
+    fn find_last_uncommented(&self, pat: &str) -> Option<usize> {
+        if let Some(left) = self.find_uncommented(pat) {
+            let mut result = left;
+            // add 1 to use find_last_uncommented for &str after pat
+            while let Some(next) = self[(result + 1)..].find_last_uncommented(pat) {
+                result += next + 1;
+            }
+            Some(result)
+        } else {
+            None
+        }
+    }
 }
 
 // Returns the first byte position after the first comment. The given string
@@ -1882,6 +1896,16 @@ mod test {
         check("\"/* abc", "abc", Some(4));
     }
 
+    #[test]
+    fn test_find_last_uncommented() {
+        fn check(haystack: &str, needle: &str, expected: Option<usize>) {
+            assert_eq!(expected, haystack.find_last_uncommented(needle));
+        }
+        check("foo test bar test", "test", Some(13));
+        check("test,", "test", Some(0));
+        check("nothing", "test", None);
+    }
+
     #[test]
     fn test_filter_normal_code() {
         let s = r#"

src/expr.rs

@@ -1881,14 +1881,13 @@ pub(crate) fn rewrite_assign_rhs<S: Into<String>, R: Rewrite>(
     rewrite_assign_rhs_with(context, lhs, ex, shape, RhsTactics::Default)
 }
 
-pub(crate) fn rewrite_assign_rhs_with<S: Into<String>, R: Rewrite>(
+pub(crate) fn rewrite_assign_rhs_expr<R: Rewrite>(
     context: &RewriteContext<'_>,
-    lhs: S,
+    lhs: &str,
     ex: &R,
     shape: Shape,
     rhs_tactics: RhsTactics,
 ) -> Option<String> {
-    let lhs = lhs.into();
     let last_line_width = last_line_width(&lhs).saturating_sub(if lhs.contains('\n') {
         shape.indent.width()
     } else {
@@ -1900,13 +1899,31 @@ pub(crate) fn rewrite_assign_rhs_with<S: Into<String>, R: Rewrite>(
         offset: shape.offset + last_line_width + 1,
         ..shape
     });
-    let rhs = choose_rhs(
+    let has_rhs_comment = if let Some(offset) = lhs.find_last_uncommented("=") {
+        lhs.trim_end().len() > offset + 1
+    } else {
+        false
+    };
+
+    choose_rhs(
         context,
         ex,
         orig_shape,
         ex.rewrite(context, orig_shape),
         rhs_tactics,
-    )?;
+        has_rhs_comment,
+    )
+}
+
+pub(crate) fn rewrite_assign_rhs_with<S: Into<String>, R: Rewrite>(
+    context: &RewriteContext<'_>,
+    lhs: S,
+    ex: &R,
+    shape: Shape,
+    rhs_tactics: RhsTactics,
+) -> Option<String> {
+    let lhs = lhs.into();
+    let rhs = rewrite_assign_rhs_expr(context, &lhs, ex, shape, rhs_tactics)?;
     Some(lhs + &rhs)
 }
 
@@ -1916,6 +1933,7 @@ fn choose_rhs<R: Rewrite>(
     shape: Shape,
     orig_rhs: Option<String>,
     rhs_tactics: RhsTactics,
+    has_rhs_comment: bool,
 ) -> Option<String> {
     match orig_rhs {
         Some(ref new_str)
@@ -1932,13 +1950,14 @@ fn choose_rhs<R: Rewrite>(
                 .indent
                 .block_indent(context.config)
                 .to_string_with_newline(context.config);
+            let before_space_str = if has_rhs_comment { "" } else { " " };
 
             match (orig_rhs, new_rhs) {
                 (Some(ref orig_rhs), Some(ref new_rhs))
                     if wrap_str(new_rhs.clone(), context.config.max_width(), new_shape)
                         .is_none() =>
                 {
-                    Some(format!(" {}", orig_rhs))
+                    Some(format!("{}{}", before_space_str, orig_rhs))
                 }
                 (Some(ref orig_rhs), Some(ref new_rhs))
                     if prefer_next_line(orig_rhs, new_rhs, rhs_tactics) =>
@@ -1948,10 +1967,11 @@ fn choose_rhs<R: Rewrite>(
                 (None, Some(ref new_rhs)) => Some(format!("{}{}", new_indent_str, new_rhs)),
                 (None, None) if rhs_tactics == RhsTactics::AllowOverflow => {
                     let shape = shape.infinite_width();
-                    expr.rewrite(context, shape).map(|s| format!(" {}", s))
+                    expr.rewrite(context, shape)
+                        .map(|s| format!("{}{}", before_space_str, s))
                 }
                 (None, None) => None,
-                (Some(orig_rhs), _) => Some(format!(" {}", orig_rhs)),
+                (Some(orig_rhs), _) => Some(format!("{}{}", before_space_str, orig_rhs)),
             }
         }
     }

src/items.rs

@@ -18,7 +18,8 @@ use crate::comment::{
 use crate::config::lists::*;
 use crate::config::{BraceStyle, Config, IndentStyle, Version};
 use crate::expr::{
-    is_empty_block, is_simple_block_stmt, rewrite_assign_rhs, rewrite_assign_rhs_with, RhsTactics,
+    is_empty_block, is_simple_block_stmt, rewrite_assign_rhs, rewrite_assign_rhs_expr,
+    rewrite_assign_rhs_with, RhsTactics,
 };
 use crate::lists::{definitive_tactic, itemize_list, write_list, ListFormatting, Separator};
 use crate::macros::{rewrite_macro, MacroPosition};
@@ -111,10 +112,67 @@ impl Rewrite for ast::Local {
         result.push_str(&infix);
 
         if let Some(ref ex) = self.init {
+            let base_span = if let Some(ref ty) = self.ty {
+                mk_sp(ty.span.hi(), self.span.hi())
+            } else {
+                mk_sp(self.pat.span.hi(), self.span.hi())
+            };
+
+            if let Some(offset) = context.snippet(base_span).find_uncommented("=") {
+                let base_span_lo = base_span.lo();
+
+                let assign_lo = base_span_lo + BytePos(offset as u32);
+                let comment_start_pos = if let Some(ref ty) = self.ty {
+                    ty.span.hi()
+                } else {
+                    self.pat.span.hi()
+                };
+                let comment_before_assign =
+                    context.snippet(mk_sp(comment_start_pos, assign_lo)).trim();
+
+                let assign_hi = base_span_lo + BytePos((offset + 1) as u32);
+                let rhs_span_lo = ex.span.lo();
+                let comment_end_pos = if ex.attrs.is_empty() {
+                    rhs_span_lo
+                } else {
+                    let attr_span_lo = ex.attrs.first().unwrap().span.lo();
+                    // for the case using block
+                    // ex. let x = { #![my_attr]do_something(); }
+                    if rhs_span_lo < attr_span_lo {
+                        rhs_span_lo
+                    } else {
+                        attr_span_lo
+                    }
+                };
+                let comment_after_assign =
+                    context.snippet(mk_sp(assign_hi, comment_end_pos)).trim();
+
+                if !comment_before_assign.is_empty() {
+                    let new_indent_str = &pat_shape
+                        .block_indent(0)
+                        .to_string_with_newline(context.config);
+                    result = format!("{}{}{}", comment_before_assign, new_indent_str, result);
+                }
+
+                if !comment_after_assign.is_empty() {
+                    let new_indent_str =
+                        &shape.block_indent(0).to_string_with_newline(context.config);
+                    result.push_str(new_indent_str);
+                    result.push_str(comment_after_assign);
+                    result.push_str(new_indent_str);
+                }
+            }
+
             // 1 = trailing semicolon;
             let nested_shape = shape.sub_width(1)?;
-
-            result = rewrite_assign_rhs(context, result, &**ex, nested_shape)?;
+            let rhs = rewrite_assign_rhs_expr(
+                context,
+                &result,
+                &**ex,
+                nested_shape,
+                RhsTactics::Default,
+            )?;
+            result = result + &rhs;
         }
 
         result.push(';');

tests/source/issue-3851.rs

@@ -0,0 +1,22 @@
+fn main() {
+    let before 
+      // before_comment 
+    = 
+    if true {
+     1.0
+    };
+
+    let after = 
+            // after_comment 
+    if true {
+     1.0
+    };
+
+    let both 
+      // before_comment 
+    = 
+            // after_comment 
+    if true {
+     1.0
+    };
+}

tests/target/issue-3851.rs

@@ -0,0 +1,19 @@
+fn main() {
+    // before_comment
+    let before = if true {
+        1.0
+    };
+
+    let after =
+    // after_comment
+    if true {
+        1.0
+    };
+
+    // before_comment
+    let both =
+    // after_comment
+    if true {
+        1.0
+    };
+}