//! The `TypeWalk` trait (probably to be replaced by Chalk's `Fold` and
//! `Visit`).

use std::mem;

use chalk_ir::DebruijnIndex;

use crate::{
    utils::make_mut_slice, AliasTy, Binders, CallableSig, GenericArg, GenericArgData, Interner,
    OpaqueTy, ProjectionTy, Substitution, TraitRef, Ty, TyKind, WhereClause,
};

/// This allows walking structures that contain types to do something with those
/// types, similar to Chalk's `Fold` trait.
pub trait TypeWalk {
    fn walk(&self, f: &mut impl FnMut(&Ty));
    fn walk_mut(&mut self, f: &mut impl FnMut(&mut Ty)) {
        self.walk_mut_binders(&mut |ty, _binders| f(ty), DebruijnIndex::INNERMOST);
    }
    /// Walk the type, counting entered binders.
    ///
    /// `TyKind::Bound` variables use DeBruijn indexing, which means that 0 refers
    /// to the innermost binder, 1 to the next, etc.. So when we want to
    /// substitute a certain bound variable, we can't just walk the whole type
    /// and blindly replace each instance of a certain index; when we 'enter'
    /// things that introduce new bound variables, we have to keep track of
    /// that. Currently, the only thing that introduces bound variables on our
    /// side are `TyKind::Dyn` and `TyKind::Opaque`, which each introduce a bound
    /// variable for the self type.
    fn walk_mut_binders(
        &mut self,
        f: &mut impl FnMut(&mut Ty, DebruijnIndex),
        binders: DebruijnIndex,
    );

    fn fold_binders(
        mut self,
        f: &mut impl FnMut(Ty, DebruijnIndex) -> Ty,
        binders: DebruijnIndex,
    ) -> Self
    where
        Self: Sized,
    {
        self.walk_mut_binders(
            &mut |ty_mut, binders| {
                let ty = mem::replace(ty_mut, TyKind::Unknown.intern(&Interner));
                *ty_mut = f(ty, binders);
            },
            binders,
        );
        self
    }

    fn fold(mut self, f: &mut impl FnMut(Ty) -> Ty) -> Self
    where
        Self: Sized,
    {
        self.walk_mut(&mut |ty_mut| {
            let ty = mem::replace(ty_mut, TyKind::Unknown.intern(&Interner));
            *ty_mut = f(ty);
        });
        self
    }

    /// Substitutes `TyKind::Bound` vars with the given substitution.
    fn subst_bound_vars(self, substs: &Substitution) -> Self
    where
        Self: Sized,
    {
        self.subst_bound_vars_at_depth(substs, DebruijnIndex::INNERMOST)
    }

    /// Substitutes `TyKind::Bound` vars with the given substitution.
    fn subst_bound_vars_at_depth(mut self, substs: &Substitution, depth: DebruijnIndex) -> Self
    where
        Self: Sized,
    {
        self.walk_mut_binders(
            &mut |ty, binders| {
                if let &mut TyKind::BoundVar(bound) = ty.interned_mut() {
                    if bound.debruijn >= binders {
                        *ty = substs.interned()[bound.index]
                            .assert_ty_ref(&Interner)
                            .clone()
                            .shift_bound_vars(binders);
                    }
                }
            },
            depth,
        );
        self
    }

    /// Shifts up debruijn indices of `TyKind::Bound` vars by `n`.
    fn shift_bound_vars(self, n: DebruijnIndex) -> Self
    where
        Self: Sized,
    {
        self.fold_binders(
            &mut |ty, binders| match ty.kind(&Interner) {
                TyKind::BoundVar(bound) if bound.debruijn >= binders => {
                    TyKind::BoundVar(bound.shifted_in_from(n)).intern(&Interner)
                }
                _ => ty,
            },
            DebruijnIndex::INNERMOST,
        )
    }

    /// Shifts debruijn indices of `TyKind::Bound` vars out (down) by `n`.
    fn shift_bound_vars_out(self, n: DebruijnIndex) -> Self
    where
        Self: Sized + std::fmt::Debug,
    {
        self.fold_binders(
            &mut |ty, binders| match ty.kind(&Interner) {
                TyKind::BoundVar(bound) if bound.debruijn >= binders => {
                    TyKind::BoundVar(bound.shifted_out_to(n).unwrap_or(bound.clone()))
                        .intern(&Interner)
                }
                _ => ty,
            },
            DebruijnIndex::INNERMOST,
        )
    }
}

impl TypeWalk for Ty {
    fn walk(&self, f: &mut impl FnMut(&Ty)) {
        match self.kind(&Interner) {
            TyKind::Alias(AliasTy::Projection(p_ty)) => {
                for t in p_ty.substitution.iter(&Interner) {
                    t.walk(f);
                }
            }
            TyKind::Alias(AliasTy::Opaque(o_ty)) => {
                for t in o_ty.substitution.iter(&Interner) {
                    t.walk(f);
                }
            }
            TyKind::Dyn(dyn_ty) => {
                for p in dyn_ty.bounds.value.interned().iter() {
                    p.walk(f);
                }
            }
            TyKind::Slice(ty) | TyKind::Array(ty) | TyKind::Ref(_, ty) | TyKind::Raw(_, ty) => {
                ty.walk(f);
            }
            _ => {
                if let Some(substs) = self.substs() {
                    for t in substs.iter(&Interner) {
                        t.walk(f);
                    }
                }
            }
        }
        f(self);
    }

    fn walk_mut_binders(
        &mut self,
        f: &mut impl FnMut(&mut Ty, DebruijnIndex),
        binders: DebruijnIndex,
    ) {
        match self.interned_mut() {
            TyKind::Alias(AliasTy::Projection(p_ty)) => {
                p_ty.substitution.walk_mut_binders(f, binders);
            }
            TyKind::Dyn(dyn_ty) => {
                for p in make_mut_slice(dyn_ty.bounds.value.interned_mut()) {
                    p.walk_mut_binders(f, binders.shifted_in());
                }
            }
            TyKind::Alias(AliasTy::Opaque(o_ty)) => {
                o_ty.substitution.walk_mut_binders(f, binders);
            }
            TyKind::Slice(ty) | TyKind::Array(ty) | TyKind::Ref(_, ty) | TyKind::Raw(_, ty) => {
                ty.walk_mut_binders(f, binders);
            }
            _ => {
                if let Some(substs) = self.substs_mut() {
                    substs.walk_mut_binders(f, binders);
                }
            }
        }
        f(self, binders);
    }
}

impl<T: TypeWalk> TypeWalk for Vec<T> {
    fn walk(&self, f: &mut impl FnMut(&Ty)) {
        for t in self {
            t.walk(f);
        }
    }
    fn walk_mut_binders(
        &mut self,
        f: &mut impl FnMut(&mut Ty, DebruijnIndex),
        binders: DebruijnIndex,
    ) {
        for t in self {
            t.walk_mut_binders(f, binders);
        }
    }
}

impl TypeWalk for OpaqueTy {
    fn walk(&self, f: &mut impl FnMut(&Ty)) {
        self.substitution.walk(f);
    }

    fn walk_mut_binders(
        &mut self,
        f: &mut impl FnMut(&mut Ty, DebruijnIndex),
        binders: DebruijnIndex,
    ) {
        self.substitution.walk_mut_binders(f, binders);
    }
}

impl TypeWalk for ProjectionTy {
    fn walk(&self, f: &mut impl FnMut(&Ty)) {
        self.substitution.walk(f);
    }

    fn walk_mut_binders(
        &mut self,
        f: &mut impl FnMut(&mut Ty, DebruijnIndex),
        binders: DebruijnIndex,
    ) {
        self.substitution.walk_mut_binders(f, binders);
    }
}

impl TypeWalk for AliasTy {
    fn walk(&self, f: &mut impl FnMut(&Ty)) {
        match self {
            AliasTy::Projection(it) => it.walk(f),
            AliasTy::Opaque(it) => it.walk(f),
        }
    }

    fn walk_mut_binders(
        &mut self,
        f: &mut impl FnMut(&mut Ty, DebruijnIndex),
        binders: DebruijnIndex,
    ) {
        match self {
            AliasTy::Projection(it) => it.walk_mut_binders(f, binders),
            AliasTy::Opaque(it) => it.walk_mut_binders(f, binders),
        }
    }
}

impl TypeWalk for GenericArg {
    fn walk(&self, f: &mut impl FnMut(&Ty)) {
        match &self.interned() {
            GenericArgData::Ty(ty) => {
                ty.walk(f);
            }
        }
    }

    fn walk_mut_binders(
        &mut self,
        f: &mut impl FnMut(&mut Ty, DebruijnIndex),
        binders: DebruijnIndex,
    ) {
        match self.interned_mut() {
            GenericArgData::Ty(ty) => {
                ty.walk_mut_binders(f, binders);
            }
        }
    }
}

impl TypeWalk for Substitution {
    fn walk(&self, f: &mut impl FnMut(&Ty)) {
        for t in self.iter(&Interner) {
            t.walk(f);
        }
    }

    fn walk_mut_binders(
        &mut self,
        f: &mut impl FnMut(&mut Ty, DebruijnIndex),
        binders: DebruijnIndex,
    ) {
        for t in self.interned_mut() {
            t.walk_mut_binders(f, binders);
        }
    }
}

impl<T: TypeWalk> TypeWalk for Binders<T> {
    fn walk(&self, f: &mut impl FnMut(&Ty)) {
        self.value.walk(f);
    }

    fn walk_mut_binders(
        &mut self,
        f: &mut impl FnMut(&mut Ty, DebruijnIndex),
        binders: DebruijnIndex,
    ) {
        self.value.walk_mut_binders(f, binders.shifted_in())
    }
}

impl TypeWalk for TraitRef {
    fn walk(&self, f: &mut impl FnMut(&Ty)) {
        self.substitution.walk(f);
    }

    fn walk_mut_binders(
        &mut self,
        f: &mut impl FnMut(&mut Ty, DebruijnIndex),
        binders: DebruijnIndex,
    ) {
        self.substitution.walk_mut_binders(f, binders);
    }
}

impl TypeWalk for WhereClause {
    fn walk(&self, f: &mut impl FnMut(&Ty)) {
        match self {
            WhereClause::Implemented(trait_ref) => trait_ref.walk(f),
            WhereClause::AliasEq(alias_eq) => alias_eq.walk(f),
        }
    }

    fn walk_mut_binders(
        &mut self,
        f: &mut impl FnMut(&mut Ty, DebruijnIndex),
        binders: DebruijnIndex,
    ) {
        match self {
            WhereClause::Implemented(trait_ref) => trait_ref.walk_mut_binders(f, binders),
            WhereClause::AliasEq(alias_eq) => alias_eq.walk_mut_binders(f, binders),
        }
    }
}

impl TypeWalk for CallableSig {
    fn walk(&self, f: &mut impl FnMut(&Ty)) {
        for t in self.params_and_return.iter() {
            t.walk(f);
        }
    }

    fn walk_mut_binders(
        &mut self,
        f: &mut impl FnMut(&mut Ty, DebruijnIndex),
        binders: DebruijnIndex,
    ) {
        for t in make_mut_slice(&mut self.params_and_return) {
            t.walk_mut_binders(f, binders);
        }
    }
}