From ba1f22ca5bea0966029ffa322acf6a17e6b92538 Mon Sep 17 00:00:00 2001
From: Yacin Tmimi <yacintmimi@gmail.com>
Date: Tue, 7 Feb 2023 18:43:07 -0500
Subject: [PATCH] Adjust enum variant spans to exclude any explicit
 discriminant

Fixes 5686

For reference, explicit discriminants were proposed in [RFC-2363].

`ast::Variant` spans extend to include explicit discriminants when they
are present.

Now we'll adjust the span of enum variants to exclude any explicit
discriminant.

[RFC-2363]: https://rust-lang.github.io/rfcs/2363-arbitrary-enum-discriminant.html
---
 src/items.rs               | 25 ++++++++++++++++++++---
 tests/source/issue_5686.rs | 40 ++++++++++++++++++++++++++++++++++++
 tests/target/issue_5686.rs | 42 ++++++++++++++++++++++++++++++++++++++
 3 files changed, 104 insertions(+), 3 deletions(-)
 create mode 100644 tests/source/issue_5686.rs
 create mode 100644 tests/target/issue_5686.rs

diff --git a/src/items.rs b/src/items.rs
index 3c5293b6bf5..714bcbc10b4 100644
--- a/src/items.rs
+++ b/src/items.rs
@@ -560,7 +560,7 @@ impl<'a> FmtVisitor<'a> {
         let variant_body = match field.data {
             ast::VariantData::Tuple(..) | ast::VariantData::Struct(..) => format_struct(
                 &context,
-                &StructParts::from_variant(field),
+                &StructParts::from_variant(field, &context),
                 self.block_indent,
                 Some(one_line_width),
             )?,
@@ -951,14 +951,14 @@ impl<'a> StructParts<'a> {
         format_header(context, self.prefix, self.ident, self.vis, offset)
     }
 
-    fn from_variant(variant: &'a ast::Variant) -> Self {
+    fn from_variant(variant: &'a ast::Variant, context: &RewriteContext<'_>) -> Self {
         StructParts {
             prefix: "",
             ident: variant.ident,
             vis: &DEFAULT_VISIBILITY,
             def: &variant.data,
             generics: None,
-            span: variant.span,
+            span: enum_variant_span(variant, context),
         }
     }
 
@@ -979,6 +979,25 @@ impl<'a> StructParts<'a> {
     }
 }
 
+fn enum_variant_span(variant: &ast::Variant, context: &RewriteContext<'_>) -> Span {
+    use ast::VariantData::*;
+    if let Some(ref anon_const) = variant.disr_expr {
+        let span_before_consts = variant.span.until(anon_const.value.span);
+        let hi = match &variant.data {
+            Struct(..) => context
+                .snippet_provider
+                .span_after_last(span_before_consts, "}"),
+            Tuple(..) => context
+                .snippet_provider
+                .span_after_last(span_before_consts, ")"),
+            Unit(..) => variant.ident.span.hi(),
+        };
+        mk_sp(span_before_consts.lo(), hi)
+    } else {
+        variant.span
+    }
+}
+
 fn format_struct(
     context: &RewriteContext<'_>,
     struct_parts: &StructParts<'_>,
diff --git a/tests/source/issue_5686.rs b/tests/source/issue_5686.rs
new file mode 100644
index 00000000000..3bf96f73b2c
--- /dev/null
+++ b/tests/source/issue_5686.rs
@@ -0,0 +1,40 @@
+#[repr(u8)]
+enum MyEnum {
+    UnitWithExplicitDiscriminant = 0,
+    EmptyStructSingleLineBlockComment {
+        /* Comment */
+    } = 1,
+    EmptyStructMultiLineBlockComment {
+        /*
+         * Comment
+         */
+    } = 2,
+    EmptyStructLineComment {
+        // comment
+    } = 3,
+    EmptyTupleSingleLineBlockComment(
+        /* Comment */
+    ) = 4,
+    EmptyTupleMultiLineBlockComment(
+        /*
+         * Comment
+         */
+    ) = 5,
+    EmptyTupleLineComment(
+        // comment
+    ) = 6,
+}
+
+enum Animal {
+    Dog(/* tuple variant closer in comment -> ) */) = 1,
+    #[hello(world)]
+    Cat(/* tuple variant close in leading attribute */) = 2,
+    Bee(/* tuple variant closer on associated field attribute */ #[hello(world)] usize) = 3,
+    Fox(/* tuple variant closer on const fn call */) = some_const_fn(),
+    Ant(/* tuple variant closer on macro call */) = some_macro!(),
+    Snake {/* stuct variant closer in comment -> } */} = 6,
+    #[hell{world}]
+    Cobra {/* struct variant close in leading attribute */} = 6,
+    Eagle {/* struct variant closer on associated field attribute */ #[hell{world}]value: Sting} = 7,
+    Koala {/* struct variant closer on macro call */} = some_macro!{}
+}
diff --git a/tests/target/issue_5686.rs b/tests/target/issue_5686.rs
new file mode 100644
index 00000000000..993f12b5316
--- /dev/null
+++ b/tests/target/issue_5686.rs
@@ -0,0 +1,42 @@
+#[repr(u8)]
+enum MyEnum {
+    UnitWithExplicitDiscriminant = 0,
+    EmptyStructSingleLineBlockComment {/* Comment */} = 1,
+    EmptyStructMultiLineBlockComment {
+        /*
+         * Comment
+         */
+    } = 2,
+    EmptyStructLineComment {
+        // comment
+    } = 3,
+    EmptyTupleSingleLineBlockComment(/* Comment */) = 4,
+    EmptyTupleMultiLineBlockComment(
+        /*
+         * Comment
+         */
+    ) = 5,
+    EmptyTupleLineComment(
+        // comment
+    ) = 6,
+}
+
+enum Animal {
+    Dog(/* tuple variant closer in comment -> ) */) = 1,
+    #[hello(world)]
+    Cat(/* tuple variant close in leading attribute */) = 2,
+    Bee(
+        /* tuple variant closer on associated field attribute */ #[hello(world)] usize,
+    ) = 3,
+    Fox(/* tuple variant closer on const fn call */) = some_const_fn(),
+    Ant(/* tuple variant closer on macro call */) = some_macro!(),
+    Snake {/* stuct variant closer in comment -> } */} = 6,
+    #[hell{world}]
+    Cobra {/* struct variant close in leading attribute */} = 6,
+    Eagle {
+        /* struct variant closer on associated field attribute */
+        #[hell{world}]
+        value: Sting,
+    } = 7,
+    Koala {/* struct variant closer on macro call */} = some_macro! {},
+}