From 0f640b06ddac1ff6033fba2836f3ad0e71390db7 Mon Sep 17 00:00:00 2001
From: Marcus Klaas <mail@marcusklaas.nl>
Date: Sat, 1 Aug 2015 14:22:31 +0200
Subject: [PATCH] Properly format unsafe blocks

---
 src/expr.rs          | 31 ++++++++++++++++++++++++++++---
 src/visitor.rs       |  2 +-
 tests/source/expr.rs | 18 ++++++++++++++++++
 tests/target/expr.rs | 19 +++++++++++++++++++
 4 files changed, 66 insertions(+), 4 deletions(-)

diff --git a/src/expr.rs b/src/expr.rs
index 930f19af453..b8c37dcab4c 100644
--- a/src/expr.rs
+++ b/src/expr.rs
@@ -15,6 +15,7 @@ use StructLitStyle;
 use utils::{span_after, make_indent};
 use visitor::FmtVisitor;
 use config::BlockIndentStyle;
+use comment::{FindUncommented, rewrite_comment};
 
 use syntax::{ast, ptr};
 use syntax::codemap::{Pos, Span, BytePos, mk_sp};
@@ -104,11 +105,35 @@ impl Rewrite for ast::Expr {
 }
 
 impl Rewrite for ast::Block {
-    fn rewrite(&self, context: &RewriteContext, _: usize, _: usize) -> Option<String> {
+    fn rewrite(&self, context: &RewriteContext, width: usize, offset: usize) -> Option<String> {
         let mut visitor = FmtVisitor::from_codemap(context.codemap, context.config);
-        visitor.last_pos = self.span.lo;
         visitor.block_indent = context.block_indent;
 
+        let prefix = match self.rules {
+            ast::BlockCheckMode::PushUnsafeBlock(..) |
+            ast::BlockCheckMode::UnsafeBlock(..) => {
+                let snippet = try_opt!(context.codemap.span_to_snippet(self.span).ok());
+                let open_pos = try_opt!(snippet.find_uncommented("{"));
+                visitor.last_pos = self.span.lo + BytePos(open_pos as u32);
+
+                // Extract comment between unsafe and block start.
+                let trimmed = &snippet[6..open_pos].trim();
+
+                if trimmed.len() > 0 {
+                    // 9 = "unsafe  {".len(), 7 = "unsafe ".len()
+                    format!("unsafe {} ", rewrite_comment(trimmed, true, width - 9, offset + 7))
+                } else {
+                    "unsafe ".to_owned()
+                }
+            }
+            ast::BlockCheckMode::PopUnsafeBlock(..) |
+            ast::BlockCheckMode::DefaultBlock => {
+                visitor.last_pos = self.span.lo;
+
+                String::new()
+            }
+        };
+
         visitor.visit_block(self);
 
         // Push text between last block item and end of block
@@ -119,7 +144,7 @@ impl Rewrite for ast::Block {
         let file_name = context.codemap.span_to_filename(self.span);
         let string_buffer = visitor.changes.get(&file_name);
 
-        Some(string_buffer.to_string())
+        Some(format!("{}{}", prefix, string_buffer))
     }
 }
 
diff --git a/src/visitor.rs b/src/visitor.rs
index 0f41b94ac4b..f4ccfdd5c25 100644
--- a/src/visitor.rs
+++ b/src/visitor.rs
@@ -73,7 +73,6 @@ impl<'a, 'v> visit::Visitor<'v> for FmtVisitor<'a> {
         debug!("visit_block: {:?} {:?}",
                self.codemap.lookup_char_pos(b.span.lo),
                self.codemap.lookup_char_pos(b.span.hi));
-        self.format_missing(b.span.lo);
 
         self.changes.push_str_span(b.span, "{");
         self.last_pos = self.last_pos + BytePos(1);
@@ -82,6 +81,7 @@ impl<'a, 'v> visit::Visitor<'v> for FmtVisitor<'a> {
         for stmt in &b.stmts {
             self.visit_stmt(&stmt)
         }
+
         match b.expr {
             Some(ref e) => {
                 self.format_missing_with_indent(e.span.lo);
diff --git a/tests/source/expr.rs b/tests/source/expr.rs
index bfa1edc5311..e7862b83e50 100644
--- a/tests/source/expr.rs
+++ b/tests/source/expr.rs
@@ -73,3 +73,21 @@ fn bar() {
              aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa,
              a);
 }
+
+fn baz() {
+    unsafe    /*    {}{}{}{{{{}}   */   {
+        let foo = 1u32;
+    }
+
+    unsafe /* very looooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooong comment */ {}
+
+    unsafe // So this is a very long comment.
+           // Multi-line, too.
+           // Will it still format correctly?
+    {
+    }
+
+    unsafe {
+        // Regular unsafe block
+    }
+}
diff --git a/tests/target/expr.rs b/tests/target/expr.rs
index b76bd408ac9..a1875a0ec26 100644
--- a/tests/target/expr.rs
+++ b/tests/target/expr.rs
@@ -102,3 +102,22 @@ fn bar() {
     let x = (aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa && aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa,
              a);
 }
+
+fn baz() {
+    unsafe /* {}{}{}{{{{}} */ {
+        let foo = 1u32;
+    }
+
+    unsafe /* very looooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooong
+            * comment */ {
+    }
+
+    unsafe /* So this is a very long comment.
+            * Multi-line, too.
+            * Will it still format correctly? */ {
+    }
+
+    unsafe {
+        // Regular unsafe block
+    }
+}