Avoid having to recreate a new pub token

dtolnay · Oct 27, 2020 · 0ba7439 · 0ba7439
1 parent 559631e
commit 0ba7439
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/src/lib.rs b/src/lib.rs
@@ -156,7 +156,7 @@ use crate::parse::{
 use proc_macro::{Ident, Punct, Spacing, Span, TokenStream, TokenTree};
 use std::fmt::Write;
 
-type Visibility = Option<Span>;
+type Visibility = Option<Ident>;
 
 enum Input {
     Export(Export),
@@ -250,7 +250,7 @@ fn expand_export(export: Export, args: ExportArgs) -> TokenStream {
     let dummy = dummy_name_for_export(&export);
 
     let attrs = export.attrs;
-    let ref vis = export.vis.map(|span| Ident::new("pub", span));
+    let vis = export.vis;
     let macro_export = match vis {
         Some(_) => quote!(#[macro_export]),
         None => quote!(),

diff --git a/src/parse.rs b/src/parse.rs
@@ -139,7 +139,10 @@ fn parse_group(tokens: Iter, delimiter: Delimiter) -> Result<IterImpl, Error> {
 fn parse_visibility(tokens: Iter) -> Result<Visibility, Error> {
     if let Some(TokenTree::Ident(ident)) = tokens.peek() {
         if ident.to_string() == "pub" {
-            return Ok(Some(tokens.next().unwrap().span()));
+            match tokens.next().unwrap() {
+                TokenTree::Ident(vis) => return Ok(Some(vis)),
+                _ => unreachable!(),
+            }
         }
     }
     Ok(None)