aizu_macros/

lib.rs

//! Proc macros for Aizu SDK.
//!
//! This crate provides the `#[query]`, `#[mutation]`, `#[action]`, `#[http]`, and `#[ws]` macros.

use proc_macro::TokenStream;
use proc_macro2::TokenStream as TokenStream2;
use quote::{format_ident, quote};
use syn::{parse_macro_input, FnArg, ItemFn, Pat, PatType, ReturnType, LitStr, Token, Ident};
use syn::parse::{Parse, ParseStream};

/// **Query** - Read-only database access.
///
/// Queries are deterministic and cacheable. Use for fetching data without side effects.
///
/// | Capability | Allowed |
/// |------------|---------|
/// | Read DB    | Yes     |
/// | Write DB   | No      |
/// | HTTP       | No      |
/// | Auth       | Yes     |
///
/// # Example
///
/// ```ignore
/// #[query]
/// pub fn get_user(ctx: &Ctx, id: Id<User>) -> Option<User> {
///     ctx.db.get(id)
/// }
/// ```
#[proc_macro_attribute]
pub fn query(_attr: TokenStream, item: TokenStream) -> TokenStream {
    let input = parse_macro_input!(item as ItemFn);
    generate_export(input, "query").into()
}

/// **Mutation** - Read and write database access.
///
/// Mutations are transactional and automatically retried on conflict.
/// Use for creating, updating, or deleting data.
///
/// | Capability | Allowed |
/// |------------|---------|
/// | Read DB    | Yes     |
/// | Write DB   | Yes     |
/// | HTTP       | No      |
/// | Auth       | Yes     |
///
/// # Example
///
/// ```ignore
/// #[mutation]
/// pub fn create_todo(ctx: &Ctx, title: String) -> Id<Todo> {
///     ctx.db.insert(&Todo { title, completed: false }).unwrap()
/// }
/// ```
#[proc_macro_attribute]
pub fn mutation(_attr: TokenStream, item: TokenStream) -> TokenStream {
    let input = parse_macro_input!(item as ItemFn);
    generate_export(input, "mutation").into()
}

/// **Action** - Full access: database + HTTP.
///
/// Actions can perform external I/O like HTTP requests. Not automatically retried.
/// Use for workflows that need to call external APIs.
///
/// | Capability | Allowed |
/// |------------|---------|
/// | Read DB    | Yes     |
/// | Write DB   | Yes     |
/// | HTTP       | Yes     |
/// | Auth       | Yes     |
///
/// # Example
///
/// ```ignore
/// #[action]
/// pub fn process_payment(ctx: &Ctx, order_id: Id<Order>) -> Result<bool, String> {
///     let order = ctx.db.get(order_id).ok_or("Order not found")?;
///     ctx.http.post("https://api.stripe.com/charge", order.amount)?;
///     ctx.db.insert(&PaymentRecord { order_id, status: "paid" })?;
///     Ok(true)
/// }
/// ```
#[proc_macro_attribute]
pub fn action(_attr: TokenStream, item: TokenStream) -> TokenStream {
    let input = parse_macro_input!(item as ItemFn);
    generate_export(input, "action").into()
}

/// **HTTP** - HTTP route handler.
///
/// Registers a function as an HTTP endpoint. The function receives an `HttpRequest`
/// and should return a `RouteResponse`.
///
/// | Capability | Allowed |
/// |------------|---------|
/// | Read DB    | Yes     |
/// | Write DB   | Yes     |
/// | HTTP       | Yes     |
/// | Auth       | Yes     |
///
/// # Example
///
/// ```ignore
/// #[http(GET, "/api/todos")]
/// pub fn list_todos(ctx: &Ctx, req: HttpRequest) -> RouteResponse {
///     let todos = ctx.db.query::<Todo>().collect::<Vec<_>>();
///     RouteResponse::json(&todos)
/// }
///
/// #[http(GET, "/api/todos/:id")]
/// pub fn get_todo(ctx: &Ctx, req: HttpRequest) -> RouteResponse {
///     let id = req.param("id").unwrap();
///     match ctx.db.get::<Todo>(Id::from_str(id)) {
///         Some(todo) => RouteResponse::json(&todo),
///         None => RouteResponse::not_found("Todo not found"),
///     }
/// }
///
/// #[http(POST, "/api/todos")]
/// pub fn create_todo(ctx: &Ctx, req: HttpRequest) -> RouteResponse {
///     let body: CreateTodo = req.json().unwrap();
///     let id = ctx.db.insert(&Todo { title: body.title, completed: false }).unwrap();
///     RouteResponse::json(&id).with_status(201)
/// }
/// ```
#[proc_macro_attribute]
pub fn http(attr: TokenStream, item: TokenStream) -> TokenStream {
    let args = parse_macro_input!(attr as HttpArgs);
    let input = parse_macro_input!(item as ItemFn);
    generate_http_export(input, args).into()
}

/// Arguments for the `#[http]` macro: `#[http(METHOD, "/path")]`
struct HttpArgs {
    method: Ident,
    path: LitStr,
}

impl Parse for HttpArgs {
    fn parse(input: ParseStream) -> syn::Result<Self> {
        let method: Ident = input.parse()?;
        input.parse::<Token![,]>()?;
        let path: LitStr = input.parse()?;
        Ok(HttpArgs { method, path })
    }
}

/// Normalize a URL path to a valid identifier suffix.
///
/// - `/api/todos/:id` → `api_todos__id`
/// - `/` → `root`
fn normalize_path(path: &str) -> String {
    if path == "/" {
        return "root".to_string();
    }

    path.trim_start_matches('/')
        .replace('/', "_")
        .replace(':', "_")
        .replace('-', "_")
}

/// Generate the HTTP handler export.
fn generate_http_export(func: ItemFn, args: HttpArgs) -> TokenStream2 {
    let fn_name = &func.sig.ident;
    let fn_vis = &func.vis;
    let fn_block = &func.block;
    let fn_output = &func.sig.output;
    let fn_attrs = &func.attrs;

    let method = args.method.to_string().to_uppercase();
    let path = args.path.value();
    let normalized_path = normalize_path(&path);

    // Export name: __aizu_http_{METHOD}_{path_normalized}
    let export_name = format_ident!("__aizu_http_{}_{}", method.to_lowercase(), normalized_path);

    // The function signature should be: fn(ctx: &Ctx, req: HttpRequest) -> RouteResponse
    // We'll pass the serialized HttpRequest as args

    // Determine if return type is Result-wrapped
    let is_result = match fn_output {
        ReturnType::Default => false,
        ReturnType::Type(_, ty) => {
            let ty_str = quote!(#ty).to_string();
            ty_str.starts_with("Result") || ty_str.contains(":: Result")
        }
    };

    // Generate result handling that serializes RouteResponse as JSON
    let result_handling = if is_result {
        quote! {
            match result {
                Ok(response) => {
                    match aizu::serde_json::to_vec(&response) {
                        Ok(data) => {
                            aizu::set_result(&data);
                            0
                        }
                        Err(_) => -3,
                    }
                }
                Err(e) => {
                    let err_response = aizu::RouteResponse::error(500, &aizu::format!("{:?}", e));
                    match aizu::serde_json::to_vec(&err_response) {
                        Ok(data) => {
                            aizu::set_result(&data);
                            -100
                        }
                        Err(_) => -3,
                    }
                }
            }
        }
    } else {
        quote! {
            match aizu::serde_json::to_vec(&result) {
                Ok(data) => {
                    aizu::set_result(&data);
                    0
                }
                Err(_) => -3,
            }
        }
    };

    let expanded = quote! {
        // Original function (kept for direct calls in tests)
        #(#fn_attrs)*
        #fn_vis fn #fn_name(ctx: &aizu::Ctx, req: aizu::HttpRequest) #fn_output
        #fn_block

        // WASM export for HTTP handler
        #[unsafe(no_mangle)]
        pub extern "C" fn #export_name(
            args_ptr: i32,
            args_len: i32,
        ) -> i32 {
            aizu::with_ctx(|ctx| {
                // Deserialize HttpRequest from JSON
                let args_slice = if args_ptr == 0 || args_len <= 0 {
                    &[]
                } else {
                    unsafe {
                        core::slice::from_raw_parts(args_ptr as *const u8, args_len as usize)
                    }
                };

                let req: aizu::HttpRequest = match aizu::serde_json::from_slice(args_slice) {
                    Ok(r) => r,
                    Err(_) => return -1,
                };

                // Call the user function
                let result = #fn_name(ctx, req);

                // Serialize RouteResponse and return
                #result_handling
            })
        }
    };

    expanded
}

/// Generate the WASM export wrapper for a function.
///
/// Generates exports compatible with Tatara's calling convention:
/// - Function: `__aizu_{kind}_{name}(args_ptr, args_len) -> i32`
/// - Returns 0 on success, negative error code on failure
/// - Results written to static buffer, read via get_result_ptr/get_result_len
fn generate_export(func: ItemFn, kind: &str) -> TokenStream2 {
    let fn_name = &func.sig.ident;
    let fn_vis = &func.vis;
    let fn_block = &func.block;
    let fn_output = &func.sig.output;
    let fn_attrs = &func.attrs;

    // Generate export name: __aizu_{kind}_{name}
    let export_name = format_ident!("__aizu_{}_{}", kind, fn_name);

    // Extract function arguments (skip first `ctx: &Ctx`)
    let args: Vec<_> = func.sig.inputs.iter().skip(1).collect();

    // Extract argument names and types
    let mut arg_names = Vec::new();
    let mut arg_types = Vec::new();

    for arg in &args {
        if let FnArg::Typed(PatType { pat, ty, .. }) = arg
            && let Pat::Ident(pat_ident) = &**pat
        {
            arg_names.push(&pat_ident.ident);
            arg_types.push(&**ty);
        }
    }

    // Generate args struct name
    let args_struct_name = format_ident!("__{}_Args", fn_name);

    // Determine return type for serialization
    let (_return_type, is_result) = match fn_output {
        ReturnType::Default => (quote! { () }, false),
        ReturnType::Type(_, ty) => {
            // Check if it's a Result type
            let ty_str = quote!(#ty).to_string();
            if ty_str.starts_with("Result") || ty_str.contains(":: Result") {
                (quote! { #ty }, true)
            } else {
                (quote! { #ty }, false)
            }
        }
    };

    // Generate the result handling code
    // Use to_vec_named to serialize structs as maps with named keys (not arrays)
    let result_handling = if is_result {
        quote! {
            match result {
                Ok(value) => {
                    match aizu::rmp_serde::to_vec_named(&value) {
                        Ok(data) => {
                            aizu::set_result(&data);
                            0
                        }
                        Err(_) => -3, // Serialization error
                    }
                }
                Err(e) => {
                    // Serialize error as string (use Display, not Debug, to avoid quotes)
                    let err_msg = aizu::format!("{}", e);
                    aizu::set_result(err_msg.as_bytes());
                    -100 // Error indicator
                }
            }
        }
    } else {
        quote! {
            match aizu::rmp_serde::to_vec_named(&result) {
                Ok(data) => {
                    aizu::set_result(&data);
                    0
                }
                Err(_) => -3, // Serialization error
            }
        }
    };

    // Handle the case of zero arguments (no args struct needed)
    let args_handling = if arg_names.is_empty() {
        quote! {
            // No arguments to deserialize
            let _ = args_slice;
        }
    } else {
        quote! {
            let args: #args_struct_name = match aizu::rmp_serde::from_slice(args_slice) {
                Ok(a) => a,
                Err(_) => return -1, // Deserialization error
            };
        }
    };

    // Generate the function call with or without args
    let fn_call = if arg_names.is_empty() {
        quote! { #fn_name(ctx) }
    } else {
        quote! { #fn_name(ctx, #(args.#arg_names),*) }
    };

    // Generate the args struct only if there are arguments
    let args_struct = if arg_names.is_empty() {
        quote! {}
    } else {
        quote! {
            #[derive(aizu::serde::Deserialize)]
            #[allow(non_camel_case_types)]
            struct #args_struct_name {
                #(#arg_names: #arg_types),*
            }
        }
    };

    // Generate the wrapper
    let expanded = quote! {
        // Original function (kept for direct calls in tests)
        #(#fn_attrs)*
        #fn_vis fn #fn_name(ctx: &aizu::Ctx, #(#arg_names: #arg_types),*) #fn_output
        #fn_block

        // Args struct for deserialization
        #args_struct

        // WASM export (Tatara calling convention: args_ptr, args_len -> status)
        #[unsafe(no_mangle)]
        pub extern "C" fn #export_name(
            args_ptr: i32,
            args_len: i32,
        ) -> i32 {
            aizu::with_ctx(|ctx| {
                // Get args slice
                let args_slice = if args_ptr == 0 || args_len <= 0 {
                    &[]
                } else {
                    unsafe {
                        core::slice::from_raw_parts(args_ptr as *const u8, args_len as usize)
                    }
                };

                #args_handling

                // Call the user function
                let result = #fn_call;

                // Serialize and return
                #result_handling
            })
        }
    };

    expanded
}

/// **WebSocket** - WebSocket route handler.
///
/// Registers a function as a WebSocket endpoint. The function receives a `WsMessage`
/// and should return a `WsResponse`.
///
/// | Capability | Allowed |
/// |------------|---------|
/// | Read DB    | Yes     |
/// | Write DB   | Yes     |
/// | HTTP       | Yes     |
/// | Auth       | Yes     |
///
/// # Example
///
/// ```ignore
/// #[ws("/ws/chat")]
/// pub fn chat_handler(ctx: &Ctx, msg: WsMessage) -> WsResponse {
///     match msg {
///         WsMessage::Connect { params, .. } => {
///             // Client connected
///             WsResponse::text("Welcome!")
///         }
///         WsMessage::Text(text) => {
///             // Echo the message back
///             WsResponse::text(format!("You said: {}", text))
///         }
///         WsMessage::Close { .. } => {
///             // Client disconnecting
///             WsResponse::ok()
///         }
///         _ => WsResponse::ok(),
///     }
/// }
/// ```
#[proc_macro_attribute]
pub fn ws(attr: TokenStream, item: TokenStream) -> TokenStream {
    let args = parse_macro_input!(attr as WsArgs);
    let input = parse_macro_input!(item as ItemFn);
    generate_ws_export(input, args).into()
}

/// Arguments for the `#[ws]` macro: `#[ws("/path")]`
struct WsArgs {
    path: LitStr,
}

impl Parse for WsArgs {
    fn parse(input: ParseStream) -> syn::Result<Self> {
        let path: LitStr = input.parse()?;
        Ok(WsArgs { path })
    }
}

/// Generate the WebSocket handler export.
fn generate_ws_export(func: ItemFn, args: WsArgs) -> TokenStream2 {
    let fn_name = &func.sig.ident;
    let fn_vis = &func.vis;
    let fn_block = &func.block;
    let fn_output = &func.sig.output;
    let fn_attrs = &func.attrs;

    let path = args.path.value();
    let normalized_path = normalize_path(&path);

    // Export name: __aizu_ws_{path_normalized}
    let export_name = format_ident!("__aizu_ws_{}", normalized_path);

    // The function signature should be: fn(ctx: &Ctx, msg: WsMessage) -> WsResponse

    // Determine if return type is Result-wrapped
    let is_result = match fn_output {
        ReturnType::Default => false,
        ReturnType::Type(_, ty) => {
            let ty_str = quote!(#ty).to_string();
            ty_str.starts_with("Result") || ty_str.contains(":: Result")
        }
    };

    // Generate result handling that serializes WsResponse as JSON
    let result_handling = if is_result {
        quote! {
            match result {
                Ok(response) => {
                    match aizu::serde_json::to_vec(&response) {
                        Ok(data) => {
                            aizu::set_result(&data);
                            0
                        }
                        Err(_) => -3,
                    }
                }
                Err(e) => {
                    let err_response = aizu::WsResponse::close_with_error(&aizu::format!("{:?}", e));
                    match aizu::serde_json::to_vec(&err_response) {
                        Ok(data) => {
                            aizu::set_result(&data);
                            -100
                        }
                        Err(_) => -3,
                    }
                }
            }
        }
    } else {
        quote! {
            match aizu::serde_json::to_vec(&result) {
                Ok(data) => {
                    aizu::set_result(&data);
                    0
                }
                Err(_) => -3,
            }
        }
    };

    let expanded = quote! {
        // Original function (kept for direct calls in tests)
        #(#fn_attrs)*
        #fn_vis fn #fn_name(ctx: &aizu::Ctx, msg: aizu::WsMessage) #fn_output
        #fn_block

        // WASM export for WebSocket handler
        #[unsafe(no_mangle)]
        pub extern "C" fn #export_name(
            args_ptr: i32,
            args_len: i32,
        ) -> i32 {
            aizu::with_ctx(|ctx| {
                // Deserialize WsMessage from JSON
                let args_slice = if args_ptr == 0 || args_len <= 0 {
                    &[]
                } else {
                    unsafe {
                        core::slice::from_raw_parts(args_ptr as *const u8, args_len as usize)
                    }
                };

                let msg: aizu::WsMessage = match aizu::serde_json::from_slice(args_slice) {
                    Ok(m) => m,
                    Err(_) => return -1,
                };

                // Call the user function
                let result = #fn_name(ctx, msg);

                // Serialize WsResponse and return
                #result_handling
            })
        }
    };

    expanded
}