Merge pull request #89 from marcusklaas/enums

Implement basic enum formatting

src/config.rs

@@ -22,6 +22,7 @@ pub struct Config {
     pub fn_args_paren_newline: bool,
     pub struct_trailing_comma: bool,
     pub struct_lit_trailing_comma: ::lists::SeparatorTactic,
+    pub enum_trailing_comma: bool,
 }
 
 impl Config {

src/default.toml

@@ -8,3 +8,4 @@ fn_return_indent = "WithArgs"
 fn_args_paren_newline = true
 struct_trailing_comma = true
 struct_lit_trailing_comma = "Vertical"
+enum_trailing_comma = true

src/expr.rs

@@ -264,7 +264,6 @@ impl<'a> FmtVisitor<'a> {
             _ => {}
         }
 
-        let result = self.snippet(expr.span);
-        result
+        self.snippet(expr.span)
     }
 }

src/imports.rs

@@ -10,6 +10,7 @@
 
 use visitor::FmtVisitor;
 use lists::{write_list, ListFormatting, SeparatorTactic, ListTactic};
+use utils::format_visibility;
 
 use syntax::ast;
 use syntax::parse::token;
@@ -48,11 +49,7 @@ impl<'a> FmtVisitor<'a> {
                             path_list: &[ast::PathListItem],
                             visibility: ast::Visibility) -> String {
         let path_str = pprust::path_to_string(path);
-
-        let vis = match visibility {
-            ast::Public => "pub ",
-            _ => ""
-        };
+        let vis = format_visibility(visibility);
 
         if path_list.len() == 1 {
             return rewrite_single_use_list(path_str, path_list[0], vis);

src/items.rs

@@ -11,7 +11,7 @@
 // Formatting top-level items - functions, structs, enums, traits, impls.
 
 use {ReturnIndent, BraceStyle};
-use utils::make_indent;
+use utils::{format_visibility, make_indent};
 use lists::{write_list, ListFormatting, SeparatorTactic, ListTactic};
 use visitor::FmtVisitor;
 use syntax::{ast, abi};
@@ -109,9 +109,8 @@ impl<'a> FmtVisitor<'a> {
 
         let mut result = String::with_capacity(1024);
         // Vis unsafety abi.
-        if vis == ast::Visibility::Public {
-            result.push_str("pub ");
-        }
+        result.push_str(format_visibility(vis));
+
         if let &ast::Unsafety::Unsafe = unsafety {
             result.push_str("unsafe ");
         }
@@ -351,7 +350,7 @@ impl<'a> FmtVisitor<'a> {
     }
 
     fn compute_budgets_for_args(&self,
-                                result: &String,
+                                result: &str,
                                 indent: usize,
                                 ret_str_len: usize,
                                 newline_brace: bool)
@@ -408,6 +407,118 @@ impl<'a> FmtVisitor<'a> {
         }
     }
 
+    pub fn visit_enum(&mut self,
+                      ident: ast::Ident,
+                      vis: ast::Visibility,
+                      enum_def: &ast::EnumDef,
+                      generics: &ast::Generics,
+                      span: Span)
+    {
+        let header_str = self.format_header("enum", ident, vis);
+        self.changes.push_str_span(span, &header_str);
+
+        let enum_snippet = self.snippet(span);
+        // FIXME this will give incorrect results if there is a { in a comment.
+        let body_start = span.lo + BytePos(enum_snippet.find('{').unwrap() as u32 + 1);
+        let generics_str = self.format_generics(generics, body_start);
+        self.changes.push_str_span(span, &generics_str);
+
+        self.last_pos = body_start;
+        self.block_indent += config!(tab_spaces);
+        for (i, f) in enum_def.variants.iter().enumerate() {
+            let next_span_start: BytePos = if i == enum_def.variants.len() - 1 {
+                span.hi
+            } else {
+                enum_def.variants[i + 1].span.lo
+            };
+
+            self.visit_variant(f, i == enum_def.variants.len() - 1, next_span_start);
+        }
+        self.block_indent -= config!(tab_spaces);
+
+        self.format_missing_with_indent(span.lo + BytePos(enum_snippet.rfind('}').unwrap() as u32));
+        self.changes.push_str_span(span, "}");
+    }
+
+    // Variant of an enum
+    fn visit_variant(&mut self,
+                     field: &ast::Variant,
+                     last_field: bool,
+                     next_span_start: BytePos)
+    {
+        if self.visit_attrs(&field.node.attrs) {
+            return;
+        }
+
+        if let ast::VariantKind::TupleVariantKind(ref types) = field.node.kind {
+            self.format_missing_with_indent(field.span.lo);
+
+            let vis = format_visibility(field.node.vis);
+            self.changes.push_str_span(field.span, vis);
+            let name = field.node.name.to_string();
+            self.changes.push_str_span(field.span, &name);
+
+            let mut result = String::new();
+
+            if types.len() > 0 {
+                let comments = self.make_comments_for_list(Vec::new(),
+                                                           types.iter().map(|arg| arg.ty.span),
+                                                           ",",
+                                                           ")",
+                                                           |span| span.lo,
+                                                           |span| span.hi,
+                                                           next_span_start);
+
+                let type_strings: Vec<_> = types.iter()
+                                                .map(|arg| pprust::ty_to_string(&arg.ty))
+                                                .zip(comments.into_iter())
+                                                .collect();
+
+                result.push('(');
+
+                let indent = self.block_indent
+                             + vis.len()
+                             + field.node.name.to_string().len()
+                             + 1; // 1 = (
+
+                let comma_cost = if config!(enum_trailing_comma) { 1 } else { 0 };
+                let budget = config!(ideal_width) - indent - comma_cost - 1; // 1 = )
+
+                let fmt = ListFormatting {
+                    tactic: ListTactic::HorizontalVertical,
+                    separator: ",",
+                    trailing_separator: SeparatorTactic::Never,
+                    indent: indent,
+                    h_width: budget,
+                    v_width: budget,
+                };
+                result.push_str(&write_list(&type_strings, &fmt));
+                result.push(')');
+            }
+
+            if let Some(ref expr) = field.node.disr_expr {
+                result.push_str(" = ");
+                let expr_snippet = self.snippet(expr.span);
+                result.push_str(&expr_snippet);
+
+                // Make sure we do not exceed column limit
+                // 4 = " = ,"
+                assert!(config!(max_width) >= vis.len() + name.len() + expr_snippet.len() + 4,
+                        "Enum variant exceeded column limit");
+            }
+
+            self.changes.push_str_span(field.span, &result);
+
+            if !last_field || config!(enum_trailing_comma) {
+                self.changes.push_str_span(field.span, ",");
+            }
+        }
+
+        // TODO: deal with struct-like variants
+
+        self.last_pos = field.span.hi + BytePos(1);
+    }
+
     pub fn visit_struct(&mut self,
                         ident: ast::Ident,
                         vis: ast::Visibility,
@@ -415,7 +526,7 @@ impl<'a> FmtVisitor<'a> {
                         generics: &ast::Generics,
                         span: Span)
     {
-        let header_str = self.struct_header(ident, vis);
+        let header_str = self.format_header("struct", ident, vis);
         self.changes.push_str_span(span, &header_str);
 
         if struct_def.fields.len() == 0 {
@@ -428,24 +539,11 @@ impl<'a> FmtVisitor<'a> {
             return;
         }
 
-        let mut generics_buf = String::new();
-        let generics_str = self.rewrite_generics(generics, self.block_indent, struct_def.fields[0].span.lo);
-        generics_buf.push_str(&generics_str);
-
-        if generics.where_clause.predicates.len() > 0 {
-            generics_buf.push_str(&self.rewrite_where_clause(&generics.where_clause,
-                                                             self.block_indent,
-                                                             struct_def.fields[0].span.lo));
-            generics_buf.push_str(&make_indent(self.block_indent));
-            generics_buf.push_str("\n{");
-
-        } else {
-            generics_buf.push_str(" {");
-        }
-        self.changes.push_str_span(span, &generics_buf);
+        let generics_str = self.format_generics(generics, struct_def.fields[0].span.lo);
+        self.changes.push_str_span(span, &generics_str);
 
         let struct_snippet = self.snippet(span);
-        // FIXME this will give incorrect results if there is a { in a commet.
+        // FIXME this will give incorrect results if there is a { in a comment.
         self.last_pos = span.lo + BytePos(struct_snippet.find('{').unwrap() as u32 + 1);
 
         self.block_indent += config!(tab_spaces);
@@ -458,18 +556,34 @@ impl<'a> FmtVisitor<'a> {
         self.changes.push_str_span(span, "}");
     }
 
-    fn struct_header(&self,
+    fn format_header(&self,
+                     item_name: &str,
                      ident: ast::Ident,
                      vis: ast::Visibility)
         -> String
     {
-        let vis = if vis == ast::Visibility::Public {
-            "pub "
-        } else {
-            ""
-        };
+        format!("{}{} {}", format_visibility(vis), item_name, &token::get_ident(ident))
+    }
 
-        format!("{}struct {}", vis, &token::get_ident(ident))
+    fn format_generics(&self,
+                       generics: &ast::Generics,
+                       span_end: BytePos)
+        -> String
+    {
+        let mut result = self.rewrite_generics(generics, self.block_indent, span_end);
+
+        if generics.where_clause.predicates.len() > 0 {
+            result.push_str(&self.rewrite_where_clause(&generics.where_clause,
+                                                             self.block_indent,
+                                                             span_end));
+            result.push_str(&make_indent(self.block_indent));
+            result.push_str("\n{");
+
+        } else {
+            result.push_str(" {");
+        }
+
+        result
     }
 
     // Field of a struct
@@ -491,11 +605,7 @@ impl<'a> FmtVisitor<'a> {
         };
         let vis = match field.node.kind {
             ast::StructFieldKind::NamedField(_, vis) |
-            ast::StructFieldKind::UnnamedField(vis) => if vis == ast::Visibility::Public {
-                "pub "
-            } else {
-                ""
-            }
+            ast::StructFieldKind::UnnamedField(vis) => format_visibility(vis)
         };
         let typ = pprust::ty_to_string(&field.node.ty);
 

src/utils.rs

@@ -8,6 +8,7 @@
 // option. This file may not be copied, modified, or distributed
 // except according to those terms.
 
+use syntax::ast::Visibility;
 
 #[inline]
 pub fn prev_char(s: &str, mut i: usize) -> usize {
@@ -38,3 +39,11 @@ pub fn make_indent(width: usize) -> String {
     }
     indent
 }
+
+#[inline]
+pub fn format_visibility(vis: Visibility) -> &'static str {
+    match vis {
+        Visibility::Public => "pub ",
+        Visibility::Inherited => ""
+    }
+}

src/visitor.rs

@@ -196,6 +196,15 @@ impl<'a, 'v> visit::Visitor<'v> for FmtVisitor<'a> {
                                   item.span);
                 self.last_pos = item.span.hi;
             }
+            ast::Item_::ItemEnum(ref def, ref generics) => {
+                self.format_missing_with_indent(item.span.lo);
+                self.visit_enum(item.ident,
+                                item.vis,
+                                def,
+                                generics,
+                                item.span);
+                self.last_pos = item.span.hi;
+            }
             _ => {
                 visit::walk_item(self, item);
             }

tests/idem/enum.rs

@@ -0,0 +1,34 @@
+// Enums test
+
+#[atrr]
+pub enum Test {
+    A,
+    B(u32, A /* comment */),
+    /// Doc comment
+    C,
+}
+
+pub enum Foo<'a, Y: Baz>
+    where X: Whatever
+{
+    A,
+}
+
+enum EmtpyWithComment {
+    // Some comment
+}
+
+// C-style enum
+enum Bar {
+    A = 1,
+    #[someAttr(test)]
+    B = 2, // comment
+    C,
+}
+
+enum LongVariants {
+    First(LOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOONG, // small comment
+          VARIANT),
+    // This is the second variant
+    Second,
+}