    {-# 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, substForallTelescope, substTyVarBndrs,
      unionSubsts, unionMaybeSubsts,

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

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

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

    -- | 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 :: forall (q :: * -> *). Quasi q => DSubst -> DType -> q DType
    substTy DSubst
    vars (DForallT DForallTelescope
    tele DType
    ty) = do
      (vars', tele') <- DSubst -> DForallTelescope -> q (DSubst, DForallTelescope)
    forall (q :: * -> *).
    Quasi q =>
    DSubst -> DForallTelescope -> q (DSubst, DForallTelescope)
    substForallTelescope DSubst
    vars DForallTelescope
    tele
      ty' <- substTy vars' ty
      return $ DForallT tele' ty'
    substTy DSubst
    vars (DConstrainedT DCxt
    cxt DType
    ty) =
      DCxt -> DType -> DType
    DConstrainedT (DCxt -> DType -> DType) -> q DCxt -> q (DType -> DType)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (DType -> q DType) -> DCxt -> q DCxt
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    (a -> m b) -> t a -> m (t b)
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
    mapM (DSubst -> DType -> q DType
    forall (q :: * -> *). Quasi q => DSubst -> DType -> q DType
    substTy DSubst
    vars) DCxt
    cxt q (DType -> DType) -> q DType -> q DType
    forall a b. q (a -> b) -> q a -> q b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> DSubst -> DType -> q DType
    forall (q :: * -> *). Quasi q => DSubst -> DType -> q DType
    substTy DSubst
    vars DType
    ty
    substTy DSubst
    vars (DAppT DType
    t1 DType
    t2) =
      DType -> DType -> DType
    DAppT (DType -> DType -> DType) -> q DType -> q (DType -> DType)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> DSubst -> DType -> q DType
    forall (q :: * -> *). Quasi q => DSubst -> DType -> q DType
    substTy DSubst
    vars DType
    t1 q (DType -> DType) -> q DType -> q DType
    forall a b. q (a -> b) -> q a -> q b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> DSubst -> DType -> q DType
    forall (q :: * -> *). Quasi q => DSubst -> DType -> q DType
    substTy DSubst
    vars DType
    t2
    substTy DSubst
    vars (DAppKindT DType
    t DType
    k) =
      DType -> DType -> DType
    DAppKindT (DType -> DType -> DType) -> q DType -> q (DType -> DType)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> DSubst -> DType -> q DType
    forall (q :: * -> *). Quasi q => DSubst -> DType -> q DType
    substTy DSubst
    vars DType
    t q (DType -> DType) -> q DType -> q DType
    forall a b. q (a -> b) -> q a -> q b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> DSubst -> DType -> q DType
    forall (q :: * -> *). Quasi q => DSubst -> DType -> q DType
    substTy DSubst
    vars DType
    k
    substTy DSubst
    vars (DSigT DType
    ty DType
    ki) =
      DType -> DType -> DType
    DSigT (DType -> DType -> DType) -> q DType -> q (DType -> DType)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> DSubst -> DType -> q DType
    forall (q :: * -> *). Quasi q => DSubst -> DType -> q DType
    substTy DSubst
    vars DType
    ty q (DType -> DType) -> q DType -> q DType
    forall a b. q (a -> b) -> q a -> q b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> DSubst -> DType -> q DType
    forall (q :: * -> *). Quasi q => DSubst -> DType -> q DType
    substTy DSubst
    vars DType
    ki
    substTy DSubst
    vars (DVarT Name
    n)
      | Just DType
    ty <- Name -> DSubst -> Maybe DType
    forall k a. Ord k => k -> Map k a -> Maybe a
    M.lookup Name
    n DSubst
    vars
      = DType -> q DType
    forall a. a -> q a
    forall (m :: * -> *) a. Monad m => a -> m a
    return DType
    ty
      | Bool
    otherwise
      = DType -> q DType
    forall a. a -> q a
    forall (m :: * -> *) a. Monad m => a -> m a
    return (DType -> q DType) -> DType -> q DType
    forall a b. (a -> b) -> a -> b
    $ Name -> DType
    DVarT Name
    n
    substTy DSubst
    _ ty :: DType
    ty@(DConT Name
    _)  = DType -> q DType
    forall a. a -> q a
    forall (m :: * -> *) a. Monad m => a -> m a
    return DType
    ty
    substTy DSubst
    _ ty :: DType
    ty@DType
    DArrowT    = DType -> q DType
    forall a. a -> q a
    forall (m :: * -> *) a. Monad m => a -> m a
    return DType
    ty
    substTy DSubst
    _ ty :: DType
    ty@(DLitT TyLit
    _)  = DType -> q DType
    forall a. a -> q a
    forall (m :: * -> *) a. Monad m => a -> m a
    return DType
    ty
    substTy DSubst
    _ ty :: DType
    ty@DType
    DWildCardT = DType -> q DType
    forall a. a -> q a
    forall (m :: * -> *) a. Monad m => a -> m a
    return DType
    ty

    substForallTelescope :: Quasi q => DSubst -> DForallTelescope
                         -> q (DSubst, DForallTelescope)
    substForallTelescope :: forall (q :: * -> *).
    Quasi q =>
    DSubst -> DForallTelescope -> q (DSubst, DForallTelescope)
    substForallTelescope DSubst
    vars DForallTelescope
    tele =
      case DForallTelescope
    tele of
        DForallVis [DTyVarBndrUnit]
    tvbs -> do
          (vars', tvbs') <- DSubst -> [DTyVarBndrUnit] -> q (DSubst, [DTyVarBndrUnit])
    forall (q :: * -> *) flag.
    Quasi q =>
    DSubst -> [DTyVarBndr flag] -> q (DSubst, [DTyVarBndr flag])
    substTyVarBndrs DSubst
    vars [DTyVarBndrUnit]
    tvbs
          return (vars', DForallVis tvbs')
        DForallInvis [DTyVarBndrSpec]
    tvbs -> do
          (vars', tvbs') <- DSubst -> [DTyVarBndrSpec] -> q (DSubst, [DTyVarBndrSpec])
    forall (q :: * -> *) flag.
    Quasi q =>
    DSubst -> [DTyVarBndr flag] -> q (DSubst, [DTyVarBndr flag])
    substTyVarBndrs DSubst
    vars [DTyVarBndrSpec]
    tvbs
          return (vars', DForallInvis tvbs')

    substTyVarBndrs :: Quasi q => DSubst -> [DTyVarBndr flag]
                    -> q (DSubst, [DTyVarBndr flag])
    substTyVarBndrs :: forall (q :: * -> *) flag.
    Quasi q =>
    DSubst -> [DTyVarBndr flag] -> q (DSubst, [DTyVarBndr flag])
    substTyVarBndrs = (DSubst -> DTyVarBndr flag -> q (DSubst, DTyVarBndr flag))
    -> DSubst -> [DTyVarBndr flag] -> q (DSubst, [DTyVarBndr flag])
    forall (m :: * -> *) acc x y.
    Monad m =>
    (acc -> x -> m (acc, y)) -> acc -> [x] -> m (acc, [y])
    mapAccumLM DSubst -> DTyVarBndr flag -> q (DSubst, DTyVarBndr flag)
    forall (q :: * -> *) flag.
    Quasi q =>
    DSubst -> DTyVarBndr flag -> q (DSubst, DTyVarBndr flag)
    substTvb

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

    -- | Computes the union of two substitutions. Fails if both subsitutions map
    -- the same variable to different types.
    unionSubsts :: DSubst -> DSubst -> Maybe DSubst
    unionSubsts :: DSubst -> DSubst -> Maybe DSubst
    unionSubsts DSubst
    a DSubst
    b =
      let shared_key_set :: Set Name
    shared_key_set = DSubst -> Set Name
    forall k a. Map k a -> Set k
    M.keysSet DSubst
    a Set Name -> Set Name -> Set Name
    forall a. Ord a => Set a -> Set a -> Set a
    `S.intersection` DSubst -> Set Name
    forall k a. Map k a -> Set k
    M.keysSet DSubst
    b
          matches_up :: Bool
    matches_up     = (Name -> Bool -> Bool) -> Bool -> Set Name -> Bool
    forall a b. (a -> b -> b) -> b -> Set a -> b
    S.foldr (\Name
    name -> ((DSubst
    a DSubst -> Name -> DType
    forall k a. Ord k => Map k a -> k -> a
    M.! Name
    name) DType -> DType -> Bool
    forall a. Eq a => a -> a -> Bool
    == (DSubst
    b DSubst -> Name -> DType
    forall k a. Ord k => Map k a -> k -> a
    M.! Name
    name) Bool -> Bool -> Bool
    &&))
                                   Bool
    True Set Name
    shared_key_set
      in
      if Bool
    matches_up then DSubst -> Maybe DSubst
    forall a. a -> Maybe a
    forall (m :: * -> *) a. Monad m => a -> m a
    return (DSubst
    a DSubst -> DSubst -> DSubst
    forall k a. Ord k => Map k a -> Map k a -> Map k a
    `M.union` DSubst
    b) else Maybe DSubst
    forall a. Maybe a
    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 :: IgnoreKinds -> DType -> DType -> Maybe DSubst
    matchTy IgnoreKinds
    _   (DVarT Name
    var_name) DType
    arg = DSubst -> Maybe DSubst
    forall a. a -> Maybe a
    Just (DSubst -> Maybe DSubst) -> DSubst -> Maybe DSubst
    forall a b. (a -> b) -> a -> b
    $ Name -> DType -> DSubst
    forall k a. k -> a -> Map k a
    M.singleton Name
    var_name DType
    arg
      -- if a pattern has a kind signature, it's really easy to get
      -- the following two cases wrong.
    matchTy IgnoreKinds
    ign (DSigT DType
    ty DType
    _ki) DType
    arg = case IgnoreKinds
    ign of
      IgnoreKinds
    YesIgnore -> IgnoreKinds -> DType -> DType -> Maybe DSubst
    matchTy IgnoreKinds
    ign DType
    ty DType
    arg
      IgnoreKinds
    NoIgnore  -> Maybe DSubst
    forall a. Maybe a
    Nothing
    matchTy IgnoreKinds
    ign (DAppKindT DType
    ty DType
    _ki) DType
    arg = case IgnoreKinds
    ign of
      IgnoreKinds
    YesIgnore -> IgnoreKinds -> DType -> DType -> Maybe DSubst
    matchTy IgnoreKinds
    ign DType
    ty DType
    arg
      IgnoreKinds
    NoIgnore  -> Maybe DSubst
    forall a. Maybe a
    Nothing
      -- but we can safely ignore kind signatures on the target,
      -- as in the following two cases.
    matchTy IgnoreKinds
    ign DType
    pat (DSigT     DType
    ty DType
    _ki) = IgnoreKinds -> DType -> DType -> Maybe DSubst
    matchTy IgnoreKinds
    ign DType
    pat DType
    ty
    matchTy IgnoreKinds
    ign DType
    pat (DAppKindT DType
    ty DType
    _ki) = IgnoreKinds -> DType -> DType -> Maybe DSubst
    matchTy IgnoreKinds
    ign DType
    pat DType
    ty
    matchTy IgnoreKinds
    _   (DForallT {}) DType
    _ =
      String -> Maybe DSubst
    forall a. HasCallStack => String -> a
    error String
    "Cannot match a forall in a pattern"
    matchTy IgnoreKinds
    _   DType
    _ (DForallT {}) =
      String -> Maybe DSubst
    forall a. HasCallStack => String -> a
    error String
    "Cannot match a forall in a target"
    matchTy IgnoreKinds
    ign (DAppT DType
    pat1 DType
    pat2) (DAppT DType
    arg1 DType
    arg2) =
      [Maybe DSubst] -> Maybe DSubst
    unionMaybeSubsts [IgnoreKinds -> DType -> DType -> Maybe DSubst
    matchTy IgnoreKinds
    ign DType
    pat1 DType
    arg1, IgnoreKinds -> DType -> DType -> Maybe DSubst
    matchTy IgnoreKinds
    ign DType
    pat2 DType
    arg2]
    matchTy IgnoreKinds
    _   (DConT Name
    pat_con) (DConT Name
    arg_con)
      | Name
    pat_con Name -> Name -> Bool
    forall a. Eq a => a -> a -> Bool
    == Name
    arg_con = DSubst -> Maybe DSubst
    forall a. a -> Maybe a
    Just DSubst
    forall k a. Map k a
    M.empty
    matchTy IgnoreKinds
    _   DType
    DArrowT DType
    DArrowT = DSubst -> Maybe DSubst
    forall a. a -> Maybe a
    Just DSubst
    forall k a. Map k a
    M.empty
    matchTy IgnoreKinds
    _   (DLitT TyLit
    pat_lit) (DLitT TyLit
    arg_lit)
      | TyLit
    pat_lit TyLit -> TyLit -> Bool
    forall a. Eq a => a -> a -> Bool
    == TyLit
    arg_lit = DSubst -> Maybe DSubst
    forall a. a -> Maybe a
    Just DSubst
    forall k a. Map k a
    M.empty
    matchTy IgnoreKinds
    _ DType
    _ DType
    _ = Maybe DSubst
    forall a. Maybe a
    Nothing

    unionMaybeSubsts :: [Maybe DSubst] -> Maybe DSubst
    unionMaybeSubsts :: [Maybe DSubst] -> Maybe DSubst
    unionMaybeSubsts = (Maybe DSubst -> Maybe DSubst -> Maybe DSubst)
    -> Maybe DSubst -> [Maybe DSubst] -> Maybe DSubst
    forall b a. (b -> a -> b) -> b -> [a] -> b
    forall (t :: * -> *) b a.
    Foldable t =>
    (b -> a -> b) -> b -> t a -> b
    L.foldl' Maybe DSubst -> Maybe DSubst -> Maybe DSubst
    union_subst1 (DSubst -> Maybe DSubst
    forall a. a -> Maybe a
    Just DSubst
    forall k a. Map k a
    M.empty)
      where
        union_subst1 :: Maybe DSubst -> Maybe DSubst -> Maybe DSubst
        union_subst1 :: Maybe DSubst -> Maybe DSubst -> Maybe DSubst
    union_subst1 Maybe DSubst
    ma Maybe DSubst
    mb = do
          a <- Maybe DSubst
    ma
          b <- mb
          unionSubsts a b