batch-impl 0.6.6

use quote::ToTokens;

use crate::apply::{check_expand_limit, err_ty, err_ty_at};
use crate::ast::*;
use crate::parse::parse_primitive;
use proc_macro2::Span;

/// `N..M` / `N..=M`: calls f for every length n in the range, packing results into a list.
/// Empty range (`start >= end`) or over-limit (len > [`MAX_EXPAND`]): a typo diagnostic
pub(crate) fn map_range(
    start: usize, end: usize, inclusive: bool, span: Span, f: impl Fn(usize) -> Ty,
) -> Ty {
    let end_mark = if inclusive { "=" } else { "" };
    let ns: Vec<_> =
        if inclusive { (start..=end).collect() } else { (start..end).collect() };
    if ns.is_empty() {
        return err_ty_at(
            &format!(
                "batch-impl: range `{}..{}{}` is empty (start not below end); no impls will be generated",
                start, end, end_mark
            ),
            span,
        );
    }
    if let Some(e) = check_expand_limit(
        &format!("range `{}..{}{}`", start, end, end_mark),
        ns.len(),
    ) {
        return e;
    }
    TyArray(ns.into_iter().map(f).collect()).into()
}

/// `(...,)^N`: expands the tuple to length N (empty / single / multi-element handled separately)
/// `N` above [`MAX_EXPAND`] is a typo diagnostic (covers `()^N` / `(T,)^N`).
fn tuple_pow(mut elems: Vec<Ty>, n: usize) -> Ty {
    if let Some(e) = check_expand_limit(&format!("tuple `^{}`", n), n) {
        return e;
    }
    match elems.len() {
        0 => pow_empty(n),
        // len == 1 is guaranteed by the match; the remove(0) out-of-bounds branch is unreachable
        1 => pow_single(elems.remove(0), n),
        _ => pow_cartesian(elems, n),
    }
}

/// `()^N` => `<A,B,...,N>(A,B,...,N)` — generate N fresh generic params and wrap
fn pow_empty(n: usize) -> Ty {
    if n == 0 {
        return TyTuple(vec![]).into();
    }
    let params = fresh_params(n);
    let param_names = params.iter().map(|p| p.to_token_stream()).collect::<Vec<_>>();
    let tp: Ty = TyTypeParam {
        params: param_names.into_iter().map(|n| (n, None)).collect(),
        bindings: vec![],
    }
    .into();
    tp.apply(TyTuple(params).into())
}

/// `(T,)^N` => `(T,T,...,T)`; `(<Bound>)^N` => `(A:Bound, B:Bound, ...)`
fn pow_single(template: Ty, n: usize) -> Ty {
    let template_span = template.span;
    if let TyKind::TypeParam(tp) = template.kind.clone() {
        // From `(<Bound>)^N`: exactly one unbound param (guaranteed by parse_angle_bracket_contents)
        if tp.params.len() != 1 || tp.params[0].1.is_some() {
            return err_ty(
                "batch-impl: unexpected bound parameter in (<Trait>)⁁; this is an internal error",
            );
        }
        let params = fresh_params(n);
        let param_names =
            params.iter().map(|p| p.to_token_stream()).collect::<Vec<_>>();
        let bound_tokens = tp.params[0].0.clone().into_iter().collect::<Vec<_>>();
        return Ty::new(
            template_span,
            TyKind::TypeParam(TyTypeParam {
                params: param_names
                    .into_iter()
                    .map(|n| (n, parse_primitive(&bound_tokens, None).into()))
                    .collect(),
                bindings: vec![],
            }),
        )
        .apply(TyTuple(params).into());
    }
    TyTuple((0..n).map(|_| Ty::new(template_span, template.kind.clone())).collect())
        .into()
}

/// `(A,B,..)^N`: N-way Cartesian product, choosing one of all elements per position.
/// The product count is checked after each round (`elems^N` can far exceed [`MAX_EXPAND`]).
fn pow_cartesian(elems: Vec<Ty>, n: usize) -> Ty {
    let mut combos = vec![vec![]];
    for _ in 0..n {
        let mut next = vec![];
        for existing in &combos {
            for elem in &elems {
                let mut extended = existing.clone();
                extended.push(elem.clone());
                next.push(extended);
            }
        }
        if let Some(e) = check_expand_limit("tuple Cartesian product", next.len()) {
            return e;
        }
        combos = next;
    }
    TyArray(combos.into_iter().map(instantiate_combo).collect()).into()
}

/// Instantiate one Cartesian combination: TypeParam positions get a fresh param with the bound
/// preserved; other positions stay as-is
fn instantiate_combo(elems: Vec<Ty>) -> Ty {
    let mut tuple_elems = vec![];
    let mut param_decls = vec![];
    for elem in elems {
        let elem_span = elem.span;
        match elem.kind {
            TyKind::TypeParam(tp) => {
                let name = fresh_param();
                // Keep the original bound (previously the param name was mistaken for the bound;
                // `(A: Clone, T)^N` produced `_Param: A` instead of `_Param: Clone`)
                let params = tp
                    .params
                    .iter()
                    .map(|(_, bound)| (name.clone(), bound.clone()))
                    .collect();
                param_decls.push(TyTypeParam { params, bindings: vec![] });
                tuple_elems.push(TyPrimitive(name).to_ty().with_span(elem_span));
            }
            _ => tuple_elems.push(Ty::new(elem_span, elem.kind)),
        }
    }
    let tuple = TyTuple(tuple_elems).into();
    if param_decls.is_empty() {
        return tuple;
    }
    let mut merged = TyTypeParam { params: vec![], bindings: vec![] };
    for tp in param_decls {
        merged.extend(tp);
    }
    Ty::new(Span::call_site(), TyKind::TypeParam(merged)).apply(tuple)
}

fn fresh_params(n: usize) -> Vec<Ty> {
    (0..n).map(|_| TyPrimitive(fresh_param()).into()).collect()
}

impl TyTuple {
    /// `(A,B,)^C` => append C; `(A,)^N` => tuple-length expansion; `(A,)^N..M` => range expansion
    pub(crate) fn apply_help(mut self, o: Ty, span: Span) -> Ty {
        match o.kind {
            TyKind::Num(TyNum(n)) => tuple_pow(self.0, n),
            _ => {
                self.0.push(o);
                self.to_ty().with_span(span)
            }
        }
    }
}

impl TyGroup {
    /// `(T)` strips to the inner type, so `(T)^N` equals `T^N` (e.g.
    /// `(W)^2` = `W^2` = `W<2>`, a const-generic argument; only valid for
    /// types with a const generic — `(u8)^2` would emit `u8<2>`, which
    /// rustc rejects with E0109). Tuple generation needs `(T,)^N`.
    /// A bare `<T>` group is rejected earlier by parsing (`<` right after
    /// `(` is not a type).
    pub(crate) fn apply_help(self, o: Ty, _span: Span) -> Ty {
        self.0.apply(o)
    }
}

impl TyFn {
    /// `fn^(A,B)` => `fn(A,B)` (fills in params); `fn(A,B)-C` => `fn(A,B)->C` (adds return type).
    /// The `is_unsafe` field passes through (`unsafe fn^(A,B)` => `unsafe fn(A,B)`).
    pub(crate) fn apply_help(self, o: Ty, span: Span) -> Ty {
        match self {
            // A bare fn gets its params via `^`; the right side must be a tuple (a Group like
            // `fn^((i8,i16))` is unwrapped by the default apply's Group branch; here `o` is always plain)
            TyFn(None, None, is_unsafe) => match o.kind {
                TyKind::Tuple(t) => {
                    TyFn(t.0.into(), None, is_unsafe).to_ty().with_span(span)
                }
                _ => err_ty_at(
                    "batch-impl: the right side of the `fn` prefix must be a tuple type, e.g. fn^(i32, u32)",
                    span,
                ),
            },
            // Has params: append the return type via `-`
            TyFn(Some(params), None, is_unsafe) => {
                TyFn(params.into(), o.into(), is_unsafe).to_ty().with_span(span)
            }
            TyFn(Some(_), Some(_), _) => err_ty_at(
                "batch-impl: the `fn` type already has a return type; cannot apply again",
                span,
            ),
            // Impossible: params None but return Some
            TyFn(None, Some(_), _) => err_ty_at(
                "batch-impl: the `fn` type is missing a parameter list; internal error",
                span,
            ),
        }
    }
}

impl TyWithAttr {
    /// `#[attr]^T` => `#[attr] T` (attaches the attribute to the type)
    pub(crate) fn apply_help(self, o: Ty, span: Span) -> Ty {
        TyWithAttr(self.0, o.into()).to_ty().with_span(span)
    }
}

impl TyTypeParam {
    /// `<T>^U` => `WithType(<T>, U)` (generic parameters applied to the target type)
    pub(crate) fn apply_help(self, o: Ty, span: Span) -> Ty {
        TyWithType(self, o.into()).to_ty().with_span(span)
    }
}
impl TyNum {
    /// A number cannot be a left operand (used only on the right, e.g. `T^3`)
    pub(crate) fn apply_help(self, _: Ty, span: Span) -> Ty {
        err_ty_at(
            &format!(
                "batch-impl: number `{}` cannot be a left operand; use it on the right (e.g. T^{})",
                self.0, self.0
            ),
            span,
        )
    }
}
impl TyRange {
    /// A range cannot be a left operand (used only on the right, e.g. `T^1..3`)
    pub(crate) fn apply_help(self, _: Ty, span: Span) -> Ty {
        let end_mark = if self.inclusive { "=" } else { "" };
        err_ty_at(
            &format!(
                "batch-impl: range `{}..{}{}` cannot be a left operand; it goes on the right (e.g. T^{}..{}{})",
                self.start, self.end, end_mark, self.start, self.end, end_mark
            ),
            span,
        )
    }
}
impl TyPrimitiveArray {
    /// `[]^T` => `[T]` (empty base wraps a slice); `[T]^N` => `[T; N]` (fixed-size array)
    ///
    /// The length right side can be a numeric literal (`[u8]^3`), a const generic (`[u8]^N`), or a
    /// list/range (expanded item-wise by the top-level right-operand dispatch); re-applying to
    /// a finished array is an error.
    pub(crate) fn apply_help(self, o: Ty, span: Span) -> Ty {
        match (self.0, self.1) {
            (None, None) => Ty::new(span, TyPrimitiveArray(o.into(), None).into()),
            (Some(elem), None) => {
                TyPrimitiveArray(elem.into(), o.to_token_stream().into())
                    .to_ty()
                    .with_span(span)
            }
            _ => err_ty_at(
                "batch-impl: fixed-size array `[T; N]` cannot be a left operand",
                span,
            ),
        }
    }
}

/// Macro generating the "passthrough to the inner type then re-wrap" apply impls for the four
/// wrapper kinds (WithTrait/WithType always have an inner type, WithCode/WithWhere have an optional
/// one); their apply_help bodies are isomorphic, folded into a macro: adding a wrapper is one line,
/// and the "outer apply passes through to the inner target" semantics is declared once at the macro
/// definition site. Optional-inner: with `None`, the right operand directly takes the inner slot.
/// Note: `self.1` is written in the macro body (definition-site hygiene), not passable as
/// a macro argument — argument tokens keep the call-site context, and `self` would resolve to the
/// module's self (E0424).
macro_rules! impl_apply_optional_inner {
    ($ty:ident, $variant:ident) => {
        impl $ty {
            pub(crate) fn apply_help(self, o: Ty, span: Span) -> Ty {
                let inner = match self.0 {
                    Some(t) => t.apply(o),
                    None => o,
                };
                Ty::new(span, TyKind::$variant($ty(inner.into(), self.1)))
            }
        }
    };
}
/// Always-inner form: apply the right operand to the inner type, then re-wrap.
macro_rules! impl_apply_inner {
    ($ty:ident, $variant:ident) => {
        impl $ty {
            pub(crate) fn apply_help(self, o: Ty, span: Span) -> Ty {
                Ty::new(span, TyKind::$variant($ty(self.0, self.1.apply(o).into())))
            }
        }
    };
}

impl_apply_inner!(TyWithTrait, WithTrait);
impl_apply_inner!(TyWithType, WithType);
impl_apply_optional_inner!(TyWithCode, WithCode);
impl_apply_optional_inner!(TyWithWhere, WithWhere);