impl ToTokens for Visibility and use it in codegen where possible

use-ink · Sep 23, 2021 · 02a940c · 02a940c
1 parent c393671
commit 02a940c
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 16 deletions.
diff --git a/crates/lang/codegen/src/generator/item_impls.rs b/crates/lang/codegen/src/generator/item_impls.rs
@@ -61,10 +61,7 @@ impl ItemImpls<'_> {
     fn generate_trait_constructor(constructor: &ir::Constructor) -> TokenStream2 {
         let span = constructor.span();
         let attrs = constructor.attrs();
-        let vis = match constructor.visibility() {
-            ir::Visibility::Inherited => None,
-            ir::Visibility::Public(vis_public) => Some(vis_public),
-        };
+        let vis = constructor.visibility();
         let ident = constructor.ident();
         let output_ident = format_ident!("{}Out", ident.to_string().to_camel_case());
         let inputs = constructor.inputs();
@@ -83,10 +80,7 @@ impl ItemImpls<'_> {
     fn generate_trait_message(message: &ir::Message) -> TokenStream2 {
         let span = message.span();
         let attrs = message.attrs();
-        let vis = match message.visibility() {
-            ir::Visibility::Inherited => None,
-            ir::Visibility::Public(vis_public) => Some(vis_public),
-        };
+        let vis = message.visibility();
         let receiver = message.receiver();
         let ident = message.ident();
         let output_ident = format_ident!("{}Out", ident.to_string().to_camel_case());
@@ -161,10 +155,7 @@ impl ItemImpls<'_> {
     fn generate_inherent_constructor(constructor: &ir::Constructor) -> TokenStream2 {
         let span = constructor.span();
         let attrs = constructor.attrs();
-        let vis = match constructor.visibility() {
-            ir::Visibility::Inherited => None,
-            ir::Visibility::Public(vis_public) => Some(vis_public),
-        };
+        let vis = constructor.visibility();
         let ident = constructor.ident();
         let inputs = constructor.inputs();
         let statements = constructor.statements();
@@ -180,10 +171,7 @@ impl ItemImpls<'_> {
     fn generate_inherent_message(message: &ir::Message) -> TokenStream2 {
         let span = message.span();
         let attrs = message.attrs();
-        let vis = match message.visibility() {
-            ir::Visibility::Inherited => None,
-            ir::Visibility::Public(vis_public) => Some(vis_public),
-        };
+        let vis = message.visibility();
         let receiver = message.receiver();
         let ident = message.ident();
         let inputs = message.inputs();

diff --git a/crates/lang/ir/src/ir/item_impl/callable.rs b/crates/lang/ir/src/ir/item_impl/callable.rs
@@ -440,6 +440,15 @@ pub enum Visibility {
     Inherited,
 }
 
+impl quote::ToTokens for Visibility {
+    fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
+        match self {
+            Self::Public(vis_public) => vis_public.to_tokens(tokens),
+            Self::Inherited => (),
+        }
+    }
+}
+
 impl Visibility {
     /// Returns `true` if the visibility of the ink! message of constructor is public (`pub`).
     ///