Format loops

6 files changed, 67 insertions, 2 deletions

diff --git a/src/changes.rs b/src/changes.rs
index 3f0d985da48..4cce45e5eb3 100644
--- a/src/changes.rs
+++ b/src/changes.rs
@@ -99,6 +99,14 @@ impl<'a> ChangeSet<'a> {
         self.push_str(&file_name, text)
     }
 
+    // Fetch the output buffer for the given file name.
+    // Panics on unknown files.
+    pub fn get(&mut self, file_name: &str) -> &StringBuffer {
+        self.file_map.get(file_name).unwrap()
+    }
+
+    // Fetch a mutable reference to the output buffer for the given file name.
+    // Panics on unknown files.
     pub fn get_mut(&mut self, file_name: &str) -> &mut StringBuffer {
         self.file_map.get_mut(file_name).unwrap()
     }
diff --git a/src/expr.rs b/src/expr.rs
index da706c4afa7..e87020f4ae1 100644
--- a/src/expr.rs
+++ b/src/expr.rs
@@ -12,11 +12,13 @@ use rewrite::{Rewrite, RewriteContext};
 use lists::{write_list, itemize_list, ListFormatting, SeparatorTactic, ListTactic};
 use string::{StringFormat, rewrite_string};
 use utils::{span_after, make_indent};
+use visitor::FmtVisitor;
 
 use syntax::{ast, ptr};
-use syntax::codemap::{Pos, Span, BytePos};
+use syntax::codemap::{Pos, Span, BytePos, mk_sp};
 use syntax::parse::token;
 use syntax::print::pprust;
+use syntax::visit::Visitor;
 
 impl Rewrite for ast::Expr {
     fn rewrite(&self, context: &RewriteContext, width: usize, offset: usize) -> Option<String> {
@@ -53,11 +55,38 @@ impl Rewrite for ast::Expr {
             ast::Expr_::ExprTup(ref items) => {
                 rewrite_tuple_lit(context, items, self.span, width, offset)
             }
+            ast::Expr_::ExprLoop(ref block, _) => {
+                // FIXME: this drops any comment between "loop" and the block.
+                // TODO: format label
+                block.rewrite(context, width, offset).map(|result| {
+                    format!("loop {}", result)
+                })
+            }
             _ => context.codemap.span_to_snippet(self.span).ok()
         }
     }
 }
 
+impl Rewrite for ast::Block {
+    fn rewrite(&self, context: &RewriteContext, _: usize, _: 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;
+
+        visitor.visit_block(self);
+
+        // Push text between last block item and end of block
+        let snippet = visitor.snippet(mk_sp(visitor.last_pos, self.span.hi));
+        visitor.changes.push_str_span(self.span, &snippet);
+
+        // Stringify visitor
+        let file_name = context.codemap.span_to_filename(self.span);
+        let string_buffer = visitor.changes.get(&file_name);
+
+        Some(string_buffer.to_string())
+    }
+}
+
 fn rewrite_string_lit(context: &RewriteContext,
                       s: &str,
                       span: Span,
diff --git a/src/rewrite.rs b/src/rewrite.rs
index 78d611253dd..33e4d5fbbf3 100644
--- a/src/rewrite.rs
+++ b/src/rewrite.rs
@@ -28,4 +28,5 @@ pub trait Rewrite {
 pub struct RewriteContext<'a> {
     pub codemap: &'a CodeMap,
     pub config: &'a Config,
+    pub block_indent: usize,
 }
diff --git a/src/visitor.rs b/src/visitor.rs
index 5c15e1bf523..1b1fd9d17c6 100644
--- a/src/visitor.rs
+++ b/src/visitor.rs
@@ -38,7 +38,9 @@ impl<'a, 'v> visit::Visitor<'v> for FmtVisitor<'a> {
                self.codemap.lookup_char_pos(ex.span.hi));
         self.format_missing(ex.span.lo);
         let offset = self.changes.cur_offset_span(ex.span);
-        let context = RewriteContext { codemap: self.codemap, config: self.config };
+        let context = RewriteContext { codemap: self.codemap,
+                                       config: self.config,
+                                       block_indent: self.block_indent, };
         let rewrite = ex.rewrite(&context, self.config.max_width - offset, offset);
 
         if let Some(new_str) = rewrite {
diff --git a/tests/source/loop.rs b/tests/source/loop.rs
new file mode 100644
index 00000000000..ad5ef93521d
--- /dev/null
+++ b/tests/source/loop.rs
@@ -0,0 +1,11 @@
+
+fn main() {
+    loop    
+    {   return some_val;}
+
+let x = loop { do_forever(); };
+
+         loop {
+        // Just comments
+    }
+}
diff --git a/tests/target/loop.rs b/tests/target/loop.rs
new file mode 100644
index 00000000000..3eb506eb6ed
--- /dev/null
+++ b/tests/target/loop.rs
@@ -0,0 +1,14 @@
+
+fn main() {
+    loop {
+        return some_val;
+    }
+
+    let x = loop {
+        do_forever();
+    };
+
+    loop {
+        // Just comments
+    }
+}