{-# LANGUAGE GADTs #-}
{-# LANGUAGE TemplateHaskell #-}
module Data.Comp.Derive.Arbitrary
(
ArbitraryF(..),
makeArbitraryF,
Arbitrary(..)
)where
import Data.Comp.Derive.Utils hiding (derive)
import Language.Haskell.TH
import Test.QuickCheck
class ArbitraryF f where
arbitraryF' :: Arbitrary v => [(Int,Gen (f v))]
arbitraryF' = [(1,arbitraryF)]
arbitraryF :: Arbitrary v => Gen (f v)
arbitraryF = frequency arbitraryF'
shrinkF :: Arbitrary v => f v -> [f v]
shrinkF _ = []
makeArbitraryF :: Name -> Q [Dec]
makeArbitraryF dt = do
Just (DataInfo _cxt name args constrs _deriving) <- abstractNewtypeQ $ reify dt
let argNames = map (VarT . tyVarBndrName) (tail args)
complType = foldl AppT (ConT name) argNames
preCond = map (mkClassP ''Arbitrary . (: [])) argNames
classType = AppT (ConT ''ArbitraryF) complType
arbitraryDecl <- generateArbitraryFDecl constrs
shrinkDecl <- generateShrinkFDecl constrs
return [mkInstanceD preCond classType [arbitraryDecl, shrinkDecl]]
generateArbitraryFDecl :: [Con] -> Q Dec
generateArbitraryFDecl = generateGenDecl 'arbitraryF'
generateGenDecl :: Name -> [Con] -> Q Dec
generateGenDecl genName constrs
= do genBody <- listE $ map (addNum . constrGen . abstractConType) constrs
let genClause = Clause [] (NormalB genBody) []
return $ FunD genName [genClause]
where addNum e = [| (1,$e) |]
constrGen :: (Name,Int) -> ExpQ
constrGen (constr, n)
= do varNs <- newNames n "x"
newSizeN <- newName "newSize"
let newSizeE = varE newSizeN
let newSizeP = varP newSizeN
let constrsE = litE . IntegerL . toInteger $ n
let binds = (`map` varNs) (\var -> bindS
(varP var)
[| resize $newSizeE arbitrary |] )
let apps = appsE (conE constr: map varE varNs)
let build = doE $
binds ++
[noBindS [|return $apps|]]
if n == 0
then [|return $apps|]
else [| sized $ \ size ->
$(letE [valD
newSizeP
(normalB [|((size - 1) `div` $constrsE ) `max` 0|])
[] ]
build) |]
generateShrinkFDecl :: [Con] -> Q Dec
generateShrinkFDecl constrs
= let clauses = map (generateClause.abstractConType) constrs
in funD 'shrinkF clauses
where generateClause (constr, n)
= do varNs <- newNames n "x"
resVarNs <- newNames n "x'"
binds <- mapM (\(var,resVar) -> bindS (varP resVar) [| $(varE var) : shrink $(varE var) |]) $ zip varNs resVarNs
let ret = NoBindS $ AppE (VarE 'return) (foldl1 AppE ( ConE constr: map VarE resVarNs ))
stmtSeq = binds ++ [ret]
pat = ConP constr $ map VarP varNs
return $ Clause [pat] (NormalB $ AppE (VarE 'tail) (DoE stmtSeq)) []