From dba89b508d558139fbda38fe0f6072f4308c91cc Mon Sep 17 00:00:00 2001
From: Jianfeng Jiang <jiangjianfeng.jjf@antgroup.com>
Date: Mon, 8 May 2023 11:13:49 +0800
Subject: [PATCH] Extend Require macro to support impl block

---
 .vscode/settings.json                         |  3 ++-
 services/libs/jinux-rights-proc/src/lib.rs    | 11 +++++-----
 .../jinux-rights-proc/src/require_attr.rs     | 22 ++++++++++++++-----
 .../jinux-rights-proc/src/require_item.rs     | 18 +++++++++++++++
 4 files changed, 42 insertions(+), 12 deletions(-)
 create mode 100644 services/libs/jinux-rights-proc/src/require_item.rs

diff --git a/.vscode/settings.json b/.vscode/settings.json
index 4fe4a2e24..dd50ab5c3 100644
--- a/.vscode/settings.json
+++ b/.vscode/settings.json
@@ -6,5 +6,6 @@
         "x86_64-custom.json",
         "-Zbuild-std=core,alloc,compiler_builtins",
         "-Zbuild-std-features=compiler-builtins-mem"
-    ]
+    ],
+    "rust-analyzer.showUnlinkedFileNotification": false
 }
\ No newline at end of file
diff --git a/services/libs/jinux-rights-proc/src/lib.rs b/services/libs/jinux-rights-proc/src/lib.rs
index ba490e87b..21c7c80b1 100644
--- a/services/libs/jinux-rights-proc/src/lib.rs
+++ b/services/libs/jinux-rights-proc/src/lib.rs
@@ -15,7 +15,7 @@
 //! Below is a simple example.
 //! Suppose we have a channel that may have read and write capability.
 //!
-//! ```rust
+//! ```ignore
 //! /// A simple channel, only for demonstration.
 //! struct Channel<R: RightSet> {
 //!    rights: PhantomData<R>,
@@ -42,7 +42,7 @@
 //! ```
 //! When we initialize channels with different rights, it can check whether the function
 //! are wrongly used due to lack of capabilities at compilation time.
-//! ```rust
+//! ```ignore
 //!     let read_channel = Channel::<R>::new();
 //!     read_channel.read();
 //!     // read_channel.write();                    // compilation error!
@@ -73,20 +73,21 @@
 #![feature(proc_macro_diagnostic)]
 #![allow(dead_code)]
 
+use require_item::RequireItem;
 use syn::parse_macro_input;
-use syn::ItemFn;
 
 use crate::require_attr::expand_require;
 use crate::require_attr::RequireAttr;
 
 mod require_attr;
+mod require_item;
 
 #[proc_macro_attribute]
 pub fn require(
     attr: proc_macro::TokenStream,
     item: proc_macro::TokenStream,
 ) -> proc_macro::TokenStream {
-    let item_fn = parse_macro_input!(item as ItemFn);
+    let require_item = parse_macro_input!(item as RequireItem);
     let require_attr = parse_macro_input!(attr as RequireAttr);
-    expand_require(item_fn, require_attr).into()
+    expand_require(require_item, require_attr).into()
 }
diff --git a/services/libs/jinux-rights-proc/src/require_attr.rs b/services/libs/jinux-rights-proc/src/require_attr.rs
index 29c9da739..ffcebea3a 100644
--- a/services/libs/jinux-rights-proc/src/require_attr.rs
+++ b/services/libs/jinux-rights-proc/src/require_attr.rs
@@ -4,9 +4,11 @@ use proc_macro2::{Ident, TokenStream};
 use quote::quote;
 use syn::{
     fold::Fold, parse::Parse, parse_quote, punctuated::Punctuated, token::Comma, GenericParam,
-    Generics, ItemFn, Token, Type, WhereClause,
+    Generics, Token, Type, WhereClause,
 };
 
+use super::require_item::RequireItem;
+
 pub struct RequireAttr {
     type_set: Type,
     required_types: Punctuated<Ident, Token![|]>,
@@ -61,11 +63,19 @@ impl Fold for RequireAttr {
     }
 }
 
-pub fn expand_require(item_fn: ItemFn, mut require_attr: RequireAttr) -> TokenStream {
-    let new_item_fn = require_attr.fold_item_fn(item_fn);
-    quote!(
-        #new_item_fn
-    )
+pub fn expand_require(item: RequireItem, mut require_attr: RequireAttr) -> TokenStream {
+    match item {
+        RequireItem::Impl(item_impl) => {
+            let new_item_impl = require_attr.fold_item_impl(item_impl);
+            quote!(#new_item_impl)
+        }
+        RequireItem::Fn(item_fn) => {
+            let new_item_fn = require_attr.fold_item_fn(item_fn);
+            quote!(
+                #new_item_fn
+            )
+        }
+    }
 }
 
 fn is_generic_param(ident: Ident, generic_params: &Punctuated<GenericParam, Comma>) -> bool {
diff --git a/services/libs/jinux-rights-proc/src/require_item.rs b/services/libs/jinux-rights-proc/src/require_item.rs
new file mode 100644
index 000000000..a423539b5
--- /dev/null
+++ b/services/libs/jinux-rights-proc/src/require_item.rs
@@ -0,0 +1,18 @@
+use syn::{parse::Parse, ItemFn, ItemImpl, Token};
+
+pub enum RequireItem {
+    Impl(ItemImpl),
+    Fn(ItemFn),
+}
+
+impl Parse for RequireItem {
+    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
+        let token = input.lookahead1();
+        if token.peek(Token!(impl)) {
+            // FIXME: is there any possible token before impl?
+            input.parse().map(RequireItem::Impl)
+        } else {
+            input.parse().map(RequireItem::Fn)
+        }
+    }
+}