@@ -308,13 +308,13 @@ fn generate_vm_wrapper(
308308
309309 Ok ( quote ! {
310310 #[ allow( dead_code) ]
311- pub ( super ) fn #wrapper_name( #( #imm_wrapper_params) , * ) -> #wrapper_output {
311+ pub ( crate ) fn #wrapper_name( #( #imm_wrapper_params) , * ) -> #wrapper_output {
312312 #( #imm_extract_stmts) *
313313 #call_expr
314314 }
315315
316316 #[ allow( dead_code) ]
317- pub ( super ) fn #mutable_wrapper_name( #( #mut_wrapper_params) , * ) -> #wrapper_output {
317+ pub ( crate ) fn #mutable_wrapper_name( #( #mut_wrapper_params) , * ) -> #wrapper_output {
318318 #( #mut_extract_stmts) *
319319 #call_expr
320320 }
@@ -344,7 +344,7 @@ fn generate_async_vm_wrapper(
344344 let ty = & pat_type. ty ;
345345 if is_host_context_param ( input) {
346346 extract_stmts. push ( quote ! {
347- let #ident = <#ty as super :: CaptureAsyncHostContext >:: capture ( vm) ?;
347+ let #ident = <#ty as super :: CaptureAsyncHostContext >:: capture_with_args ( vm, args ) ?;
348348 } ) ;
349349 call_args. push ( quote ! ( #ident) ) ;
350350 continue ;
@@ -366,12 +366,14 @@ fn generate_async_vm_wrapper(
366366 } else {
367367 quote ! ( #impl_name( #( #call_args) , * ) . await )
368368 } ;
369- let body = quote ! {
370- # ( #extract_stmts ) *
371- vm . submit_host_future ( Box :: pin ( async move {
372- let value = #await_value ;
369+ let future_result = if return_is_host_future_output ( & item . sig . output ) {
370+ quote ! ( Ok ( value . map ( super :: return_one ) ) )
371+ } else {
372+ quote ! {
373373 match super :: IntoHostCallOutcome :: into_host_call_outcome( value) {
374- super :: CallOutcome :: Return ( values) => Ok ( values) ,
374+ super :: CallOutcome :: Return ( values) => {
375+ Ok ( super :: HostFutureOutput :: returning( values) )
376+ }
375377 super :: CallOutcome :: Pending ( op_id) => Err ( super :: VmError :: HostError (
376378 format!( "async host function returned nested pending operation {op_id}" ) ,
377379 ) ) ,
@@ -381,20 +383,27 @@ fn generate_async_vm_wrapper(
381383 ) ,
382384 ) ,
383385 }
386+ }
387+ } ;
388+ let body = quote ! {
389+ #( #extract_stmts) *
390+ vm. submit_host_future( Box :: pin( async move {
391+ let value = #await_value;
392+ #future_result
384393 } ) )
385394 } ;
386395
387396 Ok ( quote ! {
388397 #[ allow( dead_code) ]
389- pub ( super ) fn #wrapper_name(
398+ pub ( crate ) fn #wrapper_name(
390399 vm: & mut super :: super :: Vm ,
391400 args: & [ super :: super :: Value ] ,
392401 ) -> super :: super :: VmResult <super :: CallOutcome > {
393402 #body
394403 }
395404
396405 #[ allow( dead_code) ]
397- pub ( super ) fn #mutable_wrapper_name(
406+ pub ( crate ) fn #mutable_wrapper_name(
398407 vm: & mut super :: super :: Vm ,
399408 args: & mut [ super :: super :: Value ] ,
400409 ) -> super :: super :: VmResult <super :: CallOutcome > {
@@ -465,6 +474,20 @@ fn unwrap_vm_result_type(ty: &Type) -> Result<Option<Type>, Error> {
465474 }
466475}
467476
477+ fn return_is_host_future_output ( output : & ReturnType ) -> bool {
478+ vm_result_inner_type ( output)
479+ . expect ( "pd_host_function return type should already be validated" )
480+ . and_then ( |ty| match ty {
481+ Type :: Path ( path) => path
482+ . path
483+ . segments
484+ . last ( )
485+ . map ( |segment| segment. ident . clone ( ) ) ,
486+ _ => None ,
487+ } )
488+ . is_some_and ( |ident| ident == "HostFutureOutput" )
489+ }
490+
468491fn return_is_vm_result ( output : & ReturnType ) -> bool {
469492 vm_result_inner_type ( output)
470493 . expect ( "pd_host_function return type should already be validated" )
@@ -526,7 +549,7 @@ fn type_label(ty: &Type) -> Result<String, Error> {
526549 let inner_label = type_label ( inner) ?;
527550 Ok ( format ! ( "{inner_label} | null" ) )
528551 }
529- "VmResult" | "HostCallResult" => {
552+ "VmResult" | "HostCallResult" | "HostFutureOutput" => {
530553 let syn:: PathArguments :: AngleBracketed ( args) = & segment. arguments else {
531554 return Err ( Error :: new_spanned (
532555 & segment. arguments ,
@@ -720,9 +743,26 @@ mod tests {
720743 assert ! ( expanded. contains( "async move" ) ) ;
721744 assert ! ( expanded. contains( "borrow_arg" ) ) ;
722745 assert ! ( expanded. contains( "CaptureAsyncHostContext" ) ) ;
746+ assert ! ( expanded. contains( "capture_with_args" ) ) ;
723747 assert ! ( !expanded. contains( "pd_host_context" ) ) ;
724748 }
725749
750+ #[ test]
751+ fn async_host_future_output_maps_its_inner_value_to_call_return ( ) {
752+ let attr: Punctuated < Meta , Token ! [ , ] > = parse_quote ! ( name = "test::completion" ) ;
753+ let item: ItemFn = parse_quote ! {
754+ /// Completes after mutating VM-owned state.
755+ async fn completion( ) -> VmResult <HostFutureOutput <i64 >> {
756+ todo!( )
757+ }
758+ } ;
759+
760+ let expanded = expand_pd_host_function ( attr, item)
761+ . expect ( "host future output should be accepted" )
762+ . to_string ( ) ;
763+ assert ! ( expanded. contains( "value . map (super :: return_one)" ) ) ;
764+ }
765+
726766 #[ test]
727767 fn async_signature_rejects_borrowed_parameters ( ) {
728768 let attr: Punctuated < Meta , Token ! [ , ] > = parse_quote ! ( name = "test::borrowed" ) ;
0 commit comments