    {-# LANGUAGE CPP #-}

    -----------------------------------------------------------------------------
    -- |
    -- Module      :  Language.Haskell.TH.Desugar.Subst
    -- Copyright   :  (C) 2018 Richard Eisenberg
    -- License     :  BSD-style (see LICENSE)
    -- Maintainer  :  Ryan Scott
    -- Stability   :  experimental
    -- Portability :  non-portable
    --
    -- Capture-avoiding substitutions on 'DType's
    --
    ----------------------------------------------------------------------------

    module Language.Haskell.TH.Desugar.Subst (
      DSubst,

      -- * Capture-avoiding substitution
      substTy, substTyVarBndrs, unionSubsts, unionMaybeSubsts,

      -- * Matching a type template against a type
      IgnoreKinds(..), matchTy
      ) where

    import qualified Data.Map as M
    import qualified Data.Set as S

    import Data.Generics
    import Data.List

    import Language.Haskell.TH.Desugar.AST
    import Language.Haskell.TH.Desugar.Core
    import Language.Haskell.TH.Syntax
    import Language.Haskell.TH.Desugar.Util

    #if __GLASGOW_HASKELL__ < 710
    import Control.Applicative
    #endif

    -- | A substitution is just a map from names to types
    type DSubst = M.Map Name DType

    -- | Capture-avoiding substitution on types
    substTy :: Quasi q => DSubst -> DType -> q DType
    substTy vars (DForallT tvbs cxt ty) =
      substTyVarBndrs vars tvbs $ \vars' tvbs' -> do
        cxt' <- mapM (substPred vars') cxt
        ty' <- substTy vars' ty
        return $ DForallT tvbs' cxt' ty'
    substTy vars (DAppT t1 t2) =
      DAppT <$> substTy vars t1 <*> substTy vars t2
    substTy vars (DSigT ty ki) =
      DSigT <$> substTy vars ty <*> substTy vars ki
    substTy vars (DVarT n)
      | Just ty <- M.lookup n vars
      = return ty
      | otherwise
      = return $ DVarT n
    substTy _ ty@(DConT _)  = return ty
    substTy _ ty@DArrowT    = return ty
    substTy _ ty@(DLitT _)  = return ty
    substTy _ ty@DWildCardT = return ty

    substTyVarBndrs :: Quasi q => DSubst -> [DTyVarBndr]
                    -> (DSubst -> [DTyVarBndr] -> q a)
                    -> q a
    substTyVarBndrs vars tvbs thing = do
      (vars', tvbs') <- mapAccumLM substTvb vars tvbs
      thing vars' tvbs'

    substTvb :: Quasi q => DSubst -> DTyVarBndr
             -> q (DSubst, DTyVarBndr)
    substTvb vars (DPlainTV n) = do
      new_n <- qNewName (nameBase n)
      return (M.insert n (DVarT new_n) vars, DPlainTV new_n)
    substTvb vars (DKindedTV n k) = do
      new_n <- qNewName (nameBase n)
      k' <- substTy vars k
      return (M.insert n (DVarT new_n) vars, DKindedTV new_n k')

    substPred :: Quasi q => DSubst -> DPred -> q DPred
    substPred vars (DForallPr tvbs cxt p) =
      substTyVarBndrs vars tvbs $ \vars' tvbs' -> do
        cxt' <- mapM (substPred vars') cxt
        p'   <- substPred vars' p
        return $ DForallPr tvbs' cxt' p'
    substPred vars (DAppPr p t) = DAppPr <$> substPred vars p <*> substTy vars t
    substPred vars (DSigPr p k) = DSigPr <$> substPred vars p <*> substTy vars k
    substPred vars (DVarPr n)
      | Just ty <- M.lookup n vars
      = dTypeToDPred ty
      | otherwise
      = return $ DVarPr n
    substPred _ p@(DConPr {}) = return p
    substPred _ p@DWildCardPr = return p

    -- | Computes the union of two substitutions. Fails if both subsitutions map
    -- the same variable to different types.
    unionSubsts :: DSubst -> DSubst -> Maybe DSubst
    unionSubsts a b =
      let shared_key_set = M.keysSet a `S.intersection` M.keysSet b
          matches_up     = S.foldr (\name -> ((a M.! name) `geq` (b M.! name) &&))
                                   True shared_key_set
      in
      if matches_up then return (a `M.union` b) else Nothing

    ---------------------------
    -- Matching

    -- | Ignore kind annotations in @matchTy@?
    data IgnoreKinds = YesIgnore | NoIgnore

    -- | @matchTy ign tmpl targ@ matches a type template @tmpl@ against a type
    -- target @targ@. This returns a Map from names of type variables in the
    -- type template to types if the types indeed match up, or @Nothing@ otherwise.
    -- In the @Just@ case, it is guaranteed that every type variable mentioned
    -- in the template is mapped by the returned substitution.
    --
    -- The first argument @ign@ tells @matchTy@ whether to ignore kind signatures
    -- in the template. A kind signature in the template might mean that a type
    -- variable has a more restrictive kind than otherwise possible, and that
    -- mapping that type variable to a type of a different kind could be disastrous.
    -- So, if we don't ignore kind signatures, this function returns @Nothing@ if
    -- the template has a signature anywhere. If we do ignore kind signatures, it's
    -- possible the returned map will be ill-kinded. Use at your own risk.
    matchTy :: IgnoreKinds -> DType -> DType -> Maybe DSubst
    matchTy _   (DVarT var_name) arg = Just $ M.singleton var_name arg
      -- if a pattern has a kind signature, it's really easy to get
      -- this wrong.
    matchTy ign (DSigT ty _ki) arg = case ign of
      YesIgnore -> matchTy ign ty arg
      NoIgnore  -> Nothing
      -- but we can safely ignore kind signatures on the target
    matchTy ign pat (DSigT ty _ki) = matchTy ign pat ty
    matchTy _   (DForallT {}) _ =
      error "Cannot match a forall in a pattern"
    matchTy _   _ (DForallT {}) =
      error "Cannot match a forall in a target"
    matchTy ign (DAppT pat1 pat2) (DAppT arg1 arg2) =
      unionMaybeSubsts [matchTy ign pat1 arg1, matchTy ign pat2 arg2]
    matchTy _   (DConT pat_con) (DConT arg_con)
      | pat_con == arg_con = Just M.empty
    matchTy _   DArrowT DArrowT = Just M.empty
    matchTy _   (DLitT pat_lit) (DLitT arg_lit)
      | pat_lit == arg_lit = Just M.empty
    matchTy _ _ _ = Nothing

    unionMaybeSubsts :: [Maybe DSubst] -> Maybe DSubst
    unionMaybeSubsts = foldl' union_subst1 (Just M.empty)
      where
        union_subst1 :: Maybe DSubst -> Maybe DSubst -> Maybe DSubst
        union_subst1 ma mb = do
          a <- ma
          b <- mb
          unionSubsts a b