Add stateful call leaf runtime

unshell-macros/src/lib.rs

@@ -1,9 +1,10 @@
 //! Proc macros for `unshell` application-layer leaf declarations.
 
 use proc_macro::TokenStream;
-use quote::quote;
+use quote::{format_ident, quote};
 use syn::{
-    DeriveInput, Error, Ident, LitStr, Result, Token, parse::Parse, parse_macro_input,
+    Attribute, DeriveInput, Error, FnArg, GenericArgument, Ident, ImplItem, ImplItemFn, ItemImpl,
+    LitStr, PatType, Result, ReturnType, Token, Type, TypePath, parse::Parse, parse_macro_input,
     punctuated::Punctuated,
 };
 
@@ -15,6 +16,17 @@ pub fn derive_leaf(input: TokenStream) -> TokenStream {
     }
 }
 
+#[proc_macro_attribute]
+pub fn procedures(attr: TokenStream, item: TokenStream) -> TokenStream {
+    match expand_procedures(
+        parse_macro_input!(attr as ProceduresAttributes),
+        parse_macro_input!(item as ItemImpl),
+    ) {
+        Ok(tokens) => tokens.into(),
+        Err(error) => error.to_compile_error().into(),
+    }
+}
+
 fn expand_leaf(input: DeriveInput) -> Result<proc_macro2::TokenStream> {
     let struct_name = input.ident;
     match input.data {
@@ -28,23 +40,8 @@ fn expand_leaf(input: DeriveInput) -> Result<proc_macro2::TokenStream> {
     };
 
     let parsed = LeafAttributes::parse_from(&input.attrs)?;
-    let procedures = parsed.procedures.clone().ok_or_else(|| {
-        Error::new_spanned(&struct_name, "missing #[leaf(procedures(...))] attribute")
-    })?;
-
-    if procedures.is_empty() {
-        return Err(Error::new_spanned(
-            &struct_name,
-            "leaf must declare at least one procedure suffix",
-        ));
-    }
-
     let (impl_generics, ty_generics, where_clause) = input.generics.split_for_impl();
     let leaf_name_expr = parsed.leaf_name_expression(&struct_name);
-    let procedure_suffix_literals = procedures
-        .iter()
-        .map(|procedure| LitStr::new(&procedure.to_string(), proc_macro2::Span::call_site()))
-        .collect::<Vec<_>>();
     let warning_note = parsed
         .explicit_id_value()
         .as_ref()
@@ -59,39 +56,17 @@ fn expand_leaf(input: DeriveInput) -> Result<proc_macro2::TokenStream> {
                 proc_macro2::Span::call_site(),
             )
         })
-        .map(|note| {
-            let attr = quote! { #[deprecated(note = #note)] };
-            (attr.clone(), attr.clone(), attr)
-        });
-    let (leaf_spec_warning_attr, procedure_warning_attr, leaf_name_warning_attr) =
-        warning_note.unwrap_or_else(|| (quote! {}, quote! {}, quote! {}));
+        .map(|note| quote! { #[deprecated(note = #note)] });
+    let leaf_name_warning_attr = warning_note.unwrap_or_else(|| quote! {});
 
     Ok(quote! {
         impl #impl_generics ::unshell::protocol::tree::ProtocolLeaf for #struct_name #ty_generics #where_clause {
             fn leaf_name() -> ::unshell::alloc::string::String {
                 #leaf_name_expr
             }
-
-            fn procedure_suffixes() -> &'static [&'static str] {
-                &[#(#procedure_suffix_literals),*]
-            }
         }
 
         impl #impl_generics #struct_name #ty_generics #where_clause {
-            /// Returns the canonical protocol leaf metadata for this type.
-            #leaf_spec_warning_attr
-            pub fn protocol_leaf_spec() -> ::unshell::protocol::tree::LeafSpec {
-                <Self as ::unshell::protocol::tree::ProtocolLeaf>::leaf_spec()
-            }
-
-            /// Resolves one local procedure suffix to its full canonical `procedure_id`.
-            #procedure_warning_attr
-            pub fn protocol_procedure_id(
-                suffix: &str,
-            ) -> ::core::option::Option<::unshell::alloc::string::String> {
-                <Self as ::unshell::protocol::tree::ProtocolLeaf>::procedure_id(suffix)
-            }
-
             /// Returns the canonical dotted leaf name declared for this type.
             #leaf_name_warning_attr
             pub fn protocol_leaf_name() -> ::unshell::alloc::string::String {
@@ -101,6 +76,315 @@ fn expand_leaf(input: DeriveInput) -> Result<proc_macro2::TokenStream> {
     })
 }
 
+fn expand_procedures(
+    attr: ProceduresAttributes,
+    mut item: ItemImpl,
+) -> Result<proc_macro2::TokenStream> {
+    let self_ty = item.self_ty.clone();
+    let impl_generics = item.generics.clone();
+    let (impl_generics_tokens, ty_generics, where_clause) = impl_generics.split_for_impl();
+    let error_ty = attr.error.ok_or_else(|| {
+        Error::new_spanned(
+            &item.self_ty,
+            "missing #[procedures(error = MyError)] attribute",
+        )
+    })?;
+
+    let mut dispatch_arms = Vec::new();
+
+    for impl_item in &mut item.items {
+        let ImplItem::Fn(method) = impl_item else {
+            continue;
+        };
+        if !take_call_attr(&mut method.attrs) {
+            continue;
+        }
+
+        dispatch_arms.push(expand_call_arm(method)?);
+    }
+
+    if dispatch_arms.is_empty() {
+        return Err(Error::new_spanned(
+            &item.self_ty,
+            "#[procedures] requires at least one #[call] method",
+        ));
+    }
+
+    let suffix_literals = dispatch_arms
+        .iter()
+        .map(|arm| arm.suffix_literal.clone())
+        .collect::<Vec<_>>();
+    let procedure_matches = dispatch_arms.iter().map(|arm| {
+        let suffix = &arm.suffix_literal;
+        quote! { #suffix => <Self as ::unshell::protocol::tree::CallProcedures>::procedure_id(#suffix), }
+    });
+    let dispatch_checks = dispatch_arms.iter().map(|arm| arm.dispatch_tokens.clone());
+
+    Ok(quote! {
+        #item
+
+        impl #impl_generics_tokens ::unshell::protocol::tree::CallProcedures for #self_ty #where_clause {
+            type Error = #error_ty;
+
+            fn procedure_suffixes() -> &'static [&'static str] {
+                &[#(#suffix_literals),*]
+            }
+
+            fn dispatch_call(
+                &mut self,
+                call: ::unshell::protocol::tree::IncomingCall,
+            ) -> ::core::result::Result<
+                ::unshell::protocol::tree::CallReply,
+                ::unshell::protocol::tree::DispatchError<Self::Error>,
+            > {
+                #(#dispatch_checks)*
+                unreachable!("protocol runtime validated local procedure dispatch")
+            }
+        }
+
+        impl #impl_generics_tokens #self_ty #ty_generics #where_clause {
+            /// Returns the canonical protocol leaf metadata for this type.
+            pub fn protocol_leaf_spec() -> ::unshell::protocol::tree::LeafSpec {
+                <Self as ::unshell::protocol::tree::CallProcedures>::leaf_spec()
+            }
+
+            /// Resolves one local procedure suffix to its full canonical `procedure_id`.
+            pub fn protocol_procedure_id(
+                suffix: &str,
+            ) -> ::core::option::Option<::unshell::alloc::string::String> {
+                match suffix {
+                    #(#procedure_matches)*
+                    _ => ::core::option::Option::None,
+                }
+            }
+        }
+    })
+}
+
+struct CallArm {
+    suffix_literal: LitStr,
+    dispatch_tokens: proc_macro2::TokenStream,
+}
+
+fn expand_call_arm(method: &ImplItemFn) -> Result<CallArm> {
+    let method_name = &method.sig.ident;
+    let suffix_literal = LitStr::new(&method_name.to_string(), method_name.span());
+    let call_id_expr = quote! {
+        <Self as ::unshell::protocol::tree::CallProcedures>::procedure_id(#suffix_literal)
+            .expect("generated procedure id must exist")
+    };
+
+    let inputs = method
+        .sig
+        .inputs
+        .iter()
+        .filter(|input| !matches!(input, FnArg::Receiver(_)))
+        .collect::<Vec<_>>();
+
+    let invocation = expand_invocation(method_name, &inputs)?;
+    let return_value = expand_return_conversion(&method.sig.output, quote! { __unshell_result })?;
+
+    Ok(CallArm {
+        suffix_literal: suffix_literal.clone(),
+        dispatch_tokens: quote! {
+            if call.message.procedure_id == #call_id_expr {
+                let __unshell_result = #invocation;
+                return { #return_value };
+            }
+        },
+    })
+}
+
+fn expand_invocation(method_name: &Ident, inputs: &[&FnArg]) -> Result<proc_macro2::TokenStream> {
+    if inputs.is_empty() {
+        return Ok(quote! { self.#method_name() });
+    }
+
+    if inputs.len() == 1 {
+        let FnArg::Typed(PatType { ty, .. }) = inputs[0] else {
+            return Err(Error::new_spanned(
+                inputs[0],
+                "unsupported receiver in procedure signature",
+            ));
+        };
+
+        if let Some(inner) = extract_call_inner_type(ty) {
+            return Ok(quote! {{
+                let __unshell_input = ::unshell::protocol::tree::decode_call_input::<#inner>(
+                    call.message.data.as_slice(),
+                )
+                .map_err(::unshell::protocol::tree::DispatchError::Decode)?;
+                let __unshell_call = ::unshell::protocol::tree::Call {
+                    input: __unshell_input,
+                    caller_path: call.header.src_path.clone(),
+                    procedure_id: call.message.procedure_id.clone(),
+                    dst_leaf: call.header.dst_leaf.clone(),
+                    response_hook: call
+                        .message
+                        .response_hook
+                        .as_ref()
+                        .map(|hook| ::unshell::protocol::tree::HookKey::new(
+                            hook.return_path.clone(),
+                            hook.hook_id,
+                        )),
+                };
+                self.#method_name(__unshell_call)
+            }});
+        }
+
+        return Ok(quote! {{
+            let __unshell_input = ::unshell::protocol::tree::decode_call_input::<#ty>(
+                call.message.data.as_slice(),
+            )
+            .map_err(::unshell::protocol::tree::DispatchError::Decode)?;
+            self.#method_name(__unshell_input)
+        }});
+    }
+
+    let tuple_types = inputs
+        .iter()
+        .map(|input| match input {
+            FnArg::Typed(PatType { ty, .. }) => Ok(ty.clone()),
+            other => Err(Error::new_spanned(
+                other,
+                "unsupported receiver in procedure signature",
+            )),
+        })
+        .collect::<Result<Vec<_>>>()?;
+    let vars = (0..tuple_types.len())
+        .map(|index| format_ident!("__unshell_arg_{index}"))
+        .collect::<Vec<_>>();
+
+    Ok(quote! {{
+        let (#(#vars),*) = ::unshell::protocol::tree::decode_call_input::<(#(#tuple_types),*)>(
+            call.message.data.as_slice(),
+        )
+        .map_err(::unshell::protocol::tree::DispatchError::Decode)?;
+        self.#method_name(#(#vars),*)
+    }})
+}
+
+fn expand_return_conversion(
+    return_type: &ReturnType,
+    value: proc_macro2::TokenStream,
+) -> Result<proc_macro2::TokenStream> {
+    match return_type {
+        ReturnType::Default => Ok(quote! {
+            let _ = #value;
+            ::core::result::Result::Ok(::unshell::protocol::tree::CallReply::NoReply)
+        }),
+        ReturnType::Type(_, ty) => normalize_output_type(ty, value),
+    }
+}
+
+fn normalize_output_type(
+    ty: &Type,
+    value: proc_macro2::TokenStream,
+) -> Result<proc_macro2::TokenStream> {
+    if is_unit_type(ty) {
+        return Ok(quote! {
+            let _ = #value;
+            ::core::result::Result::Ok(::unshell::protocol::tree::CallReply::NoReply)
+        });
+    }
+
+    if let Some(inner) = extract_outer_type_argument(ty, "CallResult") {
+        let inner_conversion = normalize_reply_value(inner, quote! { __unshell_value })?;
+        return Ok(quote! {
+            match #value {
+                ::unshell::protocol::tree::CallResult::Reply(__unshell_value) => {
+                    #inner_conversion
+                }
+                ::unshell::protocol::tree::CallResult::NoReply => {
+                    ::core::result::Result::Ok(::unshell::protocol::tree::CallReply::NoReply)
+                }
+            }
+        });
+    }
+
+    if let Some((ok_ty, _error_ty)) = extract_result_type_arguments(ty) {
+        let ok_conversion = normalize_output_type(ok_ty, quote! { __unshell_value })?;
+        return Ok(quote! {
+            match #value {
+                ::core::result::Result::Ok(__unshell_value) => { #ok_conversion }
+                ::core::result::Result::Err(__unshell_error) => {
+                    ::core::result::Result::Err(
+                        ::unshell::protocol::tree::DispatchError::Handler(__unshell_error)
+                    )
+                }
+            }
+        });
+    }
+
+    normalize_reply_value(ty, value)
+}
+
+fn normalize_reply_value(
+    _ty: &Type,
+    value: proc_macro2::TokenStream,
+) -> Result<proc_macro2::TokenStream> {
+    Ok(quote! {
+        ::core::result::Result::Ok(::unshell::protocol::tree::CallReply::Reply(
+            ::unshell::protocol::tree::encode_call_reply(&#value)
+                .map_err(::unshell::protocol::tree::DispatchError::Encode)?
+        ))
+    })
+}
+
+fn extract_call_inner_type(ty: &Type) -> Option<&Type> {
+    extract_outer_type_argument(ty, "Call")
+}
+
+fn extract_outer_type_argument<'a>(ty: &'a Type, expected: &str) -> Option<&'a Type> {
+    let Type::Path(TypePath { path, .. }) = ty else {
+        return None;
+    };
+    let segment = path.segments.last()?;
+    if segment.ident != expected {
+        return None;
+    }
+    let syn::PathArguments::AngleBracketed(arguments) = &segment.arguments else {
+        return None;
+    };
+    match arguments.args.first()? {
+        GenericArgument::Type(inner) => Some(inner),
+        _ => None,
+    }
+}
+
+fn extract_result_type_arguments(ty: &Type) -> Option<(&Type, &Type)> {
+    let Type::Path(TypePath { path, .. }) = ty else {
+        return None;
+    };
+    let segment = path.segments.last()?;
+    if segment.ident != "Result" {
+        return None;
+    }
+    let syn::PathArguments::AngleBracketed(arguments) = &segment.arguments else {
+        return None;
+    };
+    let mut args = arguments.args.iter();
+    let ok = match args.next()? {
+        GenericArgument::Type(value) => value,
+        _ => return None,
+    };
+    let err = match args.next()? {
+        GenericArgument::Type(value) => value,
+        _ => return None,
+    };
+    Some((ok, err))
+}
+
+fn is_unit_type(ty: &Type) -> bool {
+    matches!(ty, Type::Tuple(tuple) if tuple.elems.is_empty())
+}
+
+fn take_call_attr(attrs: &mut Vec<Attribute>) -> bool {
+    let original_len = attrs.len();
+    attrs.retain(|attr| !attr.path().is_ident("call"));
+    original_len != attrs.len()
+}
+
 #[derive(Default)]
 struct LeafAttributes {
     name: Option<LitStr>,
@@ -109,11 +393,10 @@ struct LeafAttributes {
     product: Option<LitStr>,
     version: Option<LitStr>,
     leaf_name: Option<LitStr>,
-    procedures: Option<Vec<Ident>>,
 }
 
 impl LeafAttributes {
-    fn parse_from(attrs: &[syn::Attribute]) -> Result<Self> {
+    fn parse_from(attrs: &[Attribute]) -> Result<Self> {
         let mut parsed = Self::default();
 
         for attr in attrs {
@@ -170,16 +453,6 @@ impl LeafAttributes {
                     return Ok(());
                 }
 
-                if meta.path.is_ident("procedures") {
-                    if parsed.procedures.is_some() {
-                        return Err(meta.error("duplicate leaf procedures attribute"));
-                    }
-
-                    let nested: ProcedureList = meta.input.parse()?;
-                    parsed.procedures = Some(nested.0.into_iter().collect());
-                    return Ok(());
-                }
-
                 Err(meta.error("unsupported #[leaf(...)] attribute"))
             })?;
         }
@@ -250,6 +523,56 @@ fn looks_like_canonical_leaf_name(name: &str) -> bool {
         .all(|character| character.is_ascii_digit() || character == '_')
 }
 
+#[derive(Default)]
+struct ProceduresAttributes {
+    error: Option<Type>,
+}
+
+impl Parse for ProceduresAttributes {
+    fn parse(input: syn::parse::ParseStream<'_>) -> Result<Self> {
+        if input.is_empty() {
+            return Ok(Self::default());
+        }
+
+        let mut parsed = Self::default();
+        let assignments = Punctuated::<Assignment, Token![,]>::parse_terminated(input)?;
+        for assignment in assignments {
+            if assignment.name == "error" {
+                if parsed.error.is_some() {
+                    return Err(Error::new_spanned(
+                        assignment.name,
+                        "duplicate procedures error attribute",
+                    ));
+                }
+                parsed.error = Some(assignment.value);
+                continue;
+            }
+            return Err(Error::new_spanned(
+                assignment.name,
+                "unsupported #[procedures(...)] attribute",
+            ));
+        }
+        Ok(parsed)
+    }
+}
+
+struct Assignment {
+    name: Ident,
+    value: Type,
+}
+
+impl Parse for Assignment {
+    fn parse(input: syn::parse::ParseStream<'_>) -> Result<Self> {
+        Ok(Self {
+            name: input.parse()?,
+            value: {
+                input.parse::<Token![=]>()?;
+                input.parse()?
+            },
+        })
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::looks_like_canonical_leaf_name;
@@ -267,13 +590,3 @@ mod tests {
         assert!(!looks_like_canonical_leaf_name("Org.example.v1.echo"));
     }
 }
-
-struct ProcedureList(Punctuated<Ident, Token![,]>);
-
-impl Parse for ProcedureList {
-    fn parse(input: syn::parse::ParseStream<'_>) -> Result<Self> {
-        let content;
-        syn::parenthesized!(content in input);
-        Ok(Self(Punctuated::parse_terminated(&content)?))
-    }
-}