Add lambda and lambda_unsafe macros

Author: lucic71

libcc2rs-macros/src/lambda.rs

@@ -0,0 +1,87 @@
+// Copyright (c) 2022-present INESC-ID.
+// Distributed under the MIT license that can be found in the LICENSE file.
+
+use proc_macro::TokenStream;
+use quote::quote;
+use syn::parse::{Parse, ParseStream};
+use syn::{Block, Error, ExprClosure, Pat, Result, Stmt, Token, parse_macro_input};
+
+struct Lambda {
+    captures: Block,
+    closure: ExprClosure,
+}
+
+impl Parse for Lambda {
+    fn parse(input: ParseStream) -> Result<Self> {
+        let captures = input.parse()?;
+        input.parse::<Token![,]>()?;
+        let closure = input.parse()?;
+        input.parse::<Option<Token![,]>>()?;
+        Ok(Lambda { captures, closure })
+    }
+}
+
+fn expand_lambda(lambda: Lambda, is_unsafe: bool) -> Result<proc_macro2::TokenStream> {
+    let mut names = Vec::new();
+    let mut types = Vec::new();
+    let mut inits = Vec::new();
+    for stmt in &lambda.captures.stmts {
+        let Stmt::Local(local) = stmt else {
+            return Err(Error::new_spanned(stmt, "expected a `let` for a capture"));
+        };
+        let (Pat::Type(pat), Some(init)) = (&local.pat, &local.init) else {
+            return Err(Error::new_spanned(
+                local,
+                "expected `let <name>: <type> = <init>;` for a capture",
+            ));
+        };
+        names.push(&pat.pat);
+        types.push(&pat.ty);
+        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"));
+        };
+        params.push(pat);
+        param_types.push(&pat.ty);
+    }
+
+    let output = &lambda.closure.output;
+    let body = &lambda.closure.body;
+    let (receiver, body, constructor) = if is_unsafe {
+        (
+            quote! { &mut self },
+            quote! { unsafe { #body } },
+            quote! { from_lambda_unsafe },
+        )
+    } else {
+        (quote! { &self }, quote! { #body }, quote! { from_lambda })
+    };
+
+    Ok(quote! {{
+        struct __Lambda {
+            #(#names: #types,)*
+        }
+        impl __Lambda {
+            #[allow(unused_unsafe, clippy::too_many_arguments)]
+            fn call(#receiver #(, #params)*) #output {
+                #body
+            }
+        }
+        ::libcc2rs::FnPtr::<fn(#(#param_types),*) #output>::#constructor(
+            __Lambda { #(#names: #inits,)* },
+            __Lambda::call,
+        )
+    }})
+}
+
+pub fn expand(input: TokenStream, is_unsafe: bool) -> TokenStream {
+    let lambda = parse_macro_input!(input as Lambda);
+    expand_lambda(lambda, is_unsafe)
+        .unwrap_or_else(Error::into_compile_error)
+        .into()
+}

libcc2rs-macros/src/lib.rs

@@ -6,6 +6,7 @@ use proc_macro::TokenStream;
 mod byte_repr;
 mod fn_ptr_arg;
 mod goto;
+mod lambda;
 mod record;
 mod state_machine;
 mod switch;
@@ -76,6 +77,16 @@ pub fn goto_block(input: TokenStream) -> TokenStream {
     goto::expand(input)
 }
 
+#[proc_macro]
+pub fn lambda(input: TokenStream) -> TokenStream {
+    lambda::expand(input, false)
+}
+
+#[proc_macro]
+pub fn lambda_unsafe(input: TokenStream) -> TokenStream {
+    lambda::expand(input, true)
+}
+
 #[proc_macro]
 pub fn goto(_input: TokenStream) -> TokenStream {
     quote::quote! {

libcc2rs/src/fn_ptr.rs

@@ -4,11 +4,11 @@
 #![allow(private_bounds)]
 
 use std::any::{Any, TypeId};
-use std::cell::{RefCell, UnsafeCell};
+use std::cell::UnsafeCell;
 use std::rc::Rc;
 
 use crate::fn_ptr_arg::{ArgList, ArgRepr, FnPtrArg, FnPtrArgs, record_from_repr};
-use crate::rc::{AsPointer, Ptr};
+use crate::rc::Ptr;
 use crate::reinterpret::ByteRepr;
 use crate::void::{AnyPtr, ErasedPtr};
 
@@ -191,33 +191,33 @@ macro_rules! impl_fn_ptr_call {
             }
 
             #[allow(non_snake_case)]
-            pub fn with_captures<Captures: 'static>(
-                captures: Captures,
-                body: fn(Ptr<Captures> $(, $a)*) -> R,
+            pub fn from_lambda<Lambda: 'static>(
+                lambda: Lambda,
+                call: fn(&Lambda $(, $a)*) -> R,
             ) -> Self {
-                let captures = Rc::new(RefCell::new(captures));
+                let lambda = Rc::new(lambda);
                 FnPtr {
-                    addr: body as usize,
+                    addr: call as usize,
                     current: None,
                     original: None,
                     lambda: Some(Rc::new(move |($($a,)*): ($($a,)*)| {
-                        body(captures.as_pointer() $(, $a)*)
+                        call(&lambda $(, $a)*)
                     })),
                 }
             }
 
             #[allow(non_snake_case)]
-            pub fn with_captures_unsafe<Captures: 'static>(
-                captures: Captures,
-                body: fn(&mut Captures $(, $a)*) -> R,
+            pub fn from_lambda_unsafe<Lambda: 'static>(
+                lambda: Lambda,
+                call: fn(&mut Lambda $(, $a)*) -> R,
             ) -> Self {
-                let captures = Rc::new(UnsafeCell::new(captures));
+                let lambda = Rc::new(UnsafeCell::new(lambda));
                 FnPtr {
-                    addr: body as usize,
+                    addr: call as usize,
                     current: None,
                     original: None,
                     lambda: Some(Rc::new(move |($($a,)*): ($($a,)*)| {
-                        body(unsafe { &mut *captures.get() } $(, $a)*)
+                        call(unsafe { &mut *lambda.get() } $(, $a)*)
                     })),
                 }
             }

libcc2rs/src/lib.rs

@@ -69,4 +69,6 @@ pub use fd::*;
 mod format;
 pub use format::*;
 
-pub use libcc2rs_macros::{ByteRepr, FnPtrArg, Record, VaArg, goto, goto_block, switch};
+pub use libcc2rs_macros::{
+    ByteRepr, FnPtrArg, Record, VaArg, goto, goto_block, lambda, lambda_unsafe, switch,
+};