Allow libp2p_prefix parameter in NetworkBehaviour derive macro

libp2p · Oct 10, 2022 · ed03603 · ed03603
1 parent aba5ccb
commit ed03603
Show file tree

Hide file tree

Showing 2 changed files with 55 additions and 16 deletions.
diff --git a/swarm-derive/src/lib.rs b/swarm-derive/src/lib.rs
@@ -46,21 +46,43 @@ fn build(ast: &DeriveInput) -> TokenStream {
 fn build_struct(ast: &DeriveInput, data_struct: &DataStruct) -> TokenStream {
     let name = &ast.ident;
     let (_, ty_generics, where_clause) = ast.generics.split_for_impl();
-    let multiaddr = quote! {::libp2p::core::Multiaddr};
-    let trait_to_impl = quote! {::libp2p::swarm::NetworkBehaviour};
-    let either_ident = quote! {::libp2p::core::either::EitherOutput};
-    let network_behaviour_action = quote! {::libp2p::swarm::NetworkBehaviourAction};
-    let into_connection_handler = quote! {::libp2p::swarm::IntoConnectionHandler};
-    let connection_handler = quote! {::libp2p::swarm::ConnectionHandler};
-    let into_proto_select_ident = quote! {::libp2p::swarm::IntoConnectionHandlerSelect};
-    let peer_id = quote! {::libp2p::core::PeerId};
-    let connection_id = quote! {::libp2p::core::connection::ConnectionId};
-    let dial_errors = quote! {Option<&Vec<::libp2p::core::Multiaddr>>};
-    let connected_point = quote! {::libp2p::core::ConnectedPoint};
-    let listener_id = quote! {::libp2p::core::transport::ListenerId};
-    let dial_error = quote! {::libp2p::swarm::DialError};
-
-    let poll_parameters = quote! {::libp2p::swarm::PollParameters};
+
+    let user_provided_libp2p_prefix = ast
+        .attrs
+        .iter()
+        .filter_map(get_meta_items)
+        .flatten()
+        .filter_map(|meta_item| {
+            if let syn::NestedMeta::Meta(syn::Meta::NameValue(ref m)) = meta_item {
+                if m.path.is_ident("libp2p_prefix") {
+                    if let syn::Lit::Str(ref s) = m.lit {
+                        return Some(syn::parse_str::<syn::TypePath>(&s.value()).unwrap());
+                    }
+                }
+            }
+            None
+        })
+        .next();
+
+    let libp2p_prefix = match user_provided_libp2p_prefix {
+        Some(pfx) => quote! {#pfx},
+        None => quote! {::libp2p},
+    };
+    let multiaddr = quote! {#libp2p_prefix::core::Multiaddr};
+    let trait_to_impl = quote! {#libp2p_prefix::swarm::NetworkBehaviour};
+    let either_ident = quote! {#libp2p_prefix::core::either::EitherOutput};
+    let network_behaviour_action = quote! {#libp2p_prefix::swarm::NetworkBehaviourAction};
+    let into_connection_handler = quote! {#libp2p_prefix::swarm::IntoConnectionHandler};
+    let connection_handler = quote! {#libp2p_prefix::swarm::ConnectionHandler};
+    let into_proto_select_ident = quote! {#libp2p_prefix::swarm::IntoConnectionHandlerSelect};
+    let peer_id = quote! {#libp2p_prefix::core::PeerId};
+    let connection_id = quote! {#libp2p_prefix::core::connection::ConnectionId};
+    let dial_errors = quote! {Option<&Vec<#libp2p_prefix::core::Multiaddr>>};
+    let connected_point = quote! {#libp2p_prefix::core::ConnectedPoint};
+    let listener_id = quote! {#libp2p_prefix::core::transport::ListenerId};
+    let dial_error = quote! {#libp2p_prefix::swarm::DialError};
+
+    let poll_parameters = quote! {#libp2p_prefix::swarm::PollParameters};
 
     // Build the generics.
     let impl_generics = {
@@ -624,7 +646,7 @@ fn build_struct(ast: &DeriveInput, data_struct: &DataStruct) -> TokenStream {
             }
 
             fn poll(&mut self, cx: &mut std::task::Context, poll_params: &mut impl #poll_parameters) -> std::task::Poll<#network_behaviour_action<Self::OutEvent, Self::ConnectionHandler>> {
-                use libp2p::futures::prelude::*;
+                use #libp2p_prefix::futures::prelude::*;
                 #(#poll_stmts)*
                 std::task::Poll::Pending
             }

diff --git a/swarm-derive/tests/test.rs b/swarm-derive/tests/test.rs
@@ -133,6 +133,23 @@ fn custom_event() {
     }
 }
 
+#[test]
+fn custom_prefix() {
+    use libp2p as mylibp2p;
+    #[allow(dead_code)]
+    #[derive(NetworkBehaviour)]
+    #[behaviour(libp2p_prefix = "mylibp2p")]
+    struct Foo {
+        ping: ping::Behaviour,
+        identify: identify::Behaviour,
+    }
+
+    #[allow(dead_code)]
+    fn foo() {
+        require_net_behaviour::<Foo>();
+    }
+}
+
 #[test]
 fn custom_event_mismatching_field_names() {
     #[allow(dead_code)]