{-# 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' = [(Int
1,Gen (f v)
forall (f :: * -> *) v. (ArbitraryF f, Arbitrary v) => Gen (f v)
arbitraryF)]
arbitraryF :: Arbitrary v => Gen (f v)
arbitraryF = [(Int, Gen (f v))] -> Gen (f v)
forall a. [(Int, Gen a)] -> Gen a
frequency [(Int, Gen (f v))]
forall (f :: * -> *) v.
(ArbitraryF f, Arbitrary v) =>
[(Int, Gen (f v))]
arbitraryF'
shrinkF :: Arbitrary v => f v -> [f v]
shrinkF f v
_ = []
makeArbitraryF :: Name -> Q [Dec]
makeArbitraryF :: Name -> Q [Dec]
makeArbitraryF Name
dt = do
Just (DataInfo Cxt
_cxt Name
name [TyVarBndr]
args [Con]
constrs [DerivClause]
_deriving) <- Q Info -> Q (Maybe DataInfo)
abstractNewtypeQ (Q Info -> Q (Maybe DataInfo)) -> Q Info -> Q (Maybe DataInfo)
forall a b. (a -> b) -> a -> b
$ Name -> Q Info
reify Name
dt
let argNames :: Cxt
argNames = (TyVarBndr -> Type) -> [TyVarBndr] -> Cxt
forall a b. (a -> b) -> [a] -> [b]
map (Name -> Type
VarT (Name -> Type) -> (TyVarBndr -> Name) -> TyVarBndr -> Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TyVarBndr -> Name
tyVarBndrName) ([TyVarBndr] -> [TyVarBndr]
forall a. [a] -> [a]
tail [TyVarBndr]
args)
complType :: Type
complType = (Type -> Type -> Type) -> Type -> Cxt -> Type
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Type -> Type -> Type
AppT (Name -> Type
ConT Name
name) Cxt
argNames
preCond :: Cxt
preCond = (Type -> Type) -> Cxt -> Cxt
forall a b. (a -> b) -> [a] -> [b]
map (Name -> Cxt -> Type
mkClassP ''Arbitrary (Cxt -> Type) -> (Type -> Cxt) -> Type -> Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Type -> Cxt -> Cxt
forall a. a -> [a] -> [a]
: [])) Cxt
argNames
classType :: Type
classType = Type -> Type -> Type
AppT (Name -> Type
ConT ''ArbitraryF) Type
complType
Dec
arbitraryDecl <- [Con] -> Q Dec
generateArbitraryFDecl [Con]
constrs
Dec
shrinkDecl <- [Con] -> Q Dec
generateShrinkFDecl [Con]
constrs
[Dec] -> Q [Dec]
forall (m :: * -> *) a. Monad m => a -> m a
return [Cxt -> Type -> [Dec] -> Dec
mkInstanceD Cxt
preCond Type
classType [Dec
arbitraryDecl, Dec
shrinkDecl]]
generateArbitraryFDecl :: [Con] -> Q Dec
generateArbitraryFDecl :: [Con] -> Q Dec
generateArbitraryFDecl = Name -> [Con] -> Q Dec
generateGenDecl 'arbitraryF'
generateGenDecl :: Name -> [Con] -> Q Dec
generateGenDecl :: Name -> [Con] -> Q Dec
generateGenDecl Name
genName [Con]
constrs
= do Exp
genBody <- [ExpQ] -> ExpQ
listE ([ExpQ] -> ExpQ) -> [ExpQ] -> ExpQ
forall a b. (a -> b) -> a -> b
$ (Con -> ExpQ) -> [Con] -> [ExpQ]
forall a b. (a -> b) -> [a] -> [b]
map (ExpQ -> ExpQ
addNum (ExpQ -> ExpQ) -> (Con -> ExpQ) -> Con -> ExpQ
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, Int) -> ExpQ
constrGen ((Name, Int) -> ExpQ) -> (Con -> (Name, Int)) -> Con -> ExpQ
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Con -> (Name, Int)
abstractConType) [Con]
constrs
let genClause :: Clause
genClause = [Pat] -> Body -> [Dec] -> Clause
Clause [] (Exp -> Body
NormalB Exp
genBody) []
Dec -> Q Dec
forall (m :: * -> *) a. Monad m => a -> m a
return (Dec -> Q Dec) -> Dec -> Q Dec
forall a b. (a -> b) -> a -> b
$ Name -> [Clause] -> Dec
FunD Name
genName [Clause
genClause]
where addNum :: ExpQ -> ExpQ
addNum ExpQ
e = [| (1,$e) |]
constrGen :: (Name,Int) -> ExpQ
constrGen :: (Name, Int) -> ExpQ
constrGen (Name
constr, Int
n)
= do [Name]
varNs <- Int -> String -> Q [Name]
newNames Int
n String
"x"
Name
newSizeN <- String -> Q Name
newName String
"newSize"
let newSizeE :: ExpQ
newSizeE = Name -> ExpQ
varE Name
newSizeN
let newSizeP :: PatQ
newSizeP = Name -> PatQ
varP Name
newSizeN
let constrsE :: ExpQ
constrsE = Lit -> ExpQ
litE (Lit -> ExpQ) -> (Int -> Lit) -> Int -> ExpQ
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> Lit
IntegerL (Integer -> Lit) -> (Int -> Integer) -> Int -> Lit
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Int -> ExpQ) -> Int -> ExpQ
forall a b. (a -> b) -> a -> b
$ Int
n
let binds :: [StmtQ]
binds = ((Name -> StmtQ) -> [Name] -> [StmtQ]
forall a b. (a -> b) -> [a] -> [b]
`map` [Name]
varNs) (\Name
var -> PatQ -> ExpQ -> StmtQ
bindS
(Name -> PatQ
varP Name
var)
[| resize $newSizeE arbitrary |] )
let apps :: ExpQ
apps = [ExpQ] -> ExpQ
appsE (Name -> ExpQ
conE Name
constrExpQ -> [ExpQ] -> [ExpQ]
forall a. a -> [a] -> [a]
: (Name -> ExpQ) -> [Name] -> [ExpQ]
forall a b. (a -> b) -> [a] -> [b]
map Name -> ExpQ
varE [Name]
varNs)
let build :: ExpQ
build = [StmtQ] -> ExpQ
doE ([StmtQ] -> ExpQ) -> [StmtQ] -> ExpQ
forall a b. (a -> b) -> a -> b
$
[StmtQ]
binds [StmtQ] -> [StmtQ] -> [StmtQ]
forall a. [a] -> [a] -> [a]
++
[ExpQ -> StmtQ
noBindS [|return $apps|]]
if Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
then [|return $apps|]
else [| sized $ \ size ->
$(letE [valD
newSizeP
(normalB [|((size - 1) `div` $constrsE ) `max` 0|])
[] ]
build) |]
generateShrinkFDecl :: [Con] -> Q Dec
generateShrinkFDecl :: [Con] -> Q Dec
generateShrinkFDecl [Con]
constrs
= let clauses :: [Q Clause]
clauses = (Con -> Q Clause) -> [Con] -> [Q Clause]
forall a b. (a -> b) -> [a] -> [b]
map ((Name, Int) -> Q Clause
generateClause((Name, Int) -> Q Clause)
-> (Con -> (Name, Int)) -> Con -> Q Clause
forall b c a. (b -> c) -> (a -> b) -> a -> c
.Con -> (Name, Int)
abstractConType) [Con]
constrs
in Name -> [Q Clause] -> Q Dec
funD 'shrinkF [Q Clause]
clauses
where generateClause :: (Name, Int) -> Q Clause
generateClause (Name
constr, Int
n)
= do [Name]
varNs <- Int -> String -> Q [Name]
newNames Int
n String
"x"
[Name]
resVarNs <- Int -> String -> Q [Name]
newNames Int
n String
"x'"
[Stmt]
binds <- ((Name, Name) -> StmtQ) -> [(Name, Name)] -> Q [Stmt]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (\(Name
var,Name
resVar) -> PatQ -> ExpQ -> StmtQ
bindS (Name -> PatQ
varP Name
resVar) [| $(varE var) : shrink $(varE var) |]) ([(Name, Name)] -> Q [Stmt]) -> [(Name, Name)] -> Q [Stmt]
forall a b. (a -> b) -> a -> b
$ [Name] -> [Name] -> [(Name, Name)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Name]
varNs [Name]
resVarNs
let ret :: Stmt
ret = Exp -> Stmt
NoBindS (Exp -> Stmt) -> Exp -> Stmt
forall a b. (a -> b) -> a -> b
$ Exp -> Exp -> Exp
AppE (Name -> Exp
VarE 'return) ((Exp -> Exp -> Exp) -> [Exp] -> Exp
forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
foldl1 Exp -> Exp -> Exp
AppE ( Name -> Exp
ConE Name
constrExp -> [Exp] -> [Exp]
forall a. a -> [a] -> [a]
: (Name -> Exp) -> [Name] -> [Exp]
forall a b. (a -> b) -> [a] -> [b]
map Name -> Exp
VarE [Name]
resVarNs ))
stmtSeq :: [Stmt]
stmtSeq = [Stmt]
binds [Stmt] -> [Stmt] -> [Stmt]
forall a. [a] -> [a] -> [a]
++ [Stmt
ret]
pat :: Pat
pat = Name -> [Pat] -> Pat
ConP Name
constr ([Pat] -> Pat) -> [Pat] -> Pat
forall a b. (a -> b) -> a -> b
$ (Name -> Pat) -> [Name] -> [Pat]
forall a b. (a -> b) -> [a] -> [b]
map Name -> Pat
VarP [Name]
varNs
Clause -> Q Clause
forall (m :: * -> *) a. Monad m => a -> m a
return (Clause -> Q Clause) -> Clause -> Q Clause
forall a b. (a -> b) -> a -> b
$ [Pat] -> Body -> [Dec] -> Clause
Clause [Pat
pat] (Exp -> Body
NormalB (Exp -> Body) -> Exp -> Body
forall a b. (a -> b) -> a -> b
$ Exp -> Exp -> Exp
AppE (Name -> Exp
VarE 'tail) ([Stmt] -> Exp
DoE [Stmt]
stmtSeq)) []