I think it might be possible. The placement by return rfc describes a transformation similar to generators that would allow some functions to return a DST. That transformation could be generalized to allow (almost) any function to return a DST.
The function would be transformed into a state machine, storing the locals that are alive for each state (like with generators/futures). Locals that are dynamically sized would be stored as a (move) reference pointing to memory allocated by the calling function (or one of its ancestors). By storing a reference the function state remains Sized
. Each time the function requires memory for a DST (for one of its locals, the return value, etc.) it would yield control to the calling function returning the required Layout
. The caller may then either allocate the memory and resume the function or yield the request to its caller.
For trait functions that may return DSTs (they would be required to opt-in for functions that return associated types), the vtable for the trait would include the Layout
of the state for the generator
(which is Sized
since all dynamically sized locals are stored as references) and the relevant function pointers.
The output of such a transformation might look like:
#![feature(
const_mut_refs,
const_fn_trait_bound,
const_fn_fn_ptr_basics,
unsize,
generic_associated_types
)]
use core::{
alloc::Layout,
cell::{Cell, UnsafeCell},
fmt::Debug,
marker::{PhantomData, Unsize},
mem::{self, MaybeUninit},
ops,
pin::Pin,
ptr, slice,
};
pub enum FnState<T> {
SpaceNeeded(Layout),
Complete(T),
}
pub trait GeneratorFn<Args>: Sized {
type State<'r>;
type Output: ?Sized;
fn start<'r>(args: Args) -> (Self::State<'r>, Layout);
/// # Panics
/// If a previous call to this function or [`start`] either returned `FnState::Complete` or paniced, then
/// this function may panic.
///
/// # Safety
/// If this is the first call to this function or if the last call to to this function returned
/// `FnState::SpaceNeeded(layout)` then the caller **must** pass a slice meeting the
/// requirements specified by `layout`.
unsafe fn resume<'r>(
this: Pin<&mut Self::State<'r>>, space: &'r mut [MaybeUninit<u8>],
) -> FnState<MoveRef<'r, Self::Output>>;
}
// Vtable for fn(Args) -> R where R is a dst
pub struct GeneratorFnVtable<Args, R: ?Sized> {
layout: Layout,
drop: unsafe fn(*mut ()),
start: for<'r> unsafe fn(Args, &'r mut [MaybeUninit<u8>]) -> (*mut (), Layout),
resume: for<'r> unsafe fn(*mut (), &'r mut [MaybeUninit<u8>]) -> FnState<MoveRef<'r, R>>,
}
impl<Args, R: ?Sized> Copy for GeneratorFnVtable<Args, R> {}
impl<Args, R: ?Sized> Clone for GeneratorFnVtable<Args, R> {
fn clone(&self) -> Self {
*self
}
}
impl<Args, R: ?Sized> GeneratorFnVtable<Args, R> {
pub const fn new<F>() -> Self
where F: GeneratorFn<Args, Output = R> {
Self {
layout: Layout::new::<F::State<'_>>(),
drop: |p| unsafe { p.cast::<F::State<'_>>().drop_in_place() },
start: |args, space| unsafe {
let (s, l) = F::start(args);
let p = space.as_mut_ptr().cast::<F::State<'_>>();
p.write(s);
(p.cast::<()>(), l)
},
resume: |p, space| unsafe {
F::resume(Pin::new_unchecked(&mut *p.cast::<F::State<'_>>()), space)
},
}
}
pub fn layout(&self) -> Layout {
self.layout
}
pub fn start<'r>(
&self, args: Args, slot: &'r mut [MaybeUninit<u8>],
) -> (DynGeneratorFnState<'r, Args, R>, Layout) {
assert!(slot.len() >= self.layout.size());
assert_eq!((slot.as_ptr() as usize) % self.layout.align(), 0);
let (state, layout) = unsafe { (self.start)(args, slot) };
(
DynGeneratorFnState {
state,
vtable: *self,
_marker: PhantomData,
},
layout,
)
}
}
pub struct DynGeneratorFnState<'r, Args, R: ?Sized> {
vtable: GeneratorFnVtable<Args, R>,
state: *mut (),
_marker: PhantomData<&'r mut ()>,
}
impl<Args, R: ?Sized> Drop for DynGeneratorFnState<'_, Args, R> {
fn drop(&mut self) {
unsafe {
(self.vtable.drop)(self.state);
}
}
}
impl<'r, Args, R: ?Sized> DynGeneratorFnState<'r, Args, R> {
pub fn resume(&mut self, space: &'r mut [MaybeUninit<u8>]) -> FnState<MoveRef<'r, R>> {
unsafe { (self.vtable.resume)(self.state, space) }
}
}
pub struct MoveRef<'a, T: ?Sized>(&'a mut T);
impl<T: ?Sized> Drop for MoveRef<'_, T> {
fn drop(&mut self) {
unsafe {
ptr::drop_in_place(self.0);
}
}
}
impl<T: ?Sized> ops::Deref for MoveRef<'_, T> {
type Target = T;
fn deref(&self) -> &Self::Target {
&*self.0
}
}
impl<T: ?Sized> ops::DerefMut for MoveRef<'_, T> {
fn deref_mut(&mut self) -> &mut Self::Target {
&mut *self.0
}
}
impl<'a, T: ?Sized> MoveRef<'a, T> {
pub fn place<U: 'a>(slot: &'a mut [MaybeUninit<u8>], value: U) -> Self
where U: Unsize<T> {
assert!(slot.len() >= mem::size_of::<U>());
assert_eq!((slot.as_ptr() as usize) % mem::align_of::<U>(), 0);
let p = slot.as_mut_ptr().cast::<U>();
// SAFETY: This is safe because the size and alignment of slot were just checked.
unsafe {
p.write(value);
Self(&mut *p)
}
}
}
pub struct Foo;
pub enum FooState {
Start,
Complete,
}
// fn foo() -> [u32] {
// [0, 1, 2, 3, 4]
// }
impl GeneratorFn<()> for Foo {
type State<'r> = FooState;
type Output = [u32];
fn start<'r>(_args: ()) -> (Self::State<'r>, Layout) {
(FooState::Start, Layout::new::<[u32; 5]>())
}
unsafe fn resume<'r>(
mut this: Pin<&mut Self::State<'r>>, space: &'r mut [MaybeUninit<u8>],
) -> FnState<MoveRef<'r, Self::Output>> {
match this.as_mut().get_mut() {
FooState::Start => {
*this.get_mut() = FooState::Complete;
FnState::Complete(MoveRef::place(space, [0, 1, 2, 3, 4]))
}
FooState::Complete => panic!("resume called after complete"),
}
}
}
pub struct Bar;
pub enum BarState<'r> {
Start(GeneratorFnVtable<(), [u32]>),
WaitForSpace1(DynGeneratorFnState<'r, (), [u32]>),
WaitForSpace2(MoveRef<'r, [u32]>),
Complete,
}
// fn bar(f: fn() -> [u32]) -> dyn Debug {
// f().iter().sum::<u32>()
// }
impl GeneratorFn<(GeneratorFnVtable<(), [u32]>,)> for Bar {
type State<'r> = BarState<'r>;
type Output = dyn Debug;
fn start<'r>((args,): (GeneratorFnVtable<(), [u32]>,)) -> (Self::State<'r>, Layout) {
let layout = args.layout;
(BarState::Start(args), layout)
}
unsafe fn resume<'r>(
mut this: Pin<&mut Self::State<'r>>, space: &'r mut [MaybeUninit<u8>],
) -> FnState<MoveRef<'r, Self::Output>> {
match this.as_mut().get_mut() {
BarState::Start(f) => {
let (s, l) = f.start((), space);
*this.get_mut() = BarState::WaitForSpace1(s);
FnState::SpaceNeeded(l)
}
BarState::WaitForSpace1(s) => match s.resume(space) {
FnState::SpaceNeeded(l) => FnState::SpaceNeeded(l),
FnState::Complete(r) => {
*this.get_mut() = BarState::WaitForSpace2(r);
FnState::SpaceNeeded(Layout::new::<u32>())
}
},
BarState::WaitForSpace2(s) => {
let sum = s.iter().sum::<u32>();
*this.get_mut() = BarState::Complete;
FnState::Complete(MoveRef::place(space, sum))
}
BarState::Complete => panic!("resume called after complete"),
}
}
}
// Wrap a function that returns a non-dst in a generator
// pub struct RegularFn<F, R>(PhantomData<(F, R)>);
// pub enum RegularFnState<R> {
// Start(R),
// Complete,
// }
//
// impl<Args, R: Unpin, F> GeneratorFn<Args> for RegularFn<F, R>
// where F: Default + Fn(Args) -> R
// {
// type State<'r> = RegularFnState<R>;
// type Output = R;
//
// fn start<'r>(args: Args) -> (Self::State<'r>, Layout) {
// let r = (F::default())(args);
// (RegularFnState::Start(r), Layout::new::<R>())
// }
//
// unsafe fn resume<'r>(
// this: Pin<&mut Self::State<'r>>, space: &'r mut [MaybeUninit<u8>],
// ) -> FnState<MoveRef<'r, Self::Output>> {
// match mem::replace(this.as_mut().get_mut(), RegularFnState::Complete) {
// RegularFnState::Start(r) => FnState::Complete(MoveRef::place(space, r)),
// RegularFnState::Complete => panic!("resume called after complete"),
// }
// }
// }
// pub trait FutureDst {
// type Output: ?Sized;
//
// fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output>;
// }
// pub struct FutureVtable {
// layout: Layout,
// drop: unsafe fn(*mut ()),
// poll: GeneratorFnVtable<(*mut (), *mut Context<'static>), Poll<dyn Drop>>,
// }
struct Arena {
space: [UnsafeCell<MaybeUninit<u8>>; 16 * 1024],
// INVARIANT: next is always <= space.len()
next: Cell<usize>,
}
impl Arena {
fn new() -> Self {
Self {
space: unsafe { MaybeUninit::uninit().assume_init() },
next: Cell::new(0),
}
}
fn allocate(&self, layout: Layout) -> &mut [MaybeUninit<u8>] {
let next = self.next.get();
let remaining = self.space.len() - next;
// SAFETY: By the invariant of Arena, next is always within the bounds of space.
let p = unsafe { (*self.space.as_ptr().add(next)).get() };
let start = p.align_offset(layout.align());
let total = start.checked_add(layout.size()).expect("stack overflow");
assert!(total <= remaining, "stack overflow");
self.next.set(next + total);
// SAFETY: The bounds were just checked.
unsafe { slice::from_raw_parts_mut(p.add(start).cast(), layout.size()) }
}
}
// fn foo_bar(f: fn(fn() -> [u32]) -> dyn Debug) {
// println!("{:?}", f(foo));
// }
fn foo_bar(f: GeneratorFnVtable<(GeneratorFnVtable<(), [u32]>,), dyn Debug>) {
let arena = Arena::new();
let r = 'outer: loop {
let (mut s, mut l) = f.start(
(GeneratorFnVtable::new::<Foo>(),),
arena.allocate(f.layout()),
);
loop {
match s.resume(arena.allocate(l)) {
FnState::Complete(r) => break 'outer r,
FnState::SpaceNeeded(n) => {
l = n;
}
}
}
};
println!("{:?}", &*r);
}
fn main() {
foo_bar(GeneratorFnVtable::new::<Bar>());
}
(playground)
Note: This is just a proof of concept and there may be safety/lifetime issues. The transformation uses the memory allocated for both dynamically sized locals and the return value, which doesn't meet the needs that motivated the placement by return RFC. The transformation would need to be modified, so that each request for memory would include the purpose for the memory (storing locals vs the return value) allowing the caller to modify its allocation strategy.