mizan/cores/mizan-rust-macros/src/function.rs

//! `#[mizan(...)]` — on async fns. Generates:
//!   * a synthetic Input struct (`<camelName>Input`) when the fn has params
//!   * `MizanType` impl on the Input struct
//!   * canonical type entries (`<camelName>Input` / `<camelName>Output`)
//!   * Vec-element sub-type entries (so `Vec<T>` outputs surface `T` too)
//!   * `FunctionSpec` impl on a ZST `__MizanFn_<name>`
//!   * `FUNCTIONS` linkme registration of `&__MIZAN_FN_<NAME>_INSTANCE`

use heck::{ToLowerCamelCase, ToShoutySnakeCase};
use proc_macro2::TokenStream;
use quote::{format_ident, quote};
use syn::{
    parse::Parser,
    punctuated::Punctuated,
    spanned::Spanned,
    Expr, ExprPath, ExprTuple, FnArg, ItemFn, Meta, Pat, Path, ReturnType, Token, Type,
};

use crate::shape::{analyze_return, primitive_of, type_shape_expr, unwrap_option};

/// Parsed attribute args for `#[mizan(...)]`.
#[derive(Default)]
pub struct FunctionArgs {
    pub context: Option<Path>,
    pub affects: Vec<Path>,
    pub merge: Vec<Path>,
    pub websocket: bool,
    pub private: bool,
    /// `auth = "required" | "staff" | "superuser"` (or bare `auth` ⇒
    /// "required") — the `@client(auth=...)` guard. Bare-true and the string
    /// `"required"` both mean "must be authenticated".
    pub auth: Option<String>,
    /// `form_name = "..."` + `form_role = "schema"|"validate"|"submit"` — the
    /// Forms binding's per-endpoint metadata, mirroring the Django form
    /// `_meta` keys. Carried into the IR (`is-form`/`form-name`/`form-role`).
    pub form_name: Option<String>,
    pub form_role: Option<String>,
}

impl FunctionArgs {
    pub fn parse(attr_tokens: TokenStream) -> syn::Result<Self> {
        if attr_tokens.is_empty() {
            return Ok(Self::default());
        }
        let parser = Punctuated::<Meta, Token![,]>::parse_terminated;
        let metas = parser.parse2(attr_tokens)?;
        let mut out = Self::default();
        for meta in metas {
            match meta {
                Meta::NameValue(nv) => {
                    if nv.path.is_ident("context") {
                        out.context = Some(expect_path(&nv.value)?);
                    } else if nv.path.is_ident("affects") {
                        out.affects = collect_paths(&nv.value)?;
                    } else if nv.path.is_ident("merge") {
                        out.merge = collect_paths(&nv.value)?;
                    } else if nv.path.is_ident("auth") {
                        out.auth = Some(expect_str(&nv.value)?);
                    } else if nv.path.is_ident("form_name") {
                        out.form_name = Some(expect_str(&nv.value)?);
                    } else if nv.path.is_ident("form_role") {
                        out.form_role = Some(expect_str(&nv.value)?);
                    } else {
                        return Err(syn::Error::new_spanned(
                            nv.path,
                            "unknown attribute key; expected one of: context, affects, merge, auth, form_name, form_role",
                        ));
                    }
                }
                Meta::Path(p) => {
                    if p.is_ident("websocket") {
                        out.websocket = true;
                    } else if p.is_ident("private") {
                        out.private = true;
                    } else if p.is_ident("auth") {
                        out.auth = Some("required".to_string());
                    } else {
                        return Err(syn::Error::new_spanned(
                            p,
                            "unknown flag; expected `websocket`, `private`, or `auth`",
                        ));
                    }
                }
                Meta::List(l) => {
                    return Err(syn::Error::new_spanned(
                        l,
                        "list-shaped attribute args not supported here",
                    ));
                }
            }
        }
        if out.context.is_some() && !out.affects.is_empty() {
            return Err(syn::Error::new_spanned(
                out.context.as_ref().unwrap(),
                "`context` and `affects` are mutually exclusive — a function is either a context reader or a mutation.",
            ));
        }
        if out.context.is_some() && !out.merge.is_empty() {
            return Err(syn::Error::new_spanned(
                out.context.as_ref().unwrap(),
                "`context` and `merge` are mutually exclusive — a function is either a context reader or a mutation.",
            ));
        }
        Ok(out)
    }
}

fn expect_path(expr: &Expr) -> syn::Result<Path> {
    if let Expr::Path(ExprPath { path, .. }) = expr {
        Ok(path.clone())
    } else {
        Err(syn::Error::new_spanned(
            expr,
            "expected a type path (e.g. `UserCtx`)",
        ))
    }
}

fn expect_str(expr: &Expr) -> syn::Result<String> {
    if let Expr::Lit(syn::ExprLit {
        lit: syn::Lit::Str(s),
        ..
    }) = expr
    {
        Ok(s.value())
    } else {
        Err(syn::Error::new_spanned(
            expr,
            "expected a string literal (e.g. `\"staff\"`)",
        ))
    }
}

fn collect_paths(expr: &Expr) -> syn::Result<Vec<Path>> {
    match expr {
        Expr::Path(_) => Ok(vec![expect_path(expr)?]),
        Expr::Tuple(ExprTuple { elems, .. }) => elems.iter().map(expect_path).collect(),
        _ => Err(syn::Error::new_spanned(
            expr,
            "expected a context type or a tuple of context types (e.g. `UserCtx` or `(UserCtx, OrderCtx)`)",
        )),
    }
}

/// Information about one input parameter, extracted from the fn signature.
struct InputArg {
    ident: syn::Ident,
    ty: Type,
}

pub fn expand(args: FunctionArgs, item: ItemFn) -> TokenStream {
    if item.sig.asyncness.is_none() {
        return syn::Error::new_spanned(
            &item.sig.fn_token,
            "#[mizan] requires an `async fn`. Wrap synchronous handlers if needed.",
        )
        .to_compile_error();
    }

    let fn_name = item.sig.ident.to_string();
    let camel = fn_name.to_lower_camel_case();
    let input_type_name = format!("{camel}Input");
    let output_type_name = format!("{camel}Output");

    let input_args = match collect_input_args(&item) {
        Ok(v) => v,
        Err(e) => return e.to_compile_error(),
    };
    let has_input = !input_args.is_empty();
    let input_type_ident = format_ident!("{}", input_type_name);

    let return_ty = match &item.sig.output {
        ReturnType::Type(_, t) => (**t).clone(),
        ReturnType::Default => {
            return syn::Error::new_spanned(
                &item.sig,
                "#[mizan] requires an explicit return type. Add `-> T` to the signature.",
            )
            .to_compile_error();
        }
    };
    let analysis = analyze_return(&return_ty);

    // ─── Synthetic Input struct ────────────────────────────────────────────
    let input_struct = if has_input {
        let mut field_defs = Vec::new();
        let mut field_shapes = Vec::new();
        for arg in &input_args {
            let ident = &arg.ident;
            let ty = &arg.ty;
            // Strip a leading underscore from the wire-level field name —
            // Rust convention uses `_foo` to silence unused-arg warnings,
            // but the wire schema and the Python fixture name the param
            // `foo`. The struct field keeps its source ident (so the
            // dispatch wrapper's `validated.#ident` compiles), and a serde
            // `rename` bridges the wire-level JSON name.
            let name_str = ident.to_string();
            let wire_name = name_str.trim_start_matches('_').to_string();
            let serde_rename = if wire_name != name_str {
                quote! { #[serde(rename = #wire_name)] }
            } else {
                TokenStream::new()
            };
            field_defs.push(quote! { #serde_rename pub #ident: #ty, });
            let is_optional = unwrap_option(ty).is_some();
            let required = !is_optional;
            let shape = type_shape_expr(ty);
            field_shapes.push(quote! {
                ::mizan_core::StructField {
                    name: #wire_name,
                    required: #required,
                    default: ::std::option::Option::None,
                    shape: #shape,
                }
            });
        }
        quote! {
            // The synthetic Input is only ever *deserialized* (from the call's
            // JSON args by the dispatch wrapper); it is never serialized, so it
            // derives `Deserialize` only. Dropping `Serialize` lets binary
            // field types like `Upload` (deserialize-only) participate.
            #[derive(::std::fmt::Debug, ::std::clone::Clone, ::serde::Deserialize)]
            pub struct #input_type_ident {
                #(#field_defs)*
            }

            impl ::mizan_core::MizanType for #input_type_ident {
                const TYPE_NAME: &'static str = #input_type_name;
                fn shape() -> ::mizan_core::NamedType {
                    ::mizan_core::NamedType::Struct(::std::vec![
                        #(#field_shapes),*
                    ])
                }
            }
        }
    } else {
        TokenStream::new()
    };

    // ─── Type entry registrations ──────────────────────────────────────────
    // - Input: TypeEntry pointing at the synthetic input struct's shape_fn.
    // - Output: TypeEntry whose shape is a copy of the user's Output shape
    //   (for struct outputs) or an `Alias(List(Ref("T")))` (for Vec outputs).
    // - For Vec<T> outputs, ALSO register T's TypeEntry pointing at T's
    //   MizanType impl (so the Ref resolves in the IR).
    let mut type_registrations = Vec::new();
    if has_input {
        let static_ident =
            format_ident!("__MIZAN_TYPE_{}", input_type_name.to_shouty_snake_case());
        type_registrations.push(quote! {
            #[::mizan_core::__priv::linkme::distributed_slice(::mizan_core::TYPES)]
            #[linkme(crate = ::mizan_core::__priv::linkme)]
            static #static_ident: ::mizan_core::TypeEntry = ::mizan_core::TypeEntry {
                name: #input_type_name,
                shape_fn: <#input_type_ident as ::mizan_core::MizanType>::shape,
            };
        });
    }

    let output_static = format_ident!("__MIZAN_TYPE_{}", output_type_name.to_shouty_snake_case());
    if analysis.is_vec {
        let elem = analysis.vec_inner.as_ref().expect("vec_inner set");
        // userOrdersOutput → alias { list { ref "OrderOutput" } }
        // The Ref name is resolved via `<T as MizanType>::type_name()`.
        type_registrations.push(quote! {
            #[::mizan_core::__priv::linkme::distributed_slice(::mizan_core::TYPES)]
            #[linkme(crate = ::mizan_core::__priv::linkme)]
            static #output_static: ::mizan_core::TypeEntry = ::mizan_core::TypeEntry {
                name: #output_type_name,
                shape_fn: || ::mizan_core::NamedType::Alias(
                    ::mizan_core::TypeShape::List(::std::boxed::Box::new(
                        ::mizan_core::TypeShape::Ref(<#elem as ::mizan_core::MizanType>::TYPE_NAME)
                    ))
                ),
            };
        });
        // Also register the element type itself by its own name. `TYPE_NAME`
        // is an associated const, so this is usable in a static initializer.
        // The static ident scopes by the function name so two handlers
        // returning `Vec<Same>` don't collide; the IrSnapshot's BTreeMap
        // dedupes by the entry's `name` at emit time.
        let elem_static =
            element_type_static_ident_scoped(elem, &fn_name.to_shouty_snake_case());
        type_registrations.push(quote! {
            #[::mizan_core::__priv::linkme::distributed_slice(::mizan_core::TYPES)]
            #[linkme(crate = ::mizan_core::__priv::linkme)]
            static #elem_static: ::mizan_core::TypeEntry = ::mizan_core::TypeEntry {
                name: <#elem as ::mizan_core::MizanType>::TYPE_NAME,
                shape_fn: <#elem as ::mizan_core::MizanType>::shape,
            };
        });
    } else {
        // Non-Vec output: copy the inner type's shape under the canonical name.
        let inner_ty = &analysis.inner;
        type_registrations.push(quote! {
            #[::mizan_core::__priv::linkme::distributed_slice(::mizan_core::TYPES)]
            #[linkme(crate = ::mizan_core::__priv::linkme)]
            static #output_static: ::mizan_core::TypeEntry = ::mizan_core::TypeEntry {
                name: #output_type_name,
                shape_fn: <#inner_ty as ::mizan_core::MizanType>::shape,
            };
        });
    }

    // ─── InputParam slice (for context-builder shared-param elevation) ────
    let mut input_params = Vec::new();
    for arg in &input_args {
        // Wire-level name strips the underscore prefix — see input_struct
        // above for the rationale.
        let name_str = arg.ident.to_string();
        let name_str = name_str.trim_start_matches('_').to_string();
        let primitive = primitive_of(&arg.ty).unwrap_or_else(|| {
            // Non-primitive params don't surface in the context's `param`
            // block; they participate as opaque payloads. Using `String` as
            // the placeholder primitive matches Python's fallback in
            // `_annotation_to_primitive`.
            quote! { ::mizan_core::Primitive::String }
        });
        let is_optional = unwrap_option(&arg.ty).is_some();
        let required = !is_optional;
        input_params.push(quote! {
            ::mizan_core::InputParam {
                name: #name_str,
                primitive: #primitive,
                required: #required,
            }
        });
    }
    let params_static = format_ident!("__MIZAN_FN_{}_PARAMS", fn_name.to_shouty_snake_case());
    let params_const = quote! {
        const #params_static: &[::mizan_core::InputParam] = &[
            #(#input_params),*
        ];
    };

    // ─── AffectTarget / merge / context wiring ─────────────────────────────
    let affects_static = format_ident!("__MIZAN_FN_{}_AFFECTS", fn_name.to_shouty_snake_case());
    let affects_entries: Vec<_> = args
        .affects
        .iter()
        .map(|p| {
            quote! {
                ::mizan_core::AffectTarget::Context(<#p as ::mizan_core::ContextMarker>::NAME)
            }
        })
        .collect();
    let affects_const = quote! {
        const #affects_static: &[::mizan_core::AffectTarget] = &[
            #(#affects_entries),*
        ];
    };

    let merge_static = format_ident!("__MIZAN_FN_{}_MERGE", fn_name.to_shouty_snake_case());
    let merge_entries: Vec<_> = args
        .merge
        .iter()
        .map(|p| quote! { <#p as ::mizan_core::ContextMarker>::NAME })
        .collect();
    let merge_const = quote! {
        const #merge_static: &[&'static str] = &[
            #(#merge_entries),*
        ];
    };

    let context_value = match &args.context {
        Some(p) => quote! { ::std::option::Option::Some(<#p as ::mizan_core::ContextMarker>::NAME) },
        None => quote! { ::std::option::Option::None },
    };

    let transport_value = if args.websocket {
        quote! { ::mizan_core::Transport::Websocket }
    } else {
        quote! { ::mizan_core::Transport::Http }
    };

    // ─── Dispatch wrapper + FunctionSpec impl ──────────────────────────────
    let inner_fn_ident = item.sig.ident.clone();
    let spec_struct = format_ident!("__MizanFn_{}", inner_fn_ident);
    let spec_const = format_ident!("__MIZAN_FN_{}_SPEC", fn_name.to_shouty_snake_case());
    let register_static =
        format_ident!("__MIZAN_FN_{}_REGISTER", fn_name.to_shouty_snake_case());

    let input_type_opt = if has_input {
        quote! { ::std::option::Option::Some(#input_type_name) }
    } else {
        quote! { ::std::option::Option::None }
    };

    let output_nullable = analysis.nullable;
    let private = args.private;

    let auth_value = match &args.auth {
        Some(a) => quote! { ::std::option::Option::Some(#a) },
        None => quote! { ::std::option::Option::None },
    };
    let is_form = args.form_name.is_some() || args.form_role.is_some();
    let form_name_value = match &args.form_name {
        Some(n) => quote! { ::std::option::Option::Some(#n) },
        None => quote! { ::std::option::Option::None },
    };
    let form_role_value = match &args.form_role {
        Some(r) => quote! { ::std::option::Option::Some(#r) },
        None => quote! { ::std::option::Option::None },
    };

    let dispatch_body = build_dispatch(
        &item,
        &input_args,
        has_input,
        &input_type_ident,
        analysis.returns_result,
    );

    quote! {
        // Keep the user's original fn intact — the macro never rewrites the
        // body, only wraps it for dispatch.
        #item

        #input_struct

        #(#type_registrations)*

        #params_const
        #affects_const
        #merge_const

        #[allow(non_camel_case_types)]
        pub struct #spec_struct;

        impl ::mizan_core::FunctionSpec for #spec_struct {
            fn name(&self) -> &'static str { #fn_name }
            fn camel_name(&self) -> &'static str { #camel }
            fn has_input(&self) -> bool { #has_input }
            fn input_type(&self) -> ::std::option::Option<&'static str> { #input_type_opt }
            fn output_type(&self) -> &'static str { #output_type_name }
            fn output_nullable(&self) -> bool { #output_nullable }
            fn context(&self) -> ::std::option::Option<&'static str> { #context_value }
            fn affects(&self) -> &'static [::mizan_core::AffectTarget] { #affects_static }
            fn merge(&self) -> &'static [&'static str] { #merge_static }
            fn transport(&self) -> ::mizan_core::Transport { #transport_value }
            fn private(&self) -> bool { #private }
            fn auth(&self) -> ::std::option::Option<&'static str> { #auth_value }
            fn is_form(&self) -> bool { #is_form }
            fn form_name(&self) -> ::std::option::Option<&'static str> { #form_name_value }
            fn form_role(&self) -> ::std::option::Option<&'static str> { #form_role_value }
            fn input_params(&self) -> &'static [::mizan_core::InputParam] { #params_static }

            fn dispatch<'a>(
                &'a self,
                req: ::mizan_core::RequestHandle<'a>,
                args: ::mizan_core::__priv::serde_json::Value,
            ) -> ::std::pin::Pin<::std::boxed::Box<dyn ::std::future::Future<
                Output = ::std::result::Result<
                    ::mizan_core::__priv::serde_json::Value,
                    ::mizan_core::MizanError,
                >
            > + ::std::marker::Send + 'a>> {
                ::std::boxed::Box::pin(async move {
                    #dispatch_body
                })
            }
        }

        #[allow(non_upper_case_globals)]
        static #spec_const: #spec_struct = #spec_struct;

        #[::mizan_core::__priv::linkme::distributed_slice(::mizan_core::FUNCTIONS)]
        #[linkme(crate = ::mizan_core::__priv::linkme)]
        static #register_static: &dyn ::mizan_core::FunctionSpec = &#spec_const;
    }
}

fn collect_input_args(item: &ItemFn) -> syn::Result<Vec<InputArg>> {
    let mut out = Vec::new();
    let mut iter = item.sig.inputs.iter();
    // First arg is the request handle — skip without inspection. The function
    // body uses it directly; the dispatch wrapper forwards `req`.
    if iter.next().is_none() {
        return Err(syn::Error::new(
            item.sig.span(),
            "#[mizan] functions must accept at least a request handle as the first parameter (e.g. `&Request` or `RequestHandle`).",
        ));
    }
    for arg in iter {
        match arg {
            FnArg::Typed(pat) => {
                let ident = match &*pat.pat {
                    Pat::Ident(pi) => pi.ident.clone(),
                    _ => {
                        return Err(syn::Error::new_spanned(
                            &pat.pat,
                            "#[mizan] function parameters must be plain identifiers (no destructuring).",
                        ));
                    }
                };
                out.push(InputArg {
                    ident,
                    ty: (*pat.ty).clone(),
                });
            }
            FnArg::Receiver(_) => {
                return Err(syn::Error::new_spanned(
                    arg,
                    "#[mizan] functions are free functions, not methods. `self` is not allowed.",
                ));
            }
        }
    }
    Ok(out)
}

fn build_dispatch(
    item: &ItemFn,
    input_args: &[InputArg],
    has_input: bool,
    input_type_ident: &syn::Ident,
    returns_result: bool,
) -> TokenStream {
    let inner = &item.sig.ident;
    // When the user returns `Result<T, MizanError>`, lift Err out into the
    // dispatch wrapper's outer Result so the HTTP/IPC adapter can surface
    // it as the standard error envelope. When the user returns `T`,
    // serialize directly — the substrate has no error path for them.
    let unwrap_user_result = if returns_result {
        quote! { ? }
    } else {
        TokenStream::new()
    };
    if has_input {
        let arg_names: Vec<_> = input_args.iter().map(|a| &a.ident).collect();
        quote! {
            let validated: #input_type_ident = ::mizan_core::__priv::serde_json::from_value(args)
                .map_err(|e| ::mizan_core::MizanError::ValidationFailed {
                    message: format!("input validation failed: {e}"),
                    details: ::mizan_core::__priv::serde_json::Value::Null,
                })?;
            let result = #inner(
                &req,
                #( validated.#arg_names ),*
            ).await #unwrap_user_result;
            ::mizan_core::__priv::serde_json::to_value(&result)
                .map_err(|e| ::mizan_core::MizanError::InternalError(
                    format!("output serialization failed: {e}"),
                ))
        }
    } else {
        quote! {
            let _ = args;
            let result = #inner(&req).await #unwrap_user_result;
            ::mizan_core::__priv::serde_json::to_value(&result)
                .map_err(|e| ::mizan_core::MizanError::InternalError(
                    format!("output serialization failed: {e}"),
                ))
        }
    }
}

fn element_type_static_ident_scoped(ty: &Type, fn_scope: &str) -> syn::Ident {
    // Derive a unique static-name for the type's registration entry,
    // scoped by the surrounding function so siblings returning the same
    // `Vec<T>` don't collide at the static-name layer. The IR-side
    // BTreeMap dedupes by TypeEntry.name at emission time.
    let last = match ty {
        Type::Path(tp) => tp.path.segments.last().map(|s| s.ident.to_string()),
        _ => None,
    };
    let suffix = last.unwrap_or_else(|| "ANON".to_string()).to_shouty_snake_case();
    format_ident!("__MIZAN_TYPE_ELEM_{}_FOR_{}", suffix, fn_scope)
}