Add library support for generic lambdas

Author: lucic71

libcc2rs-macros/src/lambda.rs

@@ -9,17 +9,38 @@ use syn::parse::{Parse, ParseStream};
 use syn::{Block, Error, ExprClosure, Pat, Result, Stmt, Token, parse_macro_input};
 
 struct Lambda {
+    is_generic: bool,
     captures: Block,
-    closure: ExprClosure,
+    closures: Vec<ExprClosure>,
 }
 
 impl Parse for Lambda {
     fn parse(input: ParseStream) -> Result<Self> {
+        let is_generic = input.peek(syn::Ident);
+        if is_generic {
+            let marker: Ident = input.parse()?;
+            if marker != "Generic" {
+                return Err(Error::new(marker.span(), "expected `Generic`"));
+            }
+            input.parse::<Token![,]>()?;
+        }
         let captures = input.parse()?;
         input.parse::<Token![,]>()?;
-        let closure = input.parse()?;
-        input.parse::<Option<Token![,]>>()?;
-        Ok(Lambda { captures, closure })
+        let mut closures = vec![input.parse()?];
+        while input.parse::<Option<Token![,]>>()?.is_some() && !input.is_empty() {
+            closures.push(input.parse()?);
+        }
+        if !is_generic && closures.len() != 1 {
+            return Err(Error::new_spanned(
+                &closures[1],
+                "a non-generic lambda has one closure",
+            ));
+        }
+        Ok(Lambda {
+            is_generic,
+            captures,
+            closures,
+        })
     }
 }
 
@@ -137,28 +158,49 @@ fn expand_lambda(lambda: Lambda, is_unsafe: bool) -> Result<proc_macro2::TokenSt
         inits.push(&init.expr);
     }
 
-    let mut params = Vec::new();
-    let mut param_types = Vec::new();
-    for input in &lambda.closure.inputs {
-        let Pat::Type(pat) = input else {
-            return Err(Error::new_spanned(input, "expected a typed parameter"));
+    let captured: HashSet<String> = names.iter().map(|name| name.to_string()).collect();
+    let mut calls = Vec::new();
+    let mut sigs = Vec::new();
+    let mut call_names = Vec::new();
+    for (i, closure) in lambda.closures.iter().enumerate() {
+        let mut params = Vec::new();
+        let mut param_types = Vec::new();
+        for input in &closure.inputs {
+            let Pat::Type(pat) = input else {
+                return Err(Error::new_spanned(input, "expected a typed parameter"));
+            };
+            params.push(pat);
+            param_types.push(&pat.ty);
+        }
+        let output = &closure.output;
+        let body = &closure.body;
+        let body = rewrite_captures(quote! { #body }, &captured, false)?;
+        let (receiver, body) = if is_unsafe {
+            (quote! { &mut self }, quote! { unsafe { #body } })
+        } else {
+            (quote! { &self }, body)
         };
-        params.push(pat);
-        param_types.push(&pat.ty);
+        let call_name = if lambda.is_generic {
+            Ident::new(&format!("call_{i}"), proc_macro2::Span::call_site())
+        } else {
+            Ident::new("call", proc_macro2::Span::call_site())
+        };
+        calls.push(quote! {
+            #[allow(unused_unsafe, clippy::too_many_arguments)]
+            fn #call_name(#receiver #(, #params)*) #output {
+                #body
+            }
+        });
+        sigs.push(quote! { fn(#(#param_types),*) #output });
+        call_names.push((call_name, closure.inputs.len()));
     }
-
-    let output = &lambda.closure.output;
-    let body = &lambda.closure.body;
-    let captured = names.iter().map(|name| name.to_string()).collect();
-    let body = rewrite_captures(quote! { #body }, &captured, false)?;
     // The copy and move constructors of the lambda copy and move each
     // capture, and its destructor destroys them in reverse order, as chosen
     // from their types by the traits in libcc2rs::__capture.
     let reversed: Vec<_> = names.iter().rev().collect();
-    let (receiver, body, destroy, constructor, copy_from, move_from) = if is_unsafe {
+    let (receiver, destroy, constructor, copy_from, move_from) = if is_unsafe {
         (
             quote! { &mut self },
-            quote! { unsafe { #body } },
             quote! {
                 use ::libcc2rs::__capture::{DestroyNoneUnsafe as _, DestroyWithDtorUnsafe as _};
                 unsafe {
@@ -192,7 +234,6 @@ fn expand_lambda(lambda: Lambda, is_unsafe: bool) -> Result<proc_macro2::TokenSt
     } else {
         (
             quote! { &self },
-            body,
             quote! {
                 use ::libcc2rs::__capture::{
                     DestroyElems as _, DestroyNone as _, DestroyWithDtor as _,
@@ -217,15 +258,12 @@ fn expand_lambda(lambda: Lambda, is_unsafe: bool) -> Result<proc_macro2::TokenSt
         )
     };
 
-    Ok(quote! {{
+    let lambda_struct = quote! {
         struct __Lambda {
             #(#names: #types,)*
         }
         impl __Lambda {
-            #[allow(unused_unsafe, clippy::too_many_arguments)]
-            fn call(#receiver #(, #params)*) #output {
-                #body
-            }
+            #(#calls)*
             #[allow(unused_imports, unused_unsafe)]
             fn copy_from(#receiver) -> Self {
                 #copy_from
@@ -239,12 +277,76 @@ fn expand_lambda(lambda: Lambda, is_unsafe: bool) -> Result<proc_macro2::TokenSt
                 #destroy
             }
         }
-        ::libcc2rs::FnPtr::<fn(#(#param_types),*) #output>::#constructor(
-            __Lambda { #(#names: #inits,)* },
-            __Lambda::call,
-            __Lambda::copy_from,
-            __Lambda::move_from,
-            __Lambda::destroy,
+    };
+    let lambda_init = quote! { __Lambda { #(#names: #inits,)* } };
+
+    if !lambda.is_generic {
+        let sig = &sigs[0];
+        return Ok(quote! {{
+            #lambda_struct
+            ::libcc2rs::FnPtr::<#sig>::#constructor(
+                #lambda_init,
+                __Lambda::call,
+                __Lambda::copy_from,
+                __Lambda::move_from,
+                __Lambda::destroy,
+            )
+        }});
+    }
+
+    let mut specs = Vec::new();
+    for (sig, (call_name, arity)) in sigs.iter().zip(&call_names) {
+        let args: Vec<Ident> = (0..*arity)
+            .map(|i| Ident::new(&format!("__a{i}"), proc_macro2::Span::call_site()))
+            .collect();
+        let call = if is_unsafe {
+            quote! {
+                |__l: &mut ::std::rc::Rc<::std::cell::UnsafeCell<__Lambda>> #(, #args)*| {
+                    unsafe { &mut *__l.get() }.#call_name(#(#args),*)
+                }
+            }
+        } else {
+            quote! {
+                |__l: &::std::rc::Rc<__Lambda> #(, #args)*| __l.#call_name(#(#args),*)
+            }
+        };
+        specs.push(quote! {
+            ::std::rc::Rc::new(::libcc2rs::FnPtr::<#sig>::#constructor(
+                ::std::rc::Rc::clone(__l),
+                #call,
+                |__l| ::std::rc::Rc::clone(__l),
+                |__l| ::std::rc::Rc::clone(__l),
+                |_| {},
+            )) as ::std::rc::Rc<dyn ::std::any::Any>
+        });
+    }
+    let (shared, copy_from, move_from, destroy) = if is_unsafe {
+        (
+            quote! { ::std::rc::Rc::new(::std::cell::UnsafeCell::new(#lambda_init)) },
+            quote! { |__l| ::std::rc::Rc::new(::std::cell::UnsafeCell::new(
+                unsafe { &mut *__l.get() }.copy_from()
+            )) },
+            quote! { |__l| ::std::rc::Rc::new(::std::cell::UnsafeCell::new(
+                unsafe { &mut *__l.get() }.move_from()
+            )) },
+            quote! { |__l| unsafe { &mut *__l.get() }.destroy() },
+        )
+    } else {
+        (
+            quote! { ::std::rc::Rc::new(#lambda_init) },
+            quote! { |__l| ::std::rc::Rc::new(__l.copy_from()) },
+            quote! { |__l| ::std::rc::Rc::new(__l.move_from()) },
+            quote! { |__l| __l.destroy() },
+        )
+    };
+    Ok(quote! {{
+        #lambda_struct
+        ::libcc2rs::FnPtr::<::libcc2rs::Generic>::from_generic_lambda(
+            #shared,
+            |__l| vec![#(#specs),*],
+            #copy_from,
+            #move_from,
+            #destroy,
         )
     }})
 }

libcc2rs/src/fn_ptr.rs

@@ -60,6 +60,10 @@ trait Adapted: Any {
     }
 
     fn destroy(&self) {}
+
+    fn specializations(&self) -> &[Rc<dyn Any>] {
+        panic!("ub: not a generic lambda");
+    }
 }
 
 impl<T: FnSig> Adapted for T {
@@ -373,6 +377,100 @@ macro_rules! impl_fn_ptr_call {
 }
 impl_fn_ptr_call!();
 
+#[derive(Clone, Copy)]
+pub struct Generic;
+
+impl FnSig for Generic {
+    type Args = ();
+    type Ret = ();
+    fn fn_addr(&self) -> usize {
+        panic!("ub: a generic lambda has no single function address");
+    }
+    fn call_direct(self, _args: ()) {
+        panic!("ub: calling a generic lambda without a signature");
+    }
+}
+
+struct GenericLambda<Captures> {
+    captures: Captures,
+    specializations: Vec<Rc<dyn Any>>,
+    make_specializations: fn(&Captures) -> Vec<Rc<dyn Any>>,
+    copy_from: fn(&Captures) -> Captures,
+    move_from: fn(&Captures) -> Captures,
+    destroy: fn(&Captures),
+}
+
+impl<Captures: 'static> GenericLambda<Captures> {
+    fn with_captures(&self, captures: Captures) -> Rc<dyn Adapted> {
+        Rc::new(GenericLambda {
+            specializations: (self.make_specializations)(&captures),
+            captures,
+            make_specializations: self.make_specializations,
+            copy_from: self.copy_from,
+            move_from: self.move_from,
+            destroy: self.destroy,
+        })
+    }
+}
+
+impl<Captures: 'static> Adapted for GenericLambda<Captures> {
+    fn call_adapted(&self, _args: ArgList) -> ArgRepr {
+        panic!("ub: calling a generic lambda without a signature");
+    }
+
+    fn copy_from(&self) -> Option<Rc<dyn Adapted>> {
+        Some(self.with_captures((self.copy_from)(&self.captures)))
+    }
+
+    fn move_from(&self) -> Option<Rc<dyn Adapted>> {
+        Some(self.with_captures((self.move_from)(&self.captures)))
+    }
+
+    fn destroy(&self) {
+        (self.destroy)(&self.captures)
+    }
+
+    fn specializations(&self) -> &[Rc<dyn Any>] {
+        &self.specializations
+    }
+}
+
+impl FnPtr<Generic> {
+    pub fn from_generic_lambda<Captures: 'static>(
+        captures: Captures,
+        make_specializations: fn(&Captures) -> Vec<Rc<dyn Any>>,
+        copy_from: fn(&Captures) -> Captures,
+        move_from: fn(&Captures) -> Captures,
+        destroy: fn(&Captures),
+    ) -> Self {
+        let lambda: Rc<dyn Adapted> = Rc::new(GenericLambda {
+            specializations: make_specializations(&captures),
+            captures,
+            make_specializations,
+            copy_from,
+            move_from,
+            destroy,
+        });
+        FnPtr {
+            addr: Rc::as_ptr(&lambda) as *const () as usize,
+            current: None,
+            original: Some(lambda),
+        }
+    }
+
+    pub fn spec<U: FnSig>(&self) -> FnPtr<U> {
+        assert!(!self.is_null(), "ub: null fn pointer call");
+        self.original
+            .as_ref()
+            .unwrap()
+            .specializations()
+            .iter()
+            .find_map(|spec| spec.downcast_ref::<FnPtr<U>>())
+            .expect("ub: calling through incompatible fn pointer type")
+            .clone()
+    }
+}
+
 impl<T: FnSig> Clone for FnPtr<T> {
     fn clone(&self) -> Self {
         FnPtr {

libcc2rs/src/lib.rs

@@ -35,7 +35,7 @@ mod fn_ptr_arg;
 pub use fn_ptr_arg::{ArgRepr, FnPtrArg, record_from_repr};
 
 mod fn_ptr;
-pub use fn_ptr::FnPtr;
+pub use fn_ptr::{FnPtr, Generic};
 
 mod capture;
 pub use capture::{Destructor, DestructorUnsafe, MoveCtor, MoveCtorUnsafe};