hetero-cartesian 0.1.0

// cartesian-macro/src/lib.rs
//
// Proc-macro implementation of `cartesian!`.
//
// Syntax:
//
//   cartesian! {
//       <OuterGeneric: Bound, ...>                       // optional outer generics
//       let pat: Type = expr;                            // optional env capture
//
//       func(args) => HandlerTrait::method<CallGenerics>(param: Type);
//       func(args) => HandlerTrait::method              (param: Type);  // no call-generics
//       func(args) => HandlerTrait                      (param: Type);  // method defaults to `call`
//       ...
//
//       { body }
//   }
//
// `_` inside func args is replaced with a reference to the generated handler struct.
// Multiple layers are composed so that each handler's `call` invocations are the
// cartesian product of all the values the functions pass to their callbacks.

use proc_macro::TokenStream;
use proc_macro2::TokenStream as TokenStream2;
use quote::{format_ident, quote};
use syn::{
    Block, Expr, GenericParam, Generics, Ident, Pat, Path, Result, Token, Type, parenthesized,
    parse::{Parse, ParseStream},
    parse_macro_input, token,
};

// ─── AST types ───────────────────────────────────────────────────────────────

/// One argument passed to the outer function.  `_` means "put the handler here".
enum FuncArg {
    Hole,
    Expr(Expr),
}

/// One layer: `func(args) => HandlerTrait::method<call_generics>(param: Type)`
struct Layer {
    func: Ident,
    func_args: Vec<FuncArg>,
    handler: Path,
    method: Ident, // trait method name, defaults to `call`
    call_generics: Vec<GenericParam>,
    param: Ident,
    param_ty: Type,
}

/// Optional `let pat: Type = expr;` environment capture.
struct EnvCapture {
    pat: Pat,
    ty: Type,
    expr: Expr,
}

/// The full macro input.
struct CartesianInput {
    outer_generics: Vec<GenericParam>,
    env: Option<EnvCapture>,
    layers: Vec<Layer>,
    body: Block,
}

// ─── Parsing ─────────────────────────────────────────────────────────────────

fn parse_func_arg(input: ParseStream) -> Result<FuncArg> {
    if input.peek(Token![_]) {
        input.parse::<Token![_]>()?;
        Ok(FuncArg::Hole)
    } else {
        Ok(FuncArg::Expr(input.parse()?))
    }
}

impl Parse for CartesianInput {
    fn parse(input: ParseStream) -> Result<Self> {
        // <OuterGeneric: Bound, ...>
        let outer_generics = if input.peek(Token![<]) {
            let generics: Generics = input.parse()?;
            generics.params.into_iter().collect()
        } else {
            vec![]
        };

        // let pat: Type = expr;
        let env = if input.peek(Token![let]) {
            input.parse::<Token![let]>()?;
            let pat = Pat::parse_single(input)?;
            input.parse::<Token![:]>()?;
            let ty: Type = input.parse()?;
            input.parse::<Token![=]>()?;
            let expr: Expr = input.parse()?;
            input.parse::<Token![;]>()?;
            Some(EnvCapture { pat, ty, expr })
        } else {
            None
        };

        // layers: func(args) => HandlerTrait::method<generics>(param: Type);
        let mut layers = vec![];
        while !input.peek(token::Brace) && !input.is_empty() {
            let func: Ident = input.parse()?;

            let args_buf;
            parenthesized!(args_buf in input);
            let mut func_args = vec![];
            loop {
                if args_buf.is_empty() {
                    break;
                }
                func_args.push(parse_func_arg(&args_buf)?);
                if args_buf.peek(Token![,]) {
                    args_buf.parse::<Token![,]>()?;
                } else {
                    break;
                }
            }

            input.parse::<Token![=>]>()?;

            // parse_mod_style stops before `<`, so it won't greedily consume
            // `<const N: usize>` as generic arguments of the path segment.
            let mut handler: Path = input.call(Path::parse_mod_style)?;
            let method: Ident = if handler.segments.len() > 1 {
                let seg = handler.segments.pop().unwrap().into_value();
                // pop() removes the last value but leaves the trailing `::` punct;
                // pop_punct() removes that dangling separator.
                handler.segments.pop_punct();
                seg.ident
            } else {
                format_ident!("call")
            };

            // Optional <CallGenerics> — parsed as syn::Generics so nested
            // angle brackets inside bounds (e.g. `T: Trait<U>`) are handled.
            let call_generics = if input.peek(Token![<]) {
                let generics: Generics = input.parse()?;
                generics.params.into_iter().collect()
            } else {
                vec![]
            };

            let param_buf;
            parenthesized!(param_buf in input);
            let param: Ident = param_buf.parse()?;
            param_buf.parse::<Token![:]>()?;
            let param_ty: Type = param_buf.parse()?;

            input.parse::<Token![;]>()?;

            layers.push(Layer {
                func,
                func_args,
                handler,
                method,
                call_generics,
                param,
                param_ty,
            });
        }

        let body: Block = input.parse()?;
        Ok(CartesianInput {
            outer_generics,
            env,
            layers,
            body,
        })
    }
}

// ─── Code generation helpers ─────────────────────────────────────────────────

/// Turn generic *declarations* into the corresponding *arguments* (just the idents).
fn params_to_args(params: &[&GenericParam]) -> Vec<TokenStream2> {
    params
        .iter()
        .map(|p| match p {
            GenericParam::Type(t) => {
                let id = &t.ident;
                quote! { #id }
            }
            GenericParam::Const(c) => {
                let id = &c.ident;
                quote! { #id }
            }
            GenericParam::Lifetime(l) => {
                let lt = &l.lifetime;
                quote! { #lt }
            }
        })
        .collect()
}

/// Build the PhantomData type that references the *outer* generics only.
/// Const generics are handled implicitly through field types.
fn phantom_type(outer_generics: &[GenericParam]) -> TokenStream2 {
    let tys: Vec<TokenStream2> = outer_generics
        .iter()
        .filter_map(|p| match p {
            GenericParam::Type(t) => {
                let id = &t.ident;
                Some(quote! { #id })
            }
            GenericParam::Lifetime(l) => {
                let lt = &l.lifetime;
                Some(quote! { &#lt () })
            }
            GenericParam::Const(_) => None,
        })
        .collect();
    quote! { (#(#tys,)*) }
}

/// Walk a pattern and collect every bound identifier (skips wildcards).
fn pat_idents(pat: &Pat) -> Vec<Ident> {
    match pat {
        Pat::Ident(p) if p.ident != "_" => vec![p.ident.clone()],
        Pat::Tuple(p) => p.elems.iter().flat_map(pat_idents).collect(),
        Pat::Wild(_) => vec![],
        Pat::Reference(r) => pat_idents(&r.pat),
        _ => vec![],
    }
}

/// The three shadow_env traits used to coerce `&mut FieldType` → the right reference kind.
///
/// Priority (via auto-deref in method resolution):
///   &mut &mut T  → UnwrapMutMut → &mut T     (field was &mut T)
///   &&mut &T     → UnwrapMutRef → &T          (field was &T)
///   &&&mut T     → UnwrapVal   → &mut T       (field was plain T, returned as &mut)
fn shadow_env_traits() -> TokenStream2 {
    quote! {
        #[allow(dead_code)]
        struct __CartesianWrap<T>(T);

        #[allow(dead_code)]
        trait __ShadowMutMut { type Out; fn shadow_env(self) -> Self::Out; }
        impl<'__a, '__b, T: ?Sized> __ShadowMutMut for __CartesianWrap<&'__a mut &'__b mut T> {
            type Out = &'__a mut T;
            #[inline(always)] fn shadow_env(self) -> Self::Out {
                self.0
            }
        }

        #[allow(dead_code)]
        trait __ShadowMutRef { type Out; fn shadow_env(self) -> Self::Out; }
        impl<'__a, '__b, T: ?Sized> __ShadowMutRef for __CartesianWrap<&'__a mut &'__b T> {
            type Out = &'__b T;
            #[inline(always)] fn shadow_env(self) -> Self::Out {
                *self.0
            }
        }

        #[allow(dead_code)]
        trait __ShadowVal { type Out; fn shadow_env(self) -> Self::Out; }
        impl<'__a, T: ::core::clone::Clone> __ShadowVal for &__CartesianWrap<&'__a mut T> {
            type Out = T;
            #[inline(always)] fn shadow_env(self) -> Self::Out {
                self.0.clone()
            }
        }
    }
}

// ─── Core recursive generator ─────────────────────────────────────────────────

/// State threaded through the recursive generation.
struct Ctx<'a> {
    input: &'a CartesianInput,
    depth: usize,
    /// Call-method generics accumulated from all outer layers.
    acc_call_generics: Vec<GenericParam>,
    /// For each outer layer: (struct_field_ident, user_param_ident, param_type).
    captured: Vec<(Ident, Ident, Type)>,
    /// Expression that yields the `*mut ()` env pointer for the current struct init.
    env_ptr: TokenStream2,
}

/// Generate all code for layer `ctx.depth` and (recursively) every inner layer.
///
/// The returned tokens form a sequence of statements:
///   1. `struct __CartesianL{depth} { ... }`
///   2. `impl HandlerTrait for __CartesianL{depth} { fn call(...) { <inner> } }`
///   3. Local binding + function call: `let mut __h = ...; func(..., &mut __h, ...)`
fn gen_layer(ctx: &Ctx) -> TokenStream2 {
    let depth = ctx.depth;
    let layer = &ctx.input.layers[depth];
    let struct_name = format_ident!("__CartesianL{}", depth);

    // ── Generics for the struct/impl ────────────────────────────────────────
    let outer_g = &ctx.input.outer_generics;
    let all_g: Vec<&GenericParam> = outer_g.iter().chain(ctx.acc_call_generics.iter()).collect();
    let all_g_args = params_to_args(&all_g);
    let phantom = phantom_type(outer_g);

    // ── Struct definition ────────────────────────────────────────────────────
    let field_defs: Vec<_> = ctx
        .captured
        .iter()
        .map(|(f, _, ty)| quote! { #f: #ty })
        .collect();

    let struct_def = if all_g.is_empty() {
        quote! {
            #[allow(non_local_definitions)]
            struct #struct_name {
                __env:    *mut (),
                __marker: ::core::marker::PhantomData<#phantom>,
                #(#field_defs,)*
            }
        }
    } else {
        quote! {
            #[allow(non_local_definitions)]
            struct #struct_name<#(#all_g),*> {
                __env:    *mut (),
                __marker: ::core::marker::PhantomData<#phantom>,
                #(#field_defs,)*
            }
        }
    };

    // ── Prepare next-layer state ─────────────────────────────────────────────
    let handler = &layer.handler;
    let method = &layer.method;
    let call_generics = &layer.call_generics;
    let param = &layer.param;
    let param_ty = &layer.param_ty;

    let field_name = format_ident!("__cartesian_p{}", depth);

    let mut new_captured = ctx.captured.clone();
    new_captured.push((field_name.clone(), param.clone(), param_ty.clone()));

    let mut new_acc_generics = ctx.acc_call_generics.clone();
    new_acc_generics.extend(call_generics.iter().cloned());

    // ── call() body ──────────────────────────────────────────────────────────
    // Restore outer params by cloning struct fields into local bindings.
    let clone_stmts: Vec<_> = ctx
        .captured
        .iter()
        .map(|(f, name, _)| quote! { let #name = self.#f.clone(); })
        .collect();

    let call_body = if depth + 1 == ctx.input.layers.len() {
        // Innermost layer — unpack env and run the user's body.
        let body_code = gen_body(ctx.input);
        quote! { #(#clone_stmts)* #body_code }
    } else {
        // More layers remain — recurse.
        let next = gen_layer(&Ctx {
            input: ctx.input,
            depth: depth + 1,
            acc_call_generics: new_acc_generics,
            captured: new_captured.clone(),
            env_ptr: quote! { self.__env },
        });
        quote! { #(#clone_stmts)* #next }
    };

    // ── impl block ───────────────────────────────────────────────────────────
    let call_generic_decl = if call_generics.is_empty() {
        quote! {}
    } else {
        quote! { <#(#call_generics),*> }
    };

    let impl_block = if all_g.is_empty() {
        quote! {
            #[allow(non_local_definitions)]
            impl #handler for #struct_name {
                fn #method #call_generic_decl (&mut self, #param: #param_ty) {
                    #call_body
                }
            }
        }
    } else {
        quote! {
            #[allow(non_local_definitions)]
            impl<#(#all_g),*> #handler for #struct_name<#(#all_g_args),*> {
                fn #method #call_generic_decl (&mut self, #param: #param_ty) {
                    #call_body
                }
            }
        }
    };

    // ── Struct initialiser + function call ───────────────────────────────────
    // Fields: previously captured params (now available as locals after clone_stmts
    // in the *enclosing* call body, or directly as fn args at depth 0).
    let env_ptr = &ctx.env_ptr;
    let captured_init: Vec<_> = ctx
        .captured
        .iter()
        .map(|(f, name, _)| quote! { #f: #name })
        .collect();

    // Bind to a local so the borrow is stable when passed to the function.
    let handler_binding = format_ident!("__cartesian_handler_{}", depth);

    let handler_init = if all_g.is_empty() {
        quote! {
            let mut #handler_binding = #struct_name {
                __env:    #env_ptr,
                __marker: ::core::marker::PhantomData,
                #(#captured_init,)*
            };
        }
    } else {
        quote! {
            let mut #handler_binding: #struct_name<#(#all_g_args),*> = #struct_name {
                __env:    #env_ptr,
                __marker: ::core::marker::PhantomData,
                #(#captured_init,)*
            };
        }
    };

    let func = &layer.func;
    let func_args: Vec<_> = layer
        .func_args
        .iter()
        .map(|arg| match arg {
            FuncArg::Hole => quote! { &mut #handler_binding },
            FuncArg::Expr(e) => quote! { #e },
        })
        .collect();

    quote! {
        #struct_def
        #impl_block
        #handler_init
        #func(#(#func_args),*)
    }
}

/// Generate the body of the innermost `call` method:
/// unpack the env pointer back into typed references, then run the user's block.
fn gen_body(input: &CartesianInput) -> TokenStream2 {
    let body = &input.body;

    let Some(env) = &input.env else {
        return quote! { #body };
    };

    let env_ty = &env.ty;
    let env_pat = &env.pat;
    let vars = pat_idents(env_pat);

    let traits = shadow_env_traits();

    // Cast the raw pointer back and destructure.
    // Then apply shadow_env so each variable has its original type (not &mut OrigType).
    let unpack = if vars.is_empty() {
        // Wildcard pattern — env captured but not used.
        quote! {}
    } else if vars.len() == 1 {
        quote! {
            let __cartesian_env_ref = self.__env as *mut #env_ty;
            #[allow(unused_variables)]
            let #env_pat = unsafe { &mut *__cartesian_env_ref };
            #[allow(unused_variables)]
            let #env_pat = __CartesianWrap(#env_pat).shadow_env();
        }
    } else {
        let shadow_calls: Vec<_> = vars
            .iter()
            .map(|v| quote! { __CartesianWrap(#v).shadow_env() })
            .collect();
        quote! {
            let __cartesian_env_ref = self.__env as *mut #env_ty;
            #[allow(unused_variables)]
            let #env_pat = unsafe { &mut *__cartesian_env_ref };
            #[allow(unused_variables)]
            let (#(#vars,)*) = (#(#shadow_calls,)*);
        }
    };

    quote! {
        #traits
        #unpack
        #body
    }
}

// ─── Entry point ─────────────────────────────────────────────────────────────

#[proc_macro]
pub fn cartesian(input: TokenStream) -> TokenStream {
    let parsed = parse_macro_input!(input as CartesianInput);

    if parsed.layers.is_empty() {
        return quote! { compile_error!("cartesian! requires at least one layer") }.into();
    }

    // Set up the environment storage (stack-allocated, behind a type-erased pointer).
    let env_setup = match &parsed.env {
        Some(env) => {
            let ty = &env.ty;
            let expr = &env.expr;
            quote! {
                let mut __cartesian_env_val: #ty = #expr;
                let __cartesian_env_ptr: *mut () =
                    &mut __cartesian_env_val as *mut _ as *mut ();
            }
        }
        None => quote! {
            let __cartesian_env_ptr: *mut () = ::core::ptr::null_mut();
        },
    };

    let code = gen_layer(&Ctx {
        input: &parsed,
        depth: 0,
        acc_call_generics: vec![],
        captured: vec![],
        env_ptr: quote! { __cartesian_env_ptr },
    });

    quote! {{
        #env_setup
        #code
    }}
    .into()
}