feat(zkvms_host_io): Add the "Return" type, which only contains the function return type

Currently we create an "Output" type, which contains a tuple with all
public inputs and the function return type. In most zkVMs, public and
private inputs are not distinguished from the start, public inputs are
returned by the guest program and private are not.

Thus, most of the time we need a more complicated return type. However,
sometimes this is not the case, and the public/private distinction is
done upfront (and we don't need to output inputs). So we need to add the
appropriate type.

1 files changed, 5 insertions, 1 deletions

diff --git a/zkvms_host_io/input_macros/src/lib.rs b/zkvms_host_io/input_macros/src/lib.rs
index 3deffef..72c4d11 100644
--- a/zkvms_host_io/input_macros/src/lib.rs
+++ b/zkvms_host_io/input_macros/src/lib.rs
@@ -36,6 +36,8 @@ static DERIVES: &str = "#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deseria
 /// ```rust
 /// pub type Output = (... ...);
 ///
+/// pub type Return = ...;
+///
 /// #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
 /// pub struct PublicInput {
 ///     pub ...: ...,
@@ -86,6 +88,8 @@ pub fn generate_output_type_input_struct(_: TokenStream) -> TokenStream {
         .collect::<String>();
     let output_type = format!("pub type Output = ({} {});", public_types, ret).to_string();
 
+    let return_type = format!("pub type Return = {};", ret).to_string();
+
     let (public_args, private_args) = args_split_public(&args, &public_inputs.keys().collect());
     let public_attrs = public_args
         .iter()
@@ -125,7 +129,7 @@ pub fn generate_output_type_input_struct(_: TokenStream) -> TokenStream {
     }
     struct_def += ") } }";
 
-    (output_type + &public_input_type + &private_input_type + &struct_def)
+    (output_type + &return_type + &public_input_type + &private_input_type + &struct_def)
         .parse::<TokenStream>()
         .unwrap()
 }