Print safe and explicitly unsafe foreign items

Author: dtolnay

src/item.rs

@@ -91,7 +91,11 @@ impl Printer {
         self.outer_attrs(&item.attrs);
         self.cbox(INDENT);
         self.visibility(&item.vis);
-        self.signature(&item.sig);
+        self.signature(
+            &item.sig,
+            #[cfg(feature = "verbatim")]
+            &verbatim::Safety::Disallowed,
+        );
         self.where_clause_for_body(&item.sig.generics.where_clause);
         self.word("{");
         self.hardbreak_if_nonempty();
@@ -794,7 +798,11 @@ impl Printer {
         self.outer_attrs(&foreign_item.attrs);
         self.cbox(INDENT);
         self.visibility(&foreign_item.vis);
-        self.signature(&foreign_item.sig);
+        self.signature(
+            &foreign_item.sig,
+            #[cfg(feature = "verbatim")]
+            &verbatim::Safety::Disallowed,
+        );
         self.where_clause_semi(&foreign_item.sig.generics.where_clause);
         self.end();
         self.hardbreak();
@@ -844,8 +852,10 @@ impl Printer {
     #[cfg(feature = "verbatim")]
     fn foreign_item_verbatim(&mut self, tokens: &TokenStream) {
         use syn::parse::{Parse, ParseStream, Result};
-        use syn::{Attribute, Token, Visibility};
-        use verbatim::{FlexibleItemFn, FlexibleItemStatic, FlexibleItemType, WhereClauseLocation};
+        use syn::{Abi, Attribute, Token, Visibility};
+        use verbatim::{
+            kw, FlexibleItemFn, FlexibleItemStatic, FlexibleItemType, WhereClauseLocation,
+        };
 
         enum ForeignItemVerbatim {
             Empty,
@@ -855,6 +865,16 @@ impl Printer {
             TypeFlexible(FlexibleItemType),
         }
 
+        fn peek_signature(input: ParseStream) -> bool {
+            let fork = input.fork();
+            fork.parse::<Option<Token![const]>>().is_ok()
+                && fork.parse::<Option<Token![async]>>().is_ok()
+                && ((fork.peek(kw::safe) && fork.parse::<kw::safe>().is_ok())
+                    || fork.parse::<Option<Token![unsafe]>>().is_ok())
+                && fork.parse::<Option<Abi>>().is_ok()
+                && fork.peek(Token![fn])
+        }
+
         impl Parse for ForeignItemVerbatim {
             fn parse(input: ParseStream) -> Result<Self> {
                 if input.is_empty() {
@@ -869,15 +889,13 @@ impl Printer {
                 let defaultness = false;
 
                 let lookahead = input.lookahead1();
-                if lookahead.peek(Token![const])
-                    || lookahead.peek(Token![async])
-                    || lookahead.peek(Token![unsafe])
-                    || lookahead.peek(Token![extern])
-                    || lookahead.peek(Token![fn])
-                {
+                if lookahead.peek(Token![fn]) || peek_signature(input) {
                     let flexible_item = FlexibleItemFn::parse(attrs, vis, defaultness, input)?;
                     Ok(ForeignItemVerbatim::FnFlexible(flexible_item))
-                } else if lookahead.peek(Token![static]) {
+                } else if lookahead.peek(Token![static])
+                    || ((input.peek(Token![unsafe]) || input.peek(kw::safe))
+                        && input.peek2(Token![static]))
+                {
                     let flexible_item = FlexibleItemStatic::parse(attrs, vis, input)?;
                     Ok(ForeignItemVerbatim::StaticFlexible(flexible_item))
                 } else if lookahead.peek(Token![type]) {
@@ -953,7 +971,11 @@ impl Printer {
     fn trait_item_fn(&mut self, trait_item: &TraitItemFn) {
         self.outer_attrs(&trait_item.attrs);
         self.cbox(INDENT);
-        self.signature(&trait_item.sig);
+        self.signature(
+            &trait_item.sig,
+            #[cfg(feature = "verbatim")]
+            &verbatim::Safety::Disallowed,
+        );
         if let Some(block) = &trait_item.default {
             self.where_clause_for_body(&trait_item.sig.generics.where_clause);
             self.word("{");
@@ -1149,7 +1171,11 @@ impl Printer {
         if impl_item.defaultness.is_some() {
             self.word("default ");
         }
-        self.signature(&impl_item.sig);
+        self.signature(
+            &impl_item.sig,
+            #[cfg(feature = "verbatim")]
+            &verbatim::Safety::Disallowed,
+        );
         self.where_clause_for_body(&impl_item.sig.generics.where_clause);
         self.word("{");
         self.hardbreak_if_nonempty();
@@ -1277,15 +1303,32 @@ impl Printer {
         }
     }
 
-    fn signature(&mut self, signature: &Signature) {
+    fn signature(
+        &mut self,
+        signature: &Signature,
+        #[cfg(feature = "verbatim")] safety: &verbatim::Safety,
+    ) {
         if signature.constness.is_some() {
             self.word("const ");
         }
         if signature.asyncness.is_some() {
             self.word("async ");
         }
-        if signature.unsafety.is_some() {
-            self.word("unsafe ");
+        #[cfg(feature = "verbatim")]
+        {
+            if let verbatim::Safety::Disallowed = safety {
+                if signature.unsafety.is_some() {
+                    self.word("unsafe ");
+                }
+            } else {
+                self.safety(safety);
+            }
+        }
+        #[cfg(not(feature = "verbatim"))]
+        {
+            if signature.unsafety.is_some() {
+                self.word("unsafe ");
+            }
         }
         if let Some(abi) = &signature.abi {
             self.abi(abi);
@@ -1381,12 +1424,16 @@ mod verbatim {
     use crate::iter::IterDelimited;
     use crate::INDENT;
     use syn::ext::IdentExt;
-    use syn::parse::{ParseStream, Result};
+    use syn::parse::{Parse, ParseStream, Result};
     use syn::{
         braced, token, Attribute, Block, Expr, Generics, Ident, Signature, StaticMutability, Stmt,
         Token, Type, TypeParamBound, Visibility, WhereClause,
     };
 
+    pub mod kw {
+        syn::custom_keyword!(safe);
+    }
+
     pub struct FlexibleItemConst {
         pub attrs: Vec<Attribute>,
         pub vis: Visibility,
@@ -1401,13 +1448,15 @@ mod verbatim {
         pub attrs: Vec<Attribute>,
         pub vis: Visibility,
         pub defaultness: bool,
+        pub safety: Safety,
         pub sig: Signature,
         pub body: Option<Vec<Stmt>>,
     }
 
     pub struct FlexibleItemStatic {
         pub attrs: Vec<Attribute>,
         pub vis: Visibility,
+        pub safety: Safety,
         pub mutability: StaticMutability,
         pub ident: Ident,
         pub ty: Option<Type>,
@@ -1425,6 +1474,13 @@ mod verbatim {
         pub where_clause_after_eq: Option<WhereClause>,
     }
 
+    pub enum Safety {
+        Unsafe,
+        Safe,
+        Default,
+        Disallowed,
+    }
+
     pub enum WhereClauseLocation {
         // type Ty<T> where T: 'static = T;
         BeforeEq,
@@ -1474,7 +1530,16 @@ mod verbatim {
             defaultness: bool,
             input: ParseStream,
         ) -> Result<Self> {
-            let sig: Signature = input.parse()?;
+            let constness: Option<Token![const]> = input.parse()?;
+            let asyncness: Option<Token![async]> = input.parse()?;
+            let safety: Safety = input.parse()?;
+
+            let lookahead = input.lookahead1();
+            let sig: Signature = if lookahead.peek(Token![extern]) || lookahead.peek(Token![fn]) {
+                input.parse()?
+            } else {
+                return Err(lookahead.error());
+            };
 
             let lookahead = input.lookahead1();
             let body = if lookahead.peek(Token![;]) {
@@ -1493,14 +1558,21 @@ mod verbatim {
                 attrs,
                 vis,
                 defaultness,
-                sig,
+                safety,
+                sig: Signature {
+                    constness,
+                    asyncness,
+                    unsafety: None,
+                    ..sig
+                },
                 body,
             })
         }
     }
 
     impl FlexibleItemStatic {
         pub fn parse(attrs: Vec<Attribute>, vis: Visibility, input: ParseStream) -> Result<Self> {
+            let safety: Safety = input.parse()?;
             input.parse::<Token![static]>()?;
             let mutability: StaticMutability = input.parse()?;
             let ident = input.parse()?;
@@ -1530,6 +1602,7 @@ mod verbatim {
             Ok(FlexibleItemStatic {
                 attrs,
                 vis,
+                safety,
                 mutability,
                 ident,
                 ty,
@@ -1601,6 +1674,20 @@ mod verbatim {
         }
     }
 
+    impl Parse for Safety {
+        fn parse(input: ParseStream) -> Result<Self> {
+            if input.peek(Token![unsafe]) {
+                input.parse::<Token![unsafe]>()?;
+                Ok(Safety::Unsafe)
+            } else if input.peek(kw::safe) {
+                input.parse::<kw::safe>()?;
+                Ok(Safety::Safe)
+            } else {
+                Ok(Safety::Default)
+            }
+        }
+    }
+
     impl Printer {
         pub fn flexible_item_const(&mut self, item: &FlexibleItemConst) {
             self.outer_attrs(&item.attrs);
@@ -1635,7 +1722,7 @@ mod verbatim {
             if item.defaultness {
                 self.word("default ");
             }
-            self.signature(&item.sig);
+            self.signature(&item.sig, &item.safety);
             if let Some(body) = &item.body {
                 self.where_clause_for_body(&item.sig.generics.where_clause);
                 self.word("{");
@@ -1658,6 +1745,7 @@ mod verbatim {
             self.outer_attrs(&item.attrs);
             self.cbox(0);
             self.visibility(&item.vis);
+            self.safety(&item.safety);
             self.word("static ");
             self.static_mutability(&item.mutability);
             self.ident(&item.ident);
@@ -1708,5 +1796,14 @@ mod verbatim {
             self.end();
             self.hardbreak();
         }
+
+        pub fn safety(&mut self, safety: &Safety) {
+            match safety {
+                Safety::Unsafe => self.word("unsafe "),
+                Safety::Safe => self.word("safe "),
+                Safety::Default => {}
+                Safety::Disallowed => unreachable!(),
+            }
+        }
     }
 }