#![no_std]
extern crate alloc;
pub const ALIGN: usize = core::mem::align_of::<crate::private::InlineStorage<1>>();
pub const DEFAULT_INLINE_SIZE: usize = core::mem::size_of::<[usize; 3]>();
#[doc(hidden)]
pub mod private {
pub use alloc::boxed::Box;
pub use core::{
mem::{align_of, size_of, transmute, ManuallyDrop, MaybeUninit},
ptr::{copy_nonoverlapping, drop_in_place, read, NonNull},
};
pub trait Closure {
type Inner;
}
#[repr(C)]
pub struct VTable<D = unsafe fn(), C = unsafe fn()> {
pub drop: D,
pub call: C,
}
#[repr(C, align(16))]
#[derive(Clone, Copy)]
pub struct InlineStorage<const INLINE_SIZE: usize> {
pub storage: MaybeUninit<[u8; INLINE_SIZE]>,
}
pub type StoragePtr<const INLINE_SIZE: usize> = NonNull<InlineStorage<INLINE_SIZE>>;
pub type DropFn<const INLINE_SIZE: usize> = unsafe fn(StoragePtr<INLINE_SIZE>);
}
#[doc(hidden)]
#[macro_export]
macro_rules! private_tiny_fn {
(@call $(<$($hrl:lifetime),+>)? Fn $(< $($lt:lifetime,)* $($t:ident,)* >)? ($($arg_name:ident: $arg_type:ty),* $(,)?) $( -> $ret:ty)?) => {
fn call $(< $($hrl),+ >)? (&self, $($arg_name: $arg_type),*) $(-> $ret)? {
unsafe {
match self.vtable {
None => (*self.payload.boxed.closure)($($arg_name),*),
Some(vtable) => {
let call_fn: CallFn< $($($lt,)* $($t,)*)? INLINE_SIZE> = $crate::private::transmute(vtable.call);
call_fn(
$crate::private::NonNull::from(&self.payload.inline),
$($arg_name),*
)
}
}
}
}
};
(@call $(<$($hrl:lifetime),+>)? FnMut $(< $($lt:lifetime,)* $($t:ident,)* >)? ($($arg_name:ident: $arg_type:ty),* $(,)?) $( -> $ret:ty)?) => {
fn call $(< $($hrl),+ >)? (&mut self, $($arg_name: $arg_type),*) $(-> $ret)? {
unsafe {
match self.vtable {
None => (*(*self.payload.boxed).closure)($($arg_name),*),
Some(vtable) => {
let call_fn: CallFn< $($($lt,)* $($t,)*)? INLINE_SIZE> = $crate::private::transmute(vtable.call);
call_fn(
$crate::private::NonNull::from(&mut self.payload.inline),
$($arg_name),*
)
}
}
}
}
};
(@call $(<$($hrl:lifetime),+>)? FnOnce $(< $($lt:lifetime,)* $($t:ident,)* >)? ($($arg_name:ident: $arg_type:ty),* $(,)?) $( -> $ret:ty)?) => {
fn call $(< $($hrl),+ >)? (self, $($arg_name: $arg_type),*) $(-> $ret)? {
let mut me = $crate::private::ManuallyDrop::new(self);
unsafe {
match me.vtable {
None => ($crate::private::ManuallyDrop::take(&mut me.payload.boxed).closure)($($arg_name),*),
Some(vtable) => {
let call_fn: CallFn< $($($lt,)* $($t,)*)? INLINE_SIZE> = $crate::private::transmute(vtable.call);
call_fn(
$crate::private::NonNull::from(&mut (*me).payload.inline),
$($arg_name),*
)
}
}
}
}
};
(@call_outer $(<$($hrl:lifetime),+>)? Fn $(< $($lt:lifetime,)* $($t:ident,)* >)? ($($arg_name:ident: $arg_type:ty),* $(,)?) $( -> $ret:ty)?) => {
#[allow(dead_code)]
pub fn call $(<$($hrl),+>)? (&self, $($arg_name: $arg_type),*) $(-> $ret)? {
self.inner.call($($arg_name,)*)
}
};
(@call_outer $(<$($hrl:lifetime),+>)? FnMut $(< $($lt:lifetime,)* $($t:ident,)* >)? ($($arg_name:ident: $arg_type:ty),* $(,)?) $( -> $ret:ty)?) => {
#[allow(dead_code)]
pub fn call $(<$($hrl),+>)? (&mut self, $($arg_name: $arg_type),*) $(-> $ret)? {
self.inner.call($($arg_name,)*)
}
};
(@call_outer $(<$($hrl:lifetime),+>)? FnOnce $(< $($lt:lifetime,)* $($t:ident,)* >)? ($($arg_name:ident: $arg_type:ty),* $(,)?) $( -> $ret:ty)?) => {
#[allow(dead_code)]
pub fn call $(<$($hrl),+>)? (self, $($arg_name: $arg_type),*) $(-> $ret)? {
self.inner.call($($arg_name,)*)
}
};
(@inline_call_cast Fn $ptr:ident) => {
(*(*$ptr.as_ptr()).storage.as_ptr().cast::<F>())
};
(@inline_call_cast FnMut $ptr:ident) => {
(*(*$ptr.as_ptr()).storage.as_mut_ptr().cast::<F>())
};
(@inline_call_cast FnOnce $ptr:ident) => {
$crate::private::read((*$ptr.as_ptr()).storage.as_mut_ptr().cast::<F>())
};
}
#[macro_export]
macro_rules! tiny_fn {
($(
$(#[$meta:meta])*
$vis:vis struct $name:ident $(< $($lt:lifetime),* $(,)? $($t:ident),* >)? =
$(<$($hrl:lifetime),+>)?
$fun:ident ($(
$arg_name:ident: $arg_type:ty
),* $(,)?)
$( -> $ret:ty)?
$(| $(+ $markers:path)+)?
$(where $( $wt:tt: $(+ $wtl:lifetime)* $(+ $wtb:path)* ),*)?
;
)*) => {
$(
const _: () = {
type CallFn< $($($lt,)* $($t,)*)? const INLINE_SIZE: usize> = $(for<$($hrl),+>)? unsafe fn($crate::private::StoragePtr<INLINE_SIZE>, $($arg_type),*) $( -> $ret)?;
struct BoxedFn <'closure $( $(, $lt)* $(, $t)*)?>
$(where $( $($wt: $wtl,)* $($wt: $wtb,)* ),*)?
{
closure: $crate::private::Box<dyn $(for<$($hrl),+>)? $fun($($arg_type),*) $(-> $ret)? $($(+ $markers)+)? + 'closure>,
}
#[doc(hidden)]
union TinyClosurePayload<'closure, $($($lt,)* $($t,)*)? const INLINE_SIZE: usize>
$(where $( $($wt: $wtl,)* $($wt: $wtb,)* ),*)?
{
inline: $crate::private::InlineStorage<INLINE_SIZE>,
boxed: $crate::private::ManuallyDrop<BoxedFn<'closure $($(, $lt)* $(, $t)*)?>>,
}
pub struct TinyClosure<'closure, $($($lt,)* $($t,)*)? const INLINE_SIZE: usize>
$(where $( $($wt: $wtl,)* $($wt: $wtb,)* ),*)?
{
vtable: Option<&'static $crate::private::VTable>,
payload: TinyClosurePayload<'closure, $($($lt,)* $($t,)*)? INLINE_SIZE>,
}
impl<'closure, $($($lt,)* $($t,)*)? const INLINE_SIZE: usize> Drop for TinyClosure<'closure, $($($lt,)* $($t,)*)? INLINE_SIZE>
$(where $( $($wt: $wtl,)* $($wt: $wtb,)* ),*)?
{
fn drop(&mut self) {
unsafe {
match self.vtable {
None => $crate::private::ManuallyDrop::drop(&mut self.payload.boxed),
Some(vtable) => {
let drop_fn: $crate::private::DropFn<INLINE_SIZE> = $crate::private::transmute(vtable.drop);
drop_fn($crate::private::NonNull::from(&mut self.payload.inline));
}
}
}
}
}
impl<'closure, $($($lt,)* $($t,)*)? const INLINE_SIZE: usize> TinyClosure<'closure, $($($lt,)* $($t,)*)? INLINE_SIZE>
$(where $( $($wt: $wtl,)* $($wt: $wtb,)* ),*)?
{
fn new<F>(f: F) -> Self
where
F: $(for<$($hrl),+>)? $fun ($($arg_type),*) $( -> $ret)? $($(+ $markers)+)? + 'closure,
{
let size_fits = $crate::private::size_of::<F>() <= INLINE_SIZE;
let align_fits = $crate::private::align_of::<F>() <= $crate::ALIGN;
if size_fits && align_fits {
let mut inline = $crate::private::InlineStorage {
storage: $crate::private::MaybeUninit::uninit(),
};
unsafe {
let storage_f = inline.storage.as_mut_ptr() as *mut $crate::private::MaybeUninit<F>;
(*storage_f).write(f);
}
let vtable = unsafe {
$crate::private::transmute(&$crate::private::VTable::< $crate::private::DropFn<INLINE_SIZE>, CallFn< $($($lt,)* $($t,)*)? INLINE_SIZE>> {
drop: |ptr: $crate::private::StoragePtr<INLINE_SIZE>| {
$crate::private::drop_in_place(ptr.cast::<F>().as_ptr());
},
call: |ptr: $crate::private::StoragePtr<INLINE_SIZE> $(, $arg_name)*| {
$crate::private_tiny_fn!(@inline_call_cast $fun ptr)($($arg_name),*)
}
})
};
TinyClosure {
vtable: Some(vtable),
payload: TinyClosurePayload {
inline,
}
}
} else {
let boxed_fn = BoxedFn {
closure: $crate::private::Box::new(f),
};
TinyClosure {
vtable: None,
payload: TinyClosurePayload {
boxed: $crate::private::ManuallyDrop::new(boxed_fn),
}
}
}
}
$crate::private_tiny_fn!(@call $(<$($hrl),+>)? $fun $(< $($lt,)* $($t,)* >)? ($($arg_name: $arg_type),*) $( -> $ret)?);
}
impl<'closure, $($($lt,)* $($t,)*)? const INLINE_SIZE: usize> $crate::private::Closure for $name<'closure, $($($lt,)* $($t,)*)? INLINE_SIZE>
$(where $( $($wt: $wtl,)* $($wt: $wtb,)* ),*)?
{
type Inner = TinyClosure<'closure, $($($lt,)* $($t,)*)? INLINE_SIZE>;
}
};
#[repr(transparent)]
$(#[$meta])*
$vis struct $name<'closure, $($($lt,)* $($t,)*)? const INLINE_SIZE: usize = { $crate::DEFAULT_INLINE_SIZE }>
$(where $( $($wt: $wtl,)* $($wt: $wtb,)* ),*)?
{
#[allow(dead_code)]
inner: < $name<'closure, $($($lt,)* $($t,)*)? INLINE_SIZE> as $crate::private::Closure>::Inner,
}
impl<'closure, $($($lt,)* $($t,)*)? const INLINE_SIZE: usize> $name<'closure, $($($lt,)* $($t,)*)? INLINE_SIZE>
$(where $( $($wt: $wtl,)* $($wt: $wtb,)* ),*)?
{
#[inline]
pub fn new<F>(f: F) -> Self
where
F: $(for<$($hrl),+>)? $fun ($($arg_type),*) $( -> $ret)? $($(+ $markers)+)? + 'closure,
{
$name {
inner: < < $name < $($($t,)*)? INLINE_SIZE > as $crate::private::Closure>::Inner>::new(f),
}
}
$crate::private_tiny_fn!(@call_outer $(<$($hrl),+>)? $fun $(< $($lt,)* $($t,)* >)? ($($arg_name: $arg_type),*) $( -> $ret)?);
}
)*
};
}
#[cfg(feature = "example")]
pub mod example {
tiny_fn! {
pub struct BinOp<T> = Fn(a: T, b: T) -> T;
pub struct MakeString = Fn() -> alloc::string::String;
}
}
#[cfg(test)]
mod tests {
use alloc::string::String;
#[test]
fn test() {
use alloc::rc::Rc;
tiny_fn! {
pub(crate) struct Foo<T> = FnMut(a: &T, b: T) -> T | + Send;
pub struct Bar = FnOnce(b: u8) -> alloc::string::String;
pub struct Baz<'a> = Fn(a: &'a str) -> &'a str | + Send;
pub struct Hrl = <'a> Fn(a: &'a str) -> &'a str | + Send;
struct Complex<'a, A, B> = <'b> Fn(a: &'a A, b: &'b B) -> &'a str | + Send where A: + 'a;
}
let mut x = 3;
let mut foo = Foo::<u8>::new(|a, b| {
x += a + b;
x
});
let bar: Bar<0> = Bar::new(|b| alloc::format!("{}", b));
assert_eq!(6, foo.call(&1, 2));
assert_eq!(13, foo.call(&3, 4));
assert_eq!("3", bar.call(3));
fn assert_send<T: Send>() {}
assert_send::<Foo<Rc<u8>>>();
let s = String::new();
let baz: Baz = Baz::new(|a: &str| a);
let _: &'static str = baz.call("Hello, World!");
let mut hlp: Hrl = Hrl::new(|a: &str| a);
let _: &'static str = hlp.call("Hello, World!");
hlp = Hrl::new(|_| "foo");
let _: &str = hlp.call(&s);
let complex: Complex<u8, u8> = Complex::new(|a: &u8, b: &u8| {
if *a > *b {
"a is greater"
} else {
"b is greater"
}
});
}
}