{- Language/Haskell/TH/Desugar/Expand.hs

(c) Richard Eisenberg 2013
rae@cs.brynmawr.edu
-}

{-# LANGUAGE NoMonomorphismRestriction, ScopedTypeVariables #-}

-----------------------------------------------------------------------------

-- |

-- Module      :  Language.Haskell.TH.Desugar.Expand

-- Copyright   :  (C) 2014 Richard Eisenberg

-- License     :  BSD-style (see LICENSE)

-- Maintainer  :  Ryan Scott

-- Stability   :  experimental

-- Portability :  non-portable

--

-- Expands type synonyms and type families in desugared types.

-- See also the package th-expand-syns for doing this to

-- non-desugared types.

--

----------------------------------------------------------------------------


module Language.Haskell.TH.Desugar.Expand (
  -- * Expand synonyms soundly

  expand, expandType,

  -- * Expand synonyms potentially unsoundly

  expandUnsoundly
  ) where

import qualified Data.Map as M
import Language.Haskell.TH hiding (cxt)
import Language.Haskell.TH.Syntax ( Quasi(..) )
import Data.Data
import Data.Generics
import qualified Data.Traversable as T

import Language.Haskell.TH.Desugar.AST
import Language.Haskell.TH.Desugar.Core
import Language.Haskell.TH.Desugar.Util
import Language.Haskell.TH.Desugar.Sweeten
import Language.Haskell.TH.Desugar.Reify
import Language.Haskell.TH.Desugar.Subst

-- | Expands all type synonyms in a desugared type. Also expands open type family

-- applications. (In GHCs before 7.10, this part does not work if there are any

-- variables.) Attempts to

-- expand closed type family applications, but aborts the moment it spots anything

-- strange, like a nested type family application or type variable.

expandType :: DsMonad q => DType -> q DType
expandType :: forall (q :: * -> *). DsMonad q => DType -> q DType
expandType = forall (q :: * -> *). DsMonad q => IgnoreKinds -> DType -> q DType
expand_type IgnoreKinds
NoIgnore

expand_type :: forall q. DsMonad q => IgnoreKinds -> DType -> q DType
expand_type :: forall (q :: * -> *). DsMonad q => IgnoreKinds -> DType -> q DType
expand_type IgnoreKinds
ign = [DTypeArg] -> DType -> q DType
go []
  where
    go :: [DTypeArg] -> DType -> q DType
    go :: [DTypeArg] -> DType -> q DType
go [] (DForallT DForallTelescope
tele DType
ty) =
      DForallTelescope -> DType -> DType
DForallT forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (q :: * -> *).
DsMonad q =>
IgnoreKinds -> DForallTelescope -> q DForallTelescope
expand_tele IgnoreKinds
ign DForallTelescope
tele
               forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (q :: * -> *). DsMonad q => IgnoreKinds -> DType -> q DType
expand_type IgnoreKinds
ign DType
ty
    go [DTypeArg]
_ (DForallT {}) =
      forall (q :: * -> *) a. MonadFail q => String -> q a
impossible String
"A forall type is applied to another type."
    go [] (DConstrainedT DCxt
cxt DType
ty) =
      DCxt -> DType -> DType
DConstrainedT forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall (q :: * -> *). DsMonad q => IgnoreKinds -> DType -> q DType
expand_type IgnoreKinds
ign) DCxt
cxt
                    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (q :: * -> *). DsMonad q => IgnoreKinds -> DType -> q DType
expand_type IgnoreKinds
ign DType
ty
    go [DTypeArg]
_ (DConstrainedT {}) =
      forall (q :: * -> *) a. MonadFail q => String -> q a
impossible String
"A constrained type is applied to another type."
    go [DTypeArg]
args (DAppT DType
t1 DType
t2) = do
      DType
t2' <- forall (q :: * -> *). DsMonad q => IgnoreKinds -> DType -> q DType
expand_type IgnoreKinds
ign DType
t2
      [DTypeArg] -> DType -> q DType
go (DType -> DTypeArg
DTANormal DType
t2' forall a. a -> [a] -> [a]
: [DTypeArg]
args) DType
t1
    go [DTypeArg]
args (DAppKindT DType
p DType
k) = do
      DType
k' <- forall (q :: * -> *). DsMonad q => IgnoreKinds -> DType -> q DType
expand_type IgnoreKinds
ign DType
k
      [DTypeArg] -> DType -> q DType
go (DType -> DTypeArg
DTyArg DType
k' forall a. a -> [a] -> [a]
: [DTypeArg]
args) DType
p
    go [DTypeArg]
args (DSigT DType
ty DType
ki) = do
      DType
ty' <- [DTypeArg] -> DType -> q DType
go [] DType
ty
      DType
ki' <- [DTypeArg] -> DType -> q DType
go [] DType
ki
      DType -> [DTypeArg] -> q DType
finish (DType -> DType -> DType
DSigT DType
ty' DType
ki') [DTypeArg]
args
    go [DTypeArg]
args (DConT Name
n) = forall (q :: * -> *).
DsMonad q =>
IgnoreKinds -> Name -> [DTypeArg] -> q DType
expand_con IgnoreKinds
ign Name
n [DTypeArg]
args
    go [DTypeArg]
args ty :: DType
ty@(DVarT Name
_)  = DType -> [DTypeArg] -> q DType
finish DType
ty [DTypeArg]
args
    go [DTypeArg]
args ty :: DType
ty@DType
DArrowT    = DType -> [DTypeArg] -> q DType
finish DType
ty [DTypeArg]
args
    go [DTypeArg]
args ty :: DType
ty@(DLitT TyLit
_)  = DType -> [DTypeArg] -> q DType
finish DType
ty [DTypeArg]
args
    go [DTypeArg]
args ty :: DType
ty@DType
DWildCardT = DType -> [DTypeArg] -> q DType
finish DType
ty [DTypeArg]
args

    finish :: DType -> [DTypeArg] -> q DType
    finish :: DType -> [DTypeArg] -> q DType
finish DType
ty [DTypeArg]
args = forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ DType -> [DTypeArg] -> DType
applyDType DType
ty [DTypeArg]
args

-- | Expands all type synonyms in the kinds of a @forall@ telescope.

expand_tele :: DsMonad q => IgnoreKinds -> DForallTelescope -> q DForallTelescope
expand_tele :: forall (q :: * -> *).
DsMonad q =>
IgnoreKinds -> DForallTelescope -> q DForallTelescope
expand_tele IgnoreKinds
ign (DForallVis   [DTyVarBndrUnit]
tvbs) = [DTyVarBndrUnit] -> DForallTelescope
DForallVis   forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall (q :: * -> *) flag.
DsMonad q =>
IgnoreKinds -> DTyVarBndr flag -> q (DTyVarBndr flag)
expand_tvb IgnoreKinds
ign) [DTyVarBndrUnit]
tvbs
expand_tele IgnoreKinds
ign (DForallInvis [DTyVarBndrSpec]
tvbs) = [DTyVarBndrSpec] -> DForallTelescope
DForallInvis forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall (q :: * -> *) flag.
DsMonad q =>
IgnoreKinds -> DTyVarBndr flag -> q (DTyVarBndr flag)
expand_tvb IgnoreKinds
ign) [DTyVarBndrSpec]
tvbs

-- | Expands all type synonyms in a type variable binder's kind.

expand_tvb :: DsMonad q => IgnoreKinds -> DTyVarBndr flag -> q (DTyVarBndr flag)
expand_tvb :: forall (q :: * -> *) flag.
DsMonad q =>
IgnoreKinds -> DTyVarBndr flag -> q (DTyVarBndr flag)
expand_tvb IgnoreKinds
_   tvb :: DTyVarBndr flag
tvb@DPlainTV{}       = forall (f :: * -> *) a. Applicative f => a -> f a
pure DTyVarBndr flag
tvb
expand_tvb IgnoreKinds
ign (DKindedTV Name
n flag
flag DType
k) = forall flag. Name -> flag -> DType -> DTyVarBndr flag
DKindedTV Name
n flag
flag forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (q :: * -> *). DsMonad q => IgnoreKinds -> DType -> q DType
expand_type IgnoreKinds
ign DType
k

-- | Expand a constructor with given arguments

expand_con :: forall q.
              DsMonad q
           => IgnoreKinds
           -> Name       -- ^ Tycon name

           -> [DTypeArg] -- ^ Arguments

           -> q DType    -- ^ Expanded type

expand_con :: forall (q :: * -> *).
DsMonad q =>
IgnoreKinds -> Name -> [DTypeArg] -> q DType
expand_con IgnoreKinds
ign Name
n [DTypeArg]
args = do
  Info
info <- forall (q :: * -> *). DsMonad q => Name -> q Info
reifyWithLocals Name
n
  case Info
info of
    TyConI (TySynD Name
_ [TyVarBndr ()]
_ Type
StarT)
         -- See Note [Don't expand synonyms for *]

      -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ DType -> [DTypeArg] -> DType
applyDType (Name -> DType
DConT Name
typeKindName) [DTypeArg]
args
    Info
_ -> Info -> q DType
go Info
info
  where
    -- Only the normal (i.e., non-visibly applied) arguments. These are

    -- important since we need to align these with the arguments of the type

    -- synonym/family, and visible kind arguments can mess with this.

    normal_args :: [DType]
    normal_args :: DCxt
normal_args = [DTypeArg] -> DCxt
filterDTANormals [DTypeArg]
args

    go :: Info -> q DType
    go :: Info -> q DType
go Info
info = do
      DInfo
dinfo <- forall (q :: * -> *). DsMonad q => Info -> q DInfo
dsInfo Info
info
      case DInfo
dinfo of
        DTyConI (DTySynD Name
_n [DTyVarBndrUnit]
tvbs DType
rhs) Maybe [DDec]
_
          |  forall (t :: * -> *) a. Foldable t => t a -> Int
length DCxt
normal_args forall a. Ord a => a -> a -> Bool
>= forall (t :: * -> *) a. Foldable t => t a -> Int
length [DTyVarBndrUnit]
tvbs   -- this should always be true!

          -> do
            let (DCxt
syn_args, DCxt
rest_args) = forall a b. [a] -> [b] -> ([b], [b])
splitAtList [DTyVarBndrUnit]
tvbs DCxt
normal_args
            DType
ty <- forall (q :: * -> *). Quasi q => DSubst -> DType -> q DType
substTy (forall k a. Ord k => [(k, a)] -> Map k a
M.fromList forall a b. (a -> b) -> a -> b
$ forall a b. [a] -> [b] -> [(a, b)]
zip (forall a b. (a -> b) -> [a] -> [b]
map forall flag. DTyVarBndr flag -> Name
dtvbName [DTyVarBndrUnit]
tvbs) DCxt
syn_args) DType
rhs
            DType
ty' <- forall (q :: * -> *). DsMonad q => IgnoreKinds -> DType -> q DType
expand_type IgnoreKinds
ign DType
ty
            forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ DType -> [DTypeArg] -> DType
applyDType DType
ty' forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map DType -> DTypeArg
DTANormal DCxt
rest_args

        DTyConI (DOpenTypeFamilyD (DTypeFamilyHead Name
_n [DTyVarBndrUnit]
tvbs DFamilyResultSig
_frs Maybe InjectivityAnn
_ann)) Maybe [DDec]
_
          |  forall (t :: * -> *) a. Foldable t => t a -> Int
length DCxt
normal_args forall a. Ord a => a -> a -> Bool
>= forall (t :: * -> *) a. Foldable t => t a -> Int
length [DTyVarBndrUnit]
tvbs   -- this should always be true!

          -> do
            let (DCxt
syn_args, DCxt
rest_args) = forall a b. [a] -> [b] -> ([b], [b])
splitAtList [DTyVarBndrUnit]
tvbs DCxt
normal_args
            -- We need to get the correct instance. If we fail to reify anything

            -- (e.g., if a type family is quasiquoted), then fall back by

            -- pretending that there are no instances in scope.

            [Dec]
insts <- forall (m :: * -> *) a. Quasi m => m a -> m a -> m a
qRecover (forall (m :: * -> *) a. Monad m => a -> m a
return []) forall a b. (a -> b) -> a -> b
$
                     forall (m :: * -> *). Quasi m => Name -> [Type] -> m [Dec]
qReifyInstances Name
n (forall a b. (a -> b) -> [a] -> [b]
map DType -> Type
typeToTH DCxt
syn_args)
            [DDec]
dinsts <- forall (q :: * -> *). DsMonad q => [Dec] -> q [DDec]
dsDecs [Dec]
insts
            case [DDec]
dinsts of
              [DTySynInstD (DTySynEqn Maybe [DTyVarBndrUnit]
_ DType
lhs DType
rhs)]
                |  (DType
_, [DTypeArg]
lhs_args) <- DType -> (DType, [DTypeArg])
unfoldDType DType
lhs
                ,  let lhs_normal_args :: DCxt
lhs_normal_args = [DTypeArg] -> DCxt
filterDTANormals [DTypeArg]
lhs_args
                ,  Just DSubst
subst <-
                     [Maybe DSubst] -> Maybe DSubst
unionMaybeSubsts forall a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (IgnoreKinds -> DType -> DType -> Maybe DSubst
matchTy IgnoreKinds
ign) DCxt
lhs_normal_args DCxt
syn_args
                -> do DType
ty <- forall (q :: * -> *). Quasi q => DSubst -> DType -> q DType
substTy DSubst
subst DType
rhs
                      DType
ty' <- forall (q :: * -> *). DsMonad q => IgnoreKinds -> DType -> q DType
expand_type IgnoreKinds
ign DType
ty
                      forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ DType -> [DTypeArg] -> DType
applyDType DType
ty' forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map DType -> DTypeArg
DTANormal DCxt
rest_args
              [DDec]
_ -> q DType
give_up


        DTyConI (DClosedTypeFamilyD (DTypeFamilyHead Name
_n [DTyVarBndrUnit]
tvbs DFamilyResultSig
_frs Maybe InjectivityAnn
_ann) [DTySynEqn]
eqns) Maybe [DDec]
_
          |  forall (t :: * -> *) a. Foldable t => t a -> Int
length DCxt
normal_args forall a. Ord a => a -> a -> Bool
>= forall (t :: * -> *) a. Foldable t => t a -> Int
length [DTyVarBndrUnit]
tvbs
          -> do
            let (DCxt
syn_args, DCxt
rest_args) = forall a b. [a] -> [b] -> ([b], [b])
splitAtList [DTyVarBndrUnit]
tvbs DCxt
normal_args
            DCxt
rhss <- forall (m :: * -> *) a b.
Monad m =>
(a -> m (Maybe b)) -> [a] -> m [b]
mapMaybeM (DCxt -> DTySynEqn -> q (Maybe DType)
check_eqn DCxt
syn_args) [DTySynEqn]
eqns
            case DCxt
rhss of
              (DType
rhs : DCxt
_) -> do
                DType
rhs' <- forall (q :: * -> *). DsMonad q => IgnoreKinds -> DType -> q DType
expand_type IgnoreKinds
ign DType
rhs
                forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ DType -> [DTypeArg] -> DType
applyDType DType
rhs' forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map DType -> DTypeArg
DTANormal DCxt
rest_args
              [] -> q DType
give_up

          where
             -- returns the substed rhs

            check_eqn :: [DType] -> DTySynEqn -> q (Maybe DType)
            check_eqn :: DCxt -> DTySynEqn -> q (Maybe DType)
check_eqn DCxt
arg_tys (DTySynEqn Maybe [DTyVarBndrUnit]
_ DType
lhs DType
rhs) = do
              let (DType
_, [DTypeArg]
lhs_args) = DType -> (DType, [DTypeArg])
unfoldDType DType
lhs
                  normal_lhs_args :: DCxt
normal_lhs_args = [DTypeArg] -> DCxt
filterDTANormals [DTypeArg]
lhs_args
                  m_subst :: Maybe DSubst
m_subst = [Maybe DSubst] -> Maybe DSubst
unionMaybeSubsts forall a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (IgnoreKinds -> DType -> DType -> Maybe DSubst
matchTy IgnoreKinds
ign) DCxt
normal_lhs_args DCxt
arg_tys
              forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
T.mapM (forall a b c. (a -> b -> c) -> b -> a -> c
flip forall (q :: * -> *). Quasi q => DSubst -> DType -> q DType
substTy DType
rhs) Maybe DSubst
m_subst

        DInfo
_ -> q DType
give_up

    -- Used when we can't proceed with type family instance expansion any more,

    -- and must conservatively return the orignal type family applied to its

    -- arguments.

    give_up :: q DType
    give_up :: q DType
give_up = forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ DType -> [DTypeArg] -> DType
applyDType (Name -> DType
DConT Name
n) [DTypeArg]
args

{-
Note [Don't expand synonyms for *]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
We deliberately avoid expanding type synonyms for * such as Type and ★.
Why? If you reify any such type synonym using Template Haskell, this is
what you'll get:

  TyConI (TySynD <type synonym name> [] StarT)

If you blindly charge ahead and recursively inspect the right-hand side of
this type synonym, you'll desugar StarT into (DConT ''Type), reify ''Type,
and get back another type synonym with StarT as its right-hand side. Then
you'll recursively inspect StarT and find yourself knee-deep in an infinite
loop.

To prevent these sorts of shenanigans, we simply stop whenever we see a type
synonym with StarT as its right-hand side and return Type.
-}

-- | Expand all type synonyms and type families in the desugared abstract

-- syntax tree provided, where type family simplification is on a "best effort"

-- basis. Normally, the first parameter should have a type like

-- 'DExp' or 'DLetDec'.

expand :: (DsMonad q, Data a) => a -> q a
expand :: forall (q :: * -> *) a. (DsMonad q, Data a) => a -> q a
expand = forall (q :: * -> *) a.
(DsMonad q, Data a) =>
IgnoreKinds -> a -> q a
expand_ IgnoreKinds
NoIgnore

-- | Expand all type synonyms and type families in the desugared abstract

-- syntax tree provided, where type family simplification is on a "better

-- than best effort" basis. This means that it will try so hard that it will

-- sometimes do the wrong thing. Specifically, any kind parameters to type

-- families are ignored. So, if we have

--

-- > type family F (x :: k) where

-- >   F (a :: *) = Int

--

-- 'expandUnsoundly' will expand @F 'True@ to @Int@, ignoring that the

-- expansion should only work for type of kind @*@.

--

-- This function is useful because plain old 'expand' will simply fail

-- to expand type families that make use of kinds. Sometimes, the kinds

-- are benign and we want to expand anyway. Use this function in that case.

expandUnsoundly :: (DsMonad q, Data a) => a -> q a
expandUnsoundly :: forall (q :: * -> *) a. (DsMonad q, Data a) => a -> q a
expandUnsoundly = forall (q :: * -> *) a.
(DsMonad q, Data a) =>
IgnoreKinds -> a -> q a
expand_ IgnoreKinds
YesIgnore

-- | Generalization of 'expand' that either can or won't ignore kind annotations.sx

expand_ :: (DsMonad q, Data a) => IgnoreKinds -> a -> q a
expand_ :: forall (q :: * -> *) a.
(DsMonad q, Data a) =>
IgnoreKinds -> a -> q a
expand_ IgnoreKinds
ign = forall (m :: * -> *). Monad m => GenericM m -> GenericM m
everywhereM (forall (m :: * -> *) a b.
(Monad m, Typeable a, Typeable b) =>
(b -> m b) -> a -> m a
mkM (forall (q :: * -> *). DsMonad q => IgnoreKinds -> DType -> q DType
expand_type IgnoreKinds
ign))