Add #[rustfmt::sort] for enum variants

Author: Code0x58

Diff for: src/items.rs

@@ -1,12 +1,12 @@
 // Formatting top-level items - functions, structs, enums, traits, impls.
 
-use std::borrow::Cow;
-use std::cmp::{Ordering, max, min};
-
+use itertools::Itertools;
 use regex::Regex;
-use rustc_ast::visit;
+use rustc_ast::{AttrVec, visit};
 use rustc_ast::{ast, ptr};
 use rustc_span::{BytePos, DUMMY_SP, Span, symbol};
+use std::borrow::Cow;
+use std::cmp::{Ordering, max, min};
 
 use crate::attr::filter_inline_attrs;
 use crate::comment::{
@@ -536,6 +536,7 @@ impl<'a> FmtVisitor<'a> {
         enum_def: &ast::EnumDef,
         generics: &ast::Generics,
         span: Span,
+        attrs: &AttrVec,
     ) {
         let enum_header =
             format_header(&self.get_context(), "enum ", ident, vis, self.block_indent);
@@ -563,7 +564,7 @@ impl<'a> FmtVisitor<'a> {
 
         self.last_pos = body_start;
 
-        match self.format_variant_list(enum_def, body_start, span.hi()) {
+        match self.format_variant_list(enum_def, body_start, span.hi(), attrs) {
             Some(ref s) if enum_def.variants.is_empty() => self.push_str(s),
             rw => {
                 self.push_rewrite(mk_sp(body_start, span.hi()), rw);
@@ -578,6 +579,7 @@ impl<'a> FmtVisitor<'a> {
         enum_def: &ast::EnumDef,
         body_lo: BytePos,
         body_hi: BytePos,
+        attrs: &AttrVec,
     ) -> Option<String> {
         if enum_def.variants.is_empty() {
             let mut buffer = String::with_capacity(128);
@@ -615,7 +617,7 @@ impl<'a> FmtVisitor<'a> {
             .unwrap_or(&0);
 
         let itemize_list_with = |one_line_width: usize| {
-            itemize_list(
+            let iter = itemize_list(
                 self.snippet_provider,
                 enum_def.variants.iter(),
                 "}",
@@ -635,8 +637,16 @@ impl<'a> FmtVisitor<'a> {
                 body_lo,
                 body_hi,
                 false,
-            )
-            .collect()
+            );
+            if contains_sort(attrs) {
+                // sort the items by their name as this enum has the rustfmt::sort attr
+                iter.enumerate()
+                    .sorted_by_key(|&(i, _)| enum_def.variants[i].ident.name.as_str())
+                    .map(|(_, item)| item)
+                    .collect()
+            } else {
+                iter.collect()
+            }
         };
         let mut items: Vec<_> = itemize_list_with(self.config.struct_variant_width());
 

Diff for: src/skip.rs

@@ -85,6 +85,7 @@ impl SkipNameContext {
 
 static RUSTFMT: &str = "rustfmt";
 static SKIP: &str = "skip";
+static SORT: &str = "sort";
 
 /// Say if you're playing with `rustfmt`'s skip attribute
 pub(crate) fn is_skip_attr(segments: &[ast::PathSegment]) -> bool {
@@ -103,6 +104,16 @@ pub(crate) fn is_skip_attr(segments: &[ast::PathSegment]) -> bool {
     }
 }
 
+pub(crate) fn is_sort_attr(segments: &[ast::PathSegment]) -> bool {
+    if segments.len() < 2 || segments[0].ident.to_string() != RUSTFMT {
+        return false;
+    }
+    match segments.len() {
+        2 => segments[1].ident.to_string() == SORT,
+        _ => false,
+    }
+}
+
 fn get_skip_names(kind: &str, attrs: &[ast::Attribute]) -> Vec<String> {
     let mut skip_names = vec![];
     let path = format!("{RUSTFMT}::{SKIP}::{kind}");

Diff for: src/utils.rs

@@ -24,6 +24,11 @@ pub(crate) fn skip_annotation() -> Symbol {
     Symbol::intern("rustfmt::skip")
 }
 
+#[inline]
+pub(crate) fn sort_annotation() -> Symbol {
+    Symbol::intern("rustfmt::sort")
+}
+
 pub(crate) fn rewrite_ident<'a>(context: &'a RewriteContext<'_>, ident: symbol::Ident) -> &'a str {
     context.snippet(ident.span)
 }
@@ -271,6 +276,35 @@ pub(crate) fn contains_skip(attrs: &[Attribute]) -> bool {
         .any(|a| a.meta().map_or(false, |a| is_skip(&a)))
 }
 
+#[inline]
+pub(crate) fn contains_sort(attrs: &[Attribute]) -> bool {
+    attrs
+        .iter()
+        .any(|a| a.meta().map_or(false, |a| is_sort(&a)))
+}
+
+#[inline]
+fn is_sort(meta_item: &MetaItem) -> bool {
+    match meta_item.kind {
+        MetaItemKind::Word => {
+            let path_str = pprust::path_to_string(&meta_item.path);
+            path_str == sort_annotation().as_str()
+        }
+        MetaItemKind::List(ref l) => {
+            meta_item.has_name(sym::cfg_attr) && l.len() == 2 && crate::utils::is_sort_nested(&l[1])
+        }
+        _ => false,
+    }
+}
+
+#[inline]
+fn is_sort_nested(meta_item: &NestedMetaItem) -> bool {
+    match meta_item {
+        NestedMetaItem::MetaItem(ref mi) => crate::utils::is_sort(mi),
+        NestedMetaItem::Lit(_) => false,
+    }
+}
+
 #[inline]
 pub(crate) fn semicolon_for_expr(context: &RewriteContext<'_>, expr: &ast::Expr) -> bool {
     // Never try to insert semicolons on expressions when we're inside

Diff for: src/visitor.rs

@@ -18,7 +18,7 @@ use crate::modules::Module;
 use crate::parse::session::ParseSess;
 use crate::rewrite::{Rewrite, RewriteContext};
 use crate::shape::{Indent, Shape};
-use crate::skip::{SkipContext, is_skip_attr};
+use crate::skip::{SkipContext, is_skip_attr, is_sort_attr};
 use crate::source_map::{LineRangeUtils, SpanUtils};
 use crate::spanned::Spanned;
 use crate::stmt::Stmt;
@@ -515,7 +515,7 @@ impl<'b, 'a: 'b> FmtVisitor<'a> {
                 }
                 ast::ItemKind::Enum(ref def, ref generics) => {
                     self.format_missing_with_indent(source!(self, item.span).lo());
-                    self.visit_enum(item.ident, &item.vis, def, generics, item.span);
+                    self.visit_enum(item.ident, &item.vis, def, generics, item.span, &item.attrs);
                     self.last_pos = source!(self, item.span).hi();
                 }
                 ast::ItemKind::Mod(safety, ref mod_kind) => {
@@ -858,7 +858,7 @@ impl<'b, 'a: 'b> FmtVisitor<'a> {
         if segments[0].ident.to_string() != "rustfmt" {
             return false;
         }
-        !is_skip_attr(segments)
+        !(is_skip_attr(segments) | is_sort_attr(segments))
     }
 
     fn walk_mod_items(&mut self, items: &[rustc_ast::ptr::P<ast::Item>]) {

Diff for: tests/source/enum.rs

@@ -210,3 +210,21 @@ pub enum E {
     A { a: u32 } = 0x100,
     B { field1: u32, field2: u8, field3: m::M } = 0x300 // comment
 }
+
+// #3422
+#[rustfmt::sort]
+enum SortE {
+
+    Y, // whitespace from above dropped
+
+    X, // whitespace from above dropped
+    E,
+    // something
+    D(),
+    C(),
+    /// Comment for B
+    B,
+    /// Comment for A
+    #[rustfmt::skip]
+    A,
+}

Diff for: tests/target/enum.rs

@@ -287,3 +287,19 @@ pub enum E {
         field3: m::M,
     } = 0x300, // comment
 }
+
+// #3422
+#[rustfmt::sort]
+enum SortE {
+    /// Comment for A
+    #[rustfmt::skip]
+    A,
+    /// Comment for B
+    B,
+    C(),
+    // something
+    D(),
+    E,
+    X, // whitespace from above dropped
+    Y, // whitespace from above dropped
+}