Merge pull request #4 from spastorino/erased-handle-lifetimes

Properly handle lifetimes on erased trait generation

dynosaur/tests/pass/basic-with-self.rs

@@ -0,0 +1,9 @@
+use dynosaur::dynosaur;
+
+#[dynosaur(DynMyTrait)]
+trait MyTrait {
+    type Item;
+    async fn foo(&self) -> Self::Item;
+}
+
+fn main() {}

dynosaur/tests/pass/basic-with-self.stdout

@@ -0,0 +1,24 @@
+#![feature(prelude_import)]
+#[prelude_import]
+use std::prelude::rust_2021::*;
+#[macro_use]
+extern crate std;
+use dynosaur::dynosaur;
+
+trait MyTrait {
+    type Item;
+    async fn foo(&self)
+    -> Self::Item;
+}
+trait ErasedMyTrait {
+    type Item;
+    fn foo<'life0, 'dynosaur>(&'life0 self)
+    ->
+        ::core::pin::Pin<Box<dyn ::core::future::Future<Output = Self::Item> +
+        'dynosaur>>
+    where
+    'life0: 'dynosaur,
+    Self: 'dynosaur;
+}
+
+fn main() {}

dynosaur/tests/pass/basic.rs

@@ -2,8 +2,7 @@ use dynosaur::dynosaur;
 
 #[dynosaur(DynMyTrait)]
 trait MyTrait {
-    type Item;
-    async fn foo(&self) -> Self::Item;
+    async fn foo(&self) -> i32;
 }
 
 fn main() {}

dynosaur/tests/pass/basic.stdout

@@ -6,14 +6,17 @@ extern crate std;
 use dynosaur::dynosaur;
 
 trait MyTrait {
-    type Item;
     async fn foo(&self)
-    -> Self::Item;
+    -> i32;
 }
 trait ErasedMyTrait {
-    type Item;
-    fn foo(&self)
-    -> ::core::pin::Pin<Box<dyn ::core::future::Future<Output = Self::Item>>>;
+    fn foo<'life0, 'dynosaur>(&'life0 self)
+    ->
+        ::core::pin::Pin<Box<dyn ::core::future::Future<Output = i32> +
+        'dynosaur>>
+    where
+    'life0: 'dynosaur,
+    Self: 'dynosaur;
 }
 
 fn main() {}

dynosaur/tests/pass/handle-lifetimes.rs

@@ -0,0 +1,9 @@
+use dynosaur::dynosaur;
+
+#[dynosaur(DynMyTrait)]
+trait MyTrait {
+    type Item;
+    async fn foo<T>(&self, x: &T) -> i32;
+}
+
+fn main() {}

dynosaur/tests/pass/handle-lifetimes.stdout

@@ -0,0 +1,26 @@
+#![feature(prelude_import)]
+#[prelude_import]
+use std::prelude::rust_2021::*;
+#[macro_use]
+extern crate std;
+use dynosaur::dynosaur;
+
+trait MyTrait {
+    type Item;
+    async fn foo<T>(&self, x: &T)
+    -> i32;
+}
+trait ErasedMyTrait {
+    type Item;
+    fn foo<'life0, 'life1, 'dynosaur, T>(&'life0 self, x: &'life1 T)
+    ->
+        ::core::pin::Pin<Box<dyn ::core::future::Future<Output = i32> +
+        'dynosaur>>
+    where
+    T: 'dynosaur,
+    'life0: 'dynosaur,
+    'life1: 'dynosaur,
+    Self: 'dynosaur;
+}
+
+fn main() {}

dynosaur_derive/Cargo.toml

@@ -17,4 +17,4 @@ proc-macro = true
 [dependencies]
 proc-macro2 = "1.0"
 quote = "1.0"
-syn = { version = "2.0", features = ["full"] }
+syn = { version = "2.0", features = ["full", "visit-mut"] }

dynosaur_derive/src/lib.rs

@@ -1,11 +1,20 @@
+use crate::lifetime::{AddLifetimeToImplTrait, CollectLifetimes};
+use crate::receiver::{has_self_in_sig, mut_pat};
 use proc_macro2::{Span, TokenStream};
-use quote::quote;
+use quote::{format_ident, quote};
+use std::mem;
+use syn::punctuated::Punctuated;
+use syn::visit_mut::VisitMut;
 use syn::{
     parse::{Parse, ParseStream},
-    parse_macro_input, parse_quote, Ident, ItemTrait, Result, ReturnType, Signature, Token,
-    TraitItem, TraitItemFn,
+    parse_macro_input, parse_quote, parse_quote_spanned, FnArg, GenericParam, Generics, Ident,
+    ItemTrait, Lifetime, LifetimeParam, Pat, Result, ReturnType, Signature, Token, TraitItem,
+    TraitItemFn, Type, WhereClause,
 };
 
+mod lifetime;
+mod receiver;
+
 struct Attrs {
     ident: Ident,
 }
@@ -56,45 +65,161 @@ pub fn dynosaur(
     .into()
 }
 
-fn mk_erased_trait(item: &ItemTrait) -> TokenStream {
-    let erased_trait = ItemTrait {
-        ident: Ident::new(&format!("Erased{}", item.ident), Span::call_site()),
-        items: item
-            .items
-            .iter()
-            .map(|item| {
-                if let TraitItem::Fn(
-                    trait_item_fn @ TraitItemFn {
-                        sig:
-                            Signature {
-                                asyncness: Some(..),
-                                ..
-                            },
-                        ..
-                    },
-                ) = item
-                {
-                    let (ret_arrow, ret) = match &trait_item_fn.sig.output {
-                        ReturnType::Default => (Token![->](Span::call_site()), quote!(())),
-                        ReturnType::Type(arrow, ret) => (*arrow, quote!(#ret)),
-                    };
-                    TraitItem::Fn(TraitItemFn {
-                        sig: Signature {
-                            asyncness: None,
-                            output: parse_quote! {
-                                #ret_arrow ::core::pin::Pin<Box<
-                                dyn ::core::future::Future<Output = #ret>>>
-                            },
-                            ..trait_item_fn.sig.clone()
+fn mk_erased_trait(item_trait: &ItemTrait) -> TokenStream {
+    let erased_trait =
+        ItemTrait {
+            ident: Ident::new(&format!("Erased{}", item_trait.ident), Span::call_site()),
+            items: item_trait
+                .items
+                .iter()
+                .map(|item| {
+                    if let TraitItem::Fn(
+                        trait_item_fn @ TraitItemFn {
+                            sig:
+                                Signature {
+                                    asyncness: Some(..),
+                                    ..
+                                },
+                            ..
                         },
-                        ..trait_item_fn.clone()
-                    })
-                } else {
-                    item.clone()
-                }
-            })
-            .collect(),
-        ..item.clone()
-    };
+                    ) = item
+                    {
+                        let mut sig = trait_item_fn.sig.clone();
+
+                        sig.fn_token.span = sig.asyncness.take().unwrap().span;
+
+                        let (ret_arrow, ret) = match &sig.output {
+                            ReturnType::Default => (Token![->](Span::call_site()), quote!(())),
+                            ReturnType::Type(arrow, ret) => (*arrow, quote!(#ret)),
+                        };
+
+                        let mut lifetimes = CollectLifetimes::new();
+                        for arg in &mut sig.inputs {
+                            match arg {
+                                FnArg::Receiver(arg) => lifetimes.visit_receiver_mut(arg),
+                                FnArg::Typed(arg) => lifetimes.visit_type_mut(&mut arg.ty),
+                            }
+                        }
+
+                        for param in &mut sig.generics.params {
+                            match param {
+                                GenericParam::Type(param) => {
+                                    let param_name = &param.ident;
+                                    let span = match param.colon_token.take() {
+                                        Some(colon_token) => colon_token.span,
+                                        None => param_name.span(),
+                                    };
+                                    let bounds = mem::replace(&mut param.bounds, Punctuated::new());
+                                    where_clause_or_default(&mut sig.generics.where_clause)
+                                        .predicates
+                                        .push(parse_quote_spanned!(span=> #param_name: 'dynosaur + #bounds));
+                                }
+                                GenericParam::Lifetime(param) => {
+                                    let param_name = &param.lifetime;
+                                    let span = match param.colon_token.take() {
+                                        Some(colon_token) => colon_token.span,
+                                        None => param_name.span(),
+                                    };
+                                    let bounds = mem::replace(&mut param.bounds, Punctuated::new());
+                                    where_clause_or_default(&mut sig.generics.where_clause)
+                                        .predicates
+                                        .push(parse_quote_spanned!(span=> #param: 'dynosaur + #bounds));
+                                }
+                                GenericParam::Const(_) => {}
+                            }
+                        }
+
+                        for param in used_lifetimes(&item_trait.generics, &lifetimes.explicit) {
+                            let param = &param.lifetime;
+                            let span = param.span();
+                            where_clause_or_default(&mut sig.generics.where_clause)
+                                .predicates
+                                .push(parse_quote_spanned!(span=> #param: 'dynosaur));
+                        }
+
+                        if sig.generics.lt_token.is_none() {
+                            sig.generics.lt_token = Some(Token![<](sig.ident.span()));
+                        }
+                        if sig.generics.gt_token.is_none() {
+                            sig.generics.gt_token = Some(Token![>](sig.paren_token.span.join()));
+                        }
+
+                        for elided in lifetimes.elided {
+                            sig.generics.params.push(parse_quote!(#elided));
+                            where_clause_or_default(&mut sig.generics.where_clause)
+                                .predicates
+                                .push(parse_quote_spanned!(elided.span()=> #elided: 'dynosaur));
+                        }
+
+                        sig.generics.params.push(parse_quote!('dynosaur));
+
+                        if has_self_in_sig(&mut sig) {
+                            where_clause_or_default(&mut sig.generics.where_clause)
+                                .predicates
+                                .push(parse_quote! {
+                                    Self: 'dynosaur
+                                });
+                        }
+
+                        for (i, arg) in sig.inputs.iter_mut().enumerate() {
+                            if let FnArg::Typed(arg) = arg {
+                                if match *arg.ty {
+                                    Type::Reference(_) => false,
+                                    _ => true,
+                                } {
+                                    match &*arg.pat {
+                                        Pat::Ident(_) => {}
+                                        _ => {
+                                            let positional = positional_arg(i, &arg.pat);
+                                            let m = mut_pat(&mut arg.pat);
+                                            arg.pat = parse_quote!(#m #positional);
+                                        }
+                                    }
+                                }
+                                AddLifetimeToImplTrait.visit_type_mut(&mut arg.ty);
+                            }
+                        }
+
+                        sig.output = parse_quote! {
+                            #ret_arrow ::core::pin::Pin<Box<
+                            dyn ::core::future::Future<Output = #ret> + 'dynosaur>>
+                        };
+                        TraitItem::Fn(TraitItemFn {
+                            sig,
+                            ..trait_item_fn.clone()
+                        })
+                    } else {
+                        item.clone()
+                    }
+                })
+                .collect(),
+            ..item_trait.clone()
+        };
     quote! { #erased_trait }
 }
+
+fn used_lifetimes<'a>(
+    generics: &'a Generics,
+    used: &'a [Lifetime],
+) -> impl Iterator<Item = &'a LifetimeParam> {
+    generics.params.iter().filter_map(move |param| {
+        if let GenericParam::Lifetime(param) = param {
+            if used.contains(&param.lifetime) {
+                return Some(param);
+            }
+        }
+        None
+    })
+}
+
+fn where_clause_or_default(clause: &mut Option<WhereClause>) -> &mut WhereClause {
+    clause.get_or_insert_with(|| WhereClause {
+        where_token: Default::default(),
+        predicates: Punctuated::new(),
+    })
+}
+
+fn positional_arg(i: usize, pat: &Pat) -> Ident {
+    let span = syn::spanned::Spanned::span(pat).resolved_at(Span::mixed_site());
+    format_ident!("__arg{}", i, span = span)
+}