//! Many kinds of items or constructs can have generic parameters: functions,
//! structs, impls, traits, etc. This module provides a common HIR for these
//! generic parameters. See also the `Generics` type and the `generics_of` query
//! in rustc.

use either::Either;
use hir_expand::name::{AsName, Name};
use intern::sym;
use la_arena::Arena;
use syntax::ast::{self, HasName, HasTypeBounds};
use thin_vec::ThinVec;

use crate::{
    GenericDefId, TypeOrConstParamId, TypeParamId,
    expr_store::{
        TypePtr,
        lower::{ExprCollector, LifetimeBoundScope, NamedLifetimeStore},
    },
    hir::generics::{
        ConstParamData, GenericParams, LifetimeBoundType, LifetimeParamData, TypeOrConstParamData,
        TypeParamData, TypeParamProvenance, WherePredicate,
    },
    type_ref::{LifetimeRef, LifetimeRefId, TypeBound, TypeRef, TypeRefId},
};

pub(crate) type ImplTraitLowerFn<'l> = &'l mut dyn for<'ec, 'db> FnMut(
    &'ec mut ExprCollector<'db>,
    TypePtr,
    ThinVec<TypeBound>,
) -> TypeRefId;

pub(crate) struct GenericParamsCollector {
    type_or_consts: Arena<TypeOrConstParamData>,
    lifetimes: Arena<LifetimeParamData>,
    where_predicates: Vec<WherePredicate>,
    parent: GenericDefId,
}

impl GenericParamsCollector {
    pub(crate) fn new(parent: GenericDefId) -> Self {
        Self {
            type_or_consts: Default::default(),
            lifetimes: Default::default(),
            where_predicates: Default::default(),
            parent,
        }
    }
    pub(crate) fn with_self_param(
        ec: &mut ExprCollector<'_>,
        parent: GenericDefId,
        bounds: Option<ast::TypeBoundList>,
    ) -> Self {
        let mut this = Self::new(parent);
        this.fill_self_param(ec, bounds);
        this
    }

    pub(crate) fn lower(
        &mut self,
        ec: &mut ExprCollector<'_>,
        generic_param_list: Option<ast::GenericParamList>,
        where_clause: Option<ast::WhereClause>,
    ) {
        if let Some(params) = generic_param_list {
            self.lower_param_list(ec, params)
        }
        if let Some(where_clause) = where_clause {
            self.lower_where_predicates(ec, where_clause)
        }
    }

    pub(crate) fn collect_impl_trait<R>(
        &mut self,
        ec: &mut ExprCollector<'_>,
        cb: impl FnOnce(&mut ExprCollector<'_>, ImplTraitLowerFn<'_>) -> R,
    ) -> R {
        cb(
            ec,
            &mut Self::lower_argument_impl_trait(
                &mut self.type_or_consts,
                &mut self.where_predicates,
                self.parent,
            ),
        )
    }

    pub(crate) fn finish(self) -> GenericParams {
        let Self { mut lifetimes, mut type_or_consts, where_predicates, parent: _ } = self;

        lifetimes.shrink_to_fit();
        type_or_consts.shrink_to_fit();
        GenericParams {
            type_or_consts,
            lifetimes,
            where_predicates: where_predicates.into_boxed_slice(),
        }
    }

    fn lower_param_list(&mut self, ec: &mut ExprCollector<'_>, params: ast::GenericParamList) {
        for generic_param in params.generic_params() {
            let enabled = ec.check_cfg(&generic_param);
            if !enabled {
                continue;
            }

            match generic_param {
                ast::GenericParam::TypeParam(type_param) => {
                    let name = type_param.name().map_or_else(Name::missing, |it| it.as_name());
                    let default = type_param.default_type().map(|it| {
                        ec.lower_type_ref(it, &mut ExprCollector::impl_trait_error_allocator)
                    });
                    let param = TypeParamData {
                        name: Some(name.clone()),
                        default,
                        provenance: TypeParamProvenance::TypeParamList,
                    };
                    let idx = self.type_or_consts.alloc(param.into());
                    let type_ref =
                        TypeRef::TypeParam(TypeParamId::from_unchecked(TypeOrConstParamId {
                            parent: self.parent,
                            local_id: idx,
                        }));
                    let type_ref = ec.alloc_type_ref_desugared(type_ref);
                    ec.with_lifetime_bound_scope(LifetimeBoundScope::WhereClause, |ec| {
                        self.lower_bounds(ec, type_param.type_bound_list(), Either::Left(type_ref))
                    });
                }
                ast::GenericParam::ConstParam(const_param) => {
                    let name = const_param.name().map_or_else(Name::missing, |it| it.as_name());
                    let ty = ec.lower_type_ref_opt(
                        const_param.ty(),
                        &mut ExprCollector::impl_trait_error_allocator,
                    );
                    let default = const_param.default_val().map(|it| ec.lower_const_arg(it));
                    let param = ConstParamData { name, ty, default };
                    self.type_or_consts.alloc(param.into());
                }
                ast::GenericParam::LifetimeParam(lifetime_param) => {
                    let lifetime = ec.lower_lifetime_ref_opt(lifetime_param.lifetime());
                    if let LifetimeRef::Named(name) = &ec.store.lifetimes[lifetime] {
                        let param = LifetimeParamData {
                            name: name.clone(),
                            bound_type: LifetimeBoundType::EarlyBound,
                        };
                        let _idx = self.lifetimes.alloc(param);
                        ec.with_lifetime_bound_scope(LifetimeBoundScope::WhereClause, |ec| {
                            self.lower_bounds(
                                ec,
                                lifetime_param.type_bound_list(),
                                Either::Right(lifetime),
                            )
                        });
                    }
                }
            }
        }
    }

    fn lower_where_predicates(
        &mut self,
        ec: &mut ExprCollector<'_>,
        where_clause: ast::WhereClause,
    ) {
        ec.with_lifetime_bound_scope(LifetimeBoundScope::WhereClause, |ec| {
            for pred in where_clause.predicates() {
                let target = if let Some(type_ref) = pred.ty() {
                    Either::Left(
                        ec.lower_type_ref(type_ref, &mut ExprCollector::impl_trait_error_allocator),
                    )
                } else if let Some(lifetime) = pred.lifetime() {
                    Either::Right(ec.lower_lifetime_ref(lifetime))
                } else {
                    continue;
                };

                let lifetimes: Option<Box<_>> =
                    pred.for_binder().and_then(|it| it.generic_param_list()).map(|param_list| {
                        // Higher-Ranked Trait Bounds
                        param_list
                            .lifetime_params()
                            .map(|lifetime_param| {
                                lifetime_param
                                    .lifetime()
                                    .map_or_else(Name::missing, |lt| Name::new_lifetime(lt.text()))
                            })
                            .collect()
                    });
                for bound in pred.type_bound_list().iter().flat_map(|l| l.bounds()) {
                    self.lower_type_bound_as_predicate(ec, bound, lifetimes.as_deref(), target);
                }
            }
        });
    }

    fn lower_bounds(
        &mut self,
        ec: &mut ExprCollector<'_>,
        type_bounds: Option<ast::TypeBoundList>,
        target: Either<TypeRefId, LifetimeRefId>,
    ) {
        for bound in type_bounds.iter().flat_map(|type_bound_list| type_bound_list.bounds()) {
            self.lower_type_bound_as_predicate(ec, bound, None, target);
        }
    }

    fn lower_type_bound_as_predicate(
        &mut self,
        ec: &mut ExprCollector<'_>,
        bound: ast::TypeBound,
        hrtb_lifetimes: Option<&[Name]>,
        target: Either<TypeRefId, LifetimeRefId>,
    ) {
        let bound = ec.lower_type_bound(
            bound,
            &mut Self::lower_argument_impl_trait(
                &mut self.type_or_consts,
                &mut self.where_predicates,
                self.parent,
            ),
        );
        let predicate = match (target, bound) {
            (_, TypeBound::Error | TypeBound::Use(_)) => return,
            (Either::Left(type_ref), bound) => WherePredicate::TypeBound {
                lifetimes: hrtb_lifetimes.map(|h| ThinVec::from_iter(h.iter().cloned())),
                target: type_ref,
                bound,
            },
            (Either::Right(lifetime), TypeBound::Lifetime(bound)) => {
                WherePredicate::Lifetime { target: lifetime, bound }
            }
            (Either::Right(_), TypeBound::ForLifetime(..) | TypeBound::Path(..)) => return,
        };
        if let WherePredicate::Lifetime { target, .. } = predicate {
            ec.push_named_target_lifetime(target);
        }
        self.where_predicates.push(predicate);
    }

    fn lower_argument_impl_trait(
        type_or_consts: &mut Arena<TypeOrConstParamData>,
        where_predicates: &mut Vec<WherePredicate>,
        parent: GenericDefId,
    ) -> impl for<'ec, 'db> FnMut(&'ec mut ExprCollector<'db>, TypePtr, ThinVec<TypeBound>) -> TypeRefId
    {
        move |ec, ptr, impl_trait_bounds| {
            let param = TypeParamData {
                name: None,
                default: None,
                provenance: TypeParamProvenance::ArgumentImplTrait,
            };
            let param_id = TypeRef::TypeParam(TypeParamId::from_unchecked(TypeOrConstParamId {
                parent,
                local_id: type_or_consts.alloc(param.into()),
            }));
            let type_ref = ec.alloc_type_ref(param_id, ptr);
            for bound in impl_trait_bounds {
                where_predicates.push(WherePredicate::TypeBound {
                    lifetimes: None,
                    target: type_ref,
                    bound: bound.clone(),
                });
            }
            type_ref
        }
    }

    fn fill_self_param(&mut self, ec: &mut ExprCollector<'_>, bounds: Option<ast::TypeBoundList>) {
        let self_ = Name::new_symbol_root(sym::Self_);
        let idx = self.type_or_consts.alloc(
            TypeParamData {
                name: Some(self_),
                default: None,
                provenance: TypeParamProvenance::TraitSelf,
            }
            .into(),
        );
        debug_assert_eq!(idx, GenericParams::SELF_PARAM_ID_IN_SELF);
        let type_ref = TypeRef::TypeParam(TypeParamId::from_unchecked(TypeOrConstParamId {
            parent: self.parent,
            local_id: idx,
        }));
        let self_ = ec.alloc_type_ref_desugared(type_ref);
        if let Some(bounds) = bounds {
            self.lower_bounds(ec, Some(bounds), Either::Left(self_));
        }
    }

    pub(crate) fn update_to_late_bound_lifetimes(
        &mut self,
        named_lifetime_store: &NamedLifetimeStore,
    ) {
        for (_param_id, lifetime) in self.lifetimes.iter_mut() {
            let lifetime_name = &lifetime.name;
            if named_lifetime_store.lifetimes_in_where_clause.contains(lifetime_name) {
                continue;
            }

            if !named_lifetime_store.lifetimes_constrained_by_input.contains(lifetime_name)
                && named_lifetime_store.lifetimes_in_output.contains(lifetime_name)
            {
                continue;
            }

            lifetime.bound_type = LifetimeBoundType::LateBound
        }
    }
}
