module Control.DeepSeq.TH
( deriveNFData
, deriveNFDatas
, whnfIsNf
) where
import Control.DeepSeq (NFData(rnf),deepseq)
import Control.Monad (mzero,liftM,mplus)
import Data.List
import Data.Maybe (fromMaybe, isJust, catMaybes)
import Language.Haskell.TH
whnfIsNf :: Type -> Maybe Bool
whnfIsNf (ConT x)
| x `elem` [''Int, ''Double, ''Float, ''Char, ''Bool, ''()] = Just True
whnfIsNf (AppT ListT _) = Just False
whnfIsNf (AppT (TupleT _) _) = Just False
whnfIsNf _ = Nothing
whnfIsNf' :: Type -> Bool
whnfIsNf' = fromMaybe False . whnfIsNf
deriveNFData :: Name -> Q [Dec]
deriveNFData tn = do
dec <- reify tn
case dec of
TyConI (DataD _ctx _tn tvs ctors _) -> do
clauses_names <- mapM con2rnf ctors
let clauses = map fst clauses_names
names = nub $ concat $ map snd clauses_names
ctxt = [ClassP ''NFData [VarT n] | n <- names ]
let ity = foldl (\t tvn -> AppT t (VarT tvn)) (ConT tn) $ map tyvarname tvs
return [ InstanceD ctxt (AppT (ConT ''NFData) ity) [FunD 'rnf clauses] ]
TyConI (NewtypeD {}) -> do
fail $ "deriveNFData ''" ++ show tn ++ ": please use GeneralizedNewtypeDeriving " ++
"for deriving NFData instances for newtype"
TyConI (TySynD {}) -> do
fail $ "deriveNFData ''" ++ show tn ++ ": cannot derive for type-alias"
TyConI _ -> do
fail $ "deriveNFData ''" ++ show tn ++ ": argument must be a proper 'data'-type"
_ -> do
fail $ "deriveNFData ''" ++ show tn ++ ": argument must be a type-level entity"
where
tyvarname (PlainTV n) = n
tyvarname (KindedTV n _) = n
tys2vars = mapM (\t -> if isJust t then liftM VarP (newName "x") else return WildP)
con2rnf :: Con -> Q (Clause, [Name])
con2rnf (NormalC n ts) = genCon2Rnf n ts
con2rnf (RecC n vts) = genCon2Rnf n [ (tst,tt) | (_,tst,tt) <- vts ]
con2rnf (InfixC tl n tr) = genCon2Rnf n [tl,tr]
con2rnf (ForallC {}) = fail "deriveNFData: 'forall' not supported in constructor declaration"
genCon2Rnf :: Name -> [(Strict,Type)] -> Q (Clause, [Name])
genCon2Rnf n ts = do
let vns = concatMap getFreeTyVars $ catMaybes ts'
ts' = [ if tst == NotStrict || not (whnfIsNf' tt) then Just tt else Nothing | (tst,tt) <- ts ]
vars <- tys2vars ts'
return (Clause [ConP n vars] (NormalB $ mkDeepSeqExpr [ n' | VarP n' <- vars ]) [], vns)
deriveNFDatas :: [Name] -> Q [Dec]
deriveNFDatas = liftM concat . mapM deriveNFData
getFreeTyVars :: Type -> [Name]
getFreeTyVars (AppT t1 t2) = getFreeTyVars t1 `mplus` getFreeTyVars t2
getFreeTyVars (ArrowT) = mzero
getFreeTyVars (ConT _) = mzero
getFreeTyVars (ForallT {}) = error "getFreeTyVars: ForallT not supported yet"
getFreeTyVars (ListT) = mzero
getFreeTyVars (SigT t1 _) = getFreeTyVars t1
getFreeTyVars (TupleT _) = mzero
getFreeTyVars (UnboxedTupleT _) = mzero
getFreeTyVars (VarT n) = return n
mkDeepSeqExpr :: [Name] -> Exp
mkDeepSeqExpr = foldr deepSeqE (ConE '())
where
deepSeqE :: Name -> Exp -> Exp
deepSeqE lhs rhs = AppE (AppE (VarE 'deepseq) (VarE lhs)) rhs