mockcat-1.0.0.0: src/Test/MockCat/TH/TypeUtils.hs

{-# LANGUAGE TemplateHaskellQuotes #-}
module Test.MockCat.TH.TypeUtils
  ( splitApps,
    substituteType,
    isNotConstantFunctionType,
    collectTypeVars,
    needsTypeable,
    collectTypeableTargets,
    isStandardTypeCon
  )
where

import qualified Data.Map.Strict as Map
import Language.Haskell.TH (Name, Type (..))
import Test.MockCat.Param (Param)
import Test.MockCat.Cons ((:>))

splitApps :: Type -> (Type, [Type])
splitApps ty = go ty []
  where
    go (Language.Haskell.TH.AppT t1 t2) acc = go t1 (t2 : acc)
    go t acc = (t, acc)

substituteType :: Map.Map Language.Haskell.TH.Name Language.Haskell.TH.Type -> Language.Haskell.TH.Type -> Language.Haskell.TH.Type
substituteType subMap = go
  where
    go (Language.Haskell.TH.VarT name) = Map.findWithDefault (Language.Haskell.TH.VarT name) name subMap
    go (Language.Haskell.TH.AppT t1 t2) = Language.Haskell.TH.AppT (go t1) (go t2)
    go (Language.Haskell.TH.SigT t k) = Language.Haskell.TH.SigT (go t) k
    go (Language.Haskell.TH.ParensT t) = Language.Haskell.TH.ParensT (go t)
    go (Language.Haskell.TH.InfixT t1 n t2) = Language.Haskell.TH.InfixT (go t1) n (go t2)
    go (Language.Haskell.TH.UInfixT t1 n t2) = Language.Haskell.TH.UInfixT (go t1) n (go t2)
    go (Language.Haskell.TH.ForallT tvs ctx t) = Language.Haskell.TH.ForallT tvs (map go ctx) (go t)
    go t = t

isNotConstantFunctionType :: Type -> Bool
isNotConstantFunctionType (AppT (AppT ArrowT _) _) = True
isNotConstantFunctionType (AppT t1 t2) = isNotConstantFunctionType t1 || isNotConstantFunctionType t2
isNotConstantFunctionType (TupleT _) = False
isNotConstantFunctionType (ForallT _ _ t) = isNotConstantFunctionType t
isNotConstantFunctionType _ = False

needsTypeable :: Type -> Bool
needsTypeable = go
  where
    go (ForallT _ _ t) = go t
    go (AppT t1 t2) = go t1 || go t2
    go (SigT t _) = go t
    go (VarT _) = True
    go (ParensT t) = go t
    go (InfixT t1 _ t2) = go t1 || go t2
    go (UInfixT t1 _ t2) = go t1 || go t2
    go (ImplicitParamT _ t) = go t
    go _ = False

collectTypeVars :: Type -> [Name]
collectTypeVars (VarT name) = [name]
collectTypeVars (AppT t1 t2) = collectTypeVars t1 ++ collectTypeVars t2
collectTypeVars (SigT t _) = collectTypeVars t
collectTypeVars (ParensT t) = collectTypeVars t
collectTypeVars (InfixT t1 _ t2) = collectTypeVars t1 ++ collectTypeVars t2
collectTypeVars (UInfixT t1 _ t2) = collectTypeVars t1 ++ collectTypeVars t2
collectTypeVars (ForallT _ _ t) = collectTypeVars t
collectTypeVars (ImplicitParamT _ t) = collectTypeVars t
collectTypeVars _ = []

collectTypeableTargets :: Type -> [Type]
collectTypeableTargets ty =
  case ty of
    VarT _ -> [ty]
    AppT _ _ ->
      let (f, args) = splitApps ty
      in if isStandardTypeCon f
         then concatMap collectTypeableTargets args
         else [ty]
    SigT t _ -> collectTypeableTargets t
    ParensT t -> collectTypeableTargets t
    InfixT t1 _ t2 -> collectTypeableTargets t1 ++ collectTypeableTargets t2
    UInfixT t1 _ t2 -> collectTypeableTargets t1 ++ collectTypeableTargets t2
    ForallT _ _ t -> collectTypeableTargets t
    _ -> []

isStandardTypeCon :: Type -> Bool
isStandardTypeCon ArrowT = True
isStandardTypeCon ListT = True
isStandardTypeCon (TupleT _) = True
isStandardTypeCon (ConT n) =
  n `elem`
    [ ''Maybe
    , ''IO
    , ''Either
    , ''[]
    , ''(,)
    , ''(,,)
    , ''(,,,)
    , ''(,,,,)
    , ''Param
    , ''(:>)
    ]
isStandardTypeCon _ = False