Format code block in comment

Closes rust-lang#554.
Closes rust-lang#1695.

src/comment.rs

@@ -318,41 +318,65 @@ fn rewrite_comment_inner(
 
     let mut result = String::with_capacity(orig.len() * 2);
     result.push_str(opener);
+    let mut code_block_buffer = String::with_capacity(128);
     let mut is_prev_line_multi_line = false;
     let mut inside_code_block = false;
     let comment_line_separator = format!("\n{}{}", indent_str, line_start);
+    let join_code_block_with_comment_line_separator = |s: &str| {
+        let mut result = String::with_capacity(s.len() + 128);
+        let mut iter = s.lines().peekable();
+        while let Some(line) = iter.next() {
+            result.push_str(line);
+            result.push_str(match iter.peek() {
+                Some(ref next_line) if next_line.is_empty() => comment_line_separator.trim_right(),
+                Some(..) => &comment_line_separator,
+                None => "",
+            });
+        }
+        result
+    };
+
     for (i, (line, has_leading_whitespace)) in lines.enumerate() {
         let is_last = i == count_newlines(orig);
-        if result == opener {
-            let force_leading_whitespace = opener == "/* " && count_newlines(orig) == 0;
-            if !has_leading_whitespace && !force_leading_whitespace && result.ends_with(' ') {
-                result.pop();
-            }
-            if line.is_empty() {
-                continue;
-            }
-        } else if is_prev_line_multi_line && !line.is_empty() {
-            result.push(' ')
-        } else if is_last && !closer.is_empty() && line.is_empty() {
-            result.push('\n');
-            result.push_str(&indent_str);
-        } else {
-            result.push_str(&comment_line_separator);
-            if !has_leading_whitespace && result.ends_with(' ') {
-                result.pop();
-            }
-        }
 
-        if line.starts_with("```") {
-            inside_code_block = !inside_code_block;
-        }
         if inside_code_block {
-            if line.is_empty() && result.ends_with(' ') {
-                result.pop();
-            } else {
+            if line.starts_with("```") {
+                inside_code_block = false;
+                result.push_str(&comment_line_separator);
+                let code_block = ::format_code_block(&code_block_buffer, config)
+                    .unwrap_or_else(|| code_block_buffer.to_owned());
+                result.push_str(&join_code_block_with_comment_line_separator(&code_block));
+                code_block_buffer.clear();
+                result.push_str(&comment_line_separator);
                 result.push_str(line);
+            } else {
+                code_block_buffer.push_str(line);
+                code_block_buffer.push('\n');
             }
+
             continue;
+        } else {
+            inside_code_block = line.starts_with("```");
+
+            if result == opener {
+                let force_leading_whitespace = opener == "/* " && count_newlines(orig) == 0;
+                if !has_leading_whitespace && !force_leading_whitespace && result.ends_with(' ') {
+                    result.pop();
+                }
+                if line.is_empty() {
+                    continue;
+                }
+            } else if is_prev_line_multi_line && !line.is_empty() {
+                result.push(' ')
+            } else if is_last && !closer.is_empty() && line.is_empty() {
+                result.push('\n');
+                result.push_str(&indent_str);
+            } else {
+                result.push_str(&comment_line_separator);
+                if !has_leading_whitespace && result.ends_with(' ') {
+                    result.pop();
+                }
+            }
         }
 
         if config.wrap_comments() && line.len() > fmt.shape.width && !has_url(line) {

src/lib.rs

@@ -44,6 +44,7 @@ use comment::{CharClasses, FullCodeCharKind};
 use config::Config;
 use filemap::FileMap;
 use issues::{BadIssueSeeker, Issue};
+use shape::Indent;
 use utils::use_colored_tty;
 use visitor::{FmtVisitor, SnippetProvider};
 
@@ -529,6 +530,55 @@ fn parse_input(
     }
 }
 
+/// Format the given snippet. The snippet is expected to be *complete* code.
+/// When we cannot parse the given snippet, this function returns `None`.
+pub fn format_snippet(snippet: &str, config: &Config) -> Option<String> {
+    let mut out: Vec<u8> = Vec::with_capacity(snippet.len() * 2);
+    let input = Input::Text(snippet.into());
+    let mut config = config.clone();
+    config.set().write_mode(config::WriteMode::Plain);
+    match format_input(input, &config, Some(&mut out)) {
+        // `format_input()` returns an empty string on parsing error.
+        Ok(..) if out.is_empty() && !snippet.is_empty() => None,
+        Ok(..) => String::from_utf8(out).ok(),
+        Err(..) => None,
+    }
+}
+
+/// Format the given code block. Mainly targeted for code block in comment.
+/// The code block may be incomplete (i.e. parser may be unable to parse it).
+/// To avoid panic in parser, we wrap the code block with a dummy function.
+/// The returned code block does *not* end with newline.
+pub fn format_code_block(code_snippet: &str, config: &Config) -> Option<String> {
+    // Wrap the given code block with `fn main()` if it does not have one.
+    let fn_main_prefix = "fn main() {\n";
+    let snippet = fn_main_prefix.to_owned() + code_snippet + "\n}";
+
+    // Trim "fn main() {" on the first line and "}" on the last line,
+    // then unindent the whole code block.
+    format_snippet(&snippet, config).map(|s| {
+        // 2 = "}\n"
+        s[fn_main_prefix.len()..s.len().checked_sub(2).unwrap_or(0)]
+            .lines()
+            .map(|line| {
+                if line.len() > config.tab_spaces() {
+                    // Make sure that the line has leading whitespaces.
+                    let indent_str =
+                        Indent::from_width(config, config.tab_spaces()).to_string(config);
+                    if line.starts_with(indent_str.as_ref()) {
+                        &line[config.tab_spaces()..]
+                    } else {
+                        line
+                    }
+                } else {
+                    line
+                }
+            })
+            .collect::<Vec<_>>()
+            .join("\n")
+    })
+}
+
 pub fn format_input<T: Write>(
     input: Input,
     config: &Config,
@@ -650,3 +700,86 @@ pub fn run(input: Input, config: &Config) -> Summary {
         }
     }
 }
+
+#[cfg(test)]
+mod test {
+    use super::{format_code_block, format_snippet, Config};
+
+    #[test]
+    fn test_no_panic_on_format_snippet_and_format_code_block() {
+        // `format_snippet()` and `format_code_block()` should not panic
+        // even when we cannot parse the given snippet.
+        let snippet = "let";
+        assert!(format_snippet(snippet, &Config::default()).is_none());
+        assert!(format_code_block(snippet, &Config::default()).is_none());
+    }
+
+    fn test_format_inner<F>(formatter: F, input: &str, expected: &str) -> bool
+    where
+        F: Fn(&str, &Config) -> Option<String>,
+    {
+        let output = formatter(input, &Config::default());
+        output.is_some() && output.unwrap() == expected
+    }
+
+    #[test]
+    fn test_format_snippet() {
+        let snippet = "fn main() { println!(\"hello, world\"); }";
+        let expected = "fn main() {\n    \
+                        println!(\"hello, world\");\n\
+                        }\n";
+        assert!(test_format_inner(format_snippet, snippet, expected));
+    }
+
+    #[test]
+    fn test_format_code_block() {
+        // simple code block
+        let code_block = "let x=3;";
+        let expected = "let x = 3;";
+        assert!(test_format_inner(format_code_block, code_block, expected));
+
+        // more complex code block, taken from chains.rs.
+        let code_block =
+"let (nested_shape, extend) = if !parent_rewrite_contains_newline && is_continuable(&parent) {
+(
+chain_indent(context, shape.add_offset(parent_rewrite.len())),
+context.config.indent_style() == IndentStyle::Visual || is_small_parent,
+)
+} else if is_block_expr(context, &parent, &parent_rewrite) {
+match context.config.indent_style() {
+// Try to put the first child on the same line with parent's last line
+IndentStyle::Block => (parent_shape.block_indent(context.config.tab_spaces()), true),
+// The parent is a block, so align the rest of the chain with the closing
+// brace.
+IndentStyle::Visual => (parent_shape, false),
+}
+} else {
+(
+chain_indent(context, shape.add_offset(parent_rewrite.len())),
+false,
+)
+};
+";
+        let expected =
+"let (nested_shape, extend) = if !parent_rewrite_contains_newline && is_continuable(&parent) {
+    (
+        chain_indent(context, shape.add_offset(parent_rewrite.len())),
+        context.config.indent_style() == IndentStyle::Visual || is_small_parent,
+    )
+} else if is_block_expr(context, &parent, &parent_rewrite) {
+    match context.config.indent_style() {
+        // Try to put the first child on the same line with parent's last line
+        IndentStyle::Block => (parent_shape.block_indent(context.config.tab_spaces()), true),
+        // The parent is a block, so align the rest of the chain with the closing
+        // brace.
+        IndentStyle::Visual => (parent_shape, false),
+    }
+} else {
+    (
+        chain_indent(context, shape.add_offset(parent_rewrite.len())),
+        false,
+    )
+};";
+        assert!(test_format_inner(format_code_block, code_block, expected));
+    }
+}

tests/target/issue-2197.rs

@@ -4,8 +4,14 @@
 
 /// ```rust
 /// unsafe fn sum_sse2(x: i32x4) -> i32 {
-///     let x = vendor::_mm_add_epi32(x, vendor::_mm_srli_si128(x.into(), 8).into());
-///     let x = vendor::_mm_add_epi32(x, vendor::_mm_srli_si128(x.into(), 4).into());
+///     let x = vendor::_mm_add_epi32(
+///         x,
+///         vendor::_mm_srli_si128(x.into(), 8).into(),
+///     );
+///     let x = vendor::_mm_add_epi32(
+///         x,
+///         vendor::_mm_srli_si128(x.into(), 4).into(),
+///     );
 ///     vendor::_mm_cvtsi128_si32(x)
 /// }
 /// ```