@@ -522,21 +522,22 @@ impl<'a> FmtVisitor<'a> {
522522 self . push_rewrite ( struct_parts. span , rewrite) ;
523523 }
524524
525- pub ( crate ) fn visit_enum (
525+ // TODO(ding-young) do I need to make it a separate function instead of a method of FmtVisitor?
526+ fn format_enum (
526527 & mut self ,
527528 ident : symbol:: Ident ,
528529 vis : & ast:: Visibility ,
529530 enum_def : & ast:: EnumDef ,
530531 generics : & ast:: Generics ,
531532 span : Span ,
532- ) {
533+ ) -> Option < String > {
533534 let enum_header =
534535 format_header ( & self . get_context ( ) , "enum " , ident, vis, self . block_indent ) ;
535- self . push_str ( & enum_header) ;
536536
537537 let enum_snippet = self . snippet ( span) ;
538538 let brace_pos = enum_snippet. find_uncommented ( "{" ) . unwrap ( ) ;
539539 let body_start = span. lo ( ) + BytePos ( brace_pos as u32 + 1 ) ;
540+ // TODO(ding-young) what if there is no generic?
540541 let generics_str = format_generics (
541542 & self . get_context ( ) ,
542543 generics,
@@ -550,25 +551,36 @@ impl<'a> FmtVisitor<'a> {
550551 // make a span that starts right after `enum Foo`
551552 mk_sp ( ident. span . hi ( ) , body_start) ,
552553 last_line_width ( & enum_header) ,
553- ) ;
554-
555- if let Some ( generics_str) = generics_str {
556- self . push_str ( & generics_str) ;
557- } else {
558- self . push_str ( self . snippet ( mk_sp ( ident. span . hi ( ) , body_start) ) ) ;
559- }
560-
561- self . last_pos = body_start;
554+ ) ?;
562555
563556 match self . format_variant_list ( enum_def, body_start, span. hi ( ) ) {
564- Some ( ref s) if enum_def. variants . is_empty ( ) => self . push_str ( s) ,
565- rw => {
566- self . push_rewrite ( mk_sp ( body_start, span. hi ( ) ) , rw) ;
557+ Some ( ref s) if enum_def. variants . is_empty ( ) => {
558+ Some ( format ! ( "{enum_header}{generics_str}{s}" ) )
559+ }
560+ Some ( rw) => {
561+ let indent = self . block_indent . to_string ( self . config ) ;
562+ self . block_indent = self . block_indent . block_unindent ( self . config ) ;
563+ Some ( format ! ( "{enum_header}{generics_str}\n {indent}{rw}" ) )
564+ }
565+ None => {
567566 self . block_indent = self . block_indent . block_unindent ( self . config ) ;
567+ None
568568 }
569569 }
570570 }
571571
572+ pub ( crate ) fn visit_enum (
573+ & mut self ,
574+ ident : symbol:: Ident ,
575+ vis : & ast:: Visibility ,
576+ enum_def : & ast:: EnumDef ,
577+ generics : & ast:: Generics ,
578+ span : Span ,
579+ ) {
580+ let rewrite = self . format_enum ( ident, vis, enum_def, generics, span) ;
581+ self . push_rewrite ( span, rewrite) ;
582+ }
583+
572584 // Format the body of an enum definition
573585 fn format_variant_list (
574586 & mut self ,
0 commit comments