{-# LANGUAGE GADTs           #-}
{-# LANGUAGE TemplateHaskell #-}
--------------------------------------------------------------------------------
-- |
-- Module      :  Data.Comp.Derive.Arbitrary
-- Copyright   :  (c) 2011 Patrick Bahr
-- License     :  BSD3
-- Maintainer  :  Patrick Bahr <paba@diku.dk>
-- Stability   :  experimental
-- Portability :  non-portable (GHC Extensions)
--
-- Automatically derive instances of @ArbitraryF@.
--
--------------------------------------------------------------------------------

module Data.Comp.Derive.Arbitrary
    (
     ArbitraryF(..),
     makeArbitraryF,
     Arbitrary(..)
    )where

import Data.Comp.Derive.Utils hiding (derive)
import Language.Haskell.TH
import Test.QuickCheck

{-| Signature arbitration. An instance @ArbitraryF f@ gives rise to an instance
  @Arbitrary (Term f)@. -}
class ArbitraryF f where
    arbitraryF' :: Arbitrary v => [(Int,Gen (f v))]
    arbitraryF' = [(Int
1,forall (f :: * -> *) v. (ArbitraryF f, Arbitrary v) => Gen (f v)
arbitraryF)]
    arbitraryF :: Arbitrary v => Gen (f v)
    arbitraryF = forall a. [(Int, Gen a)] -> Gen a
frequency forall (f :: * -> *) v.
(ArbitraryF f, Arbitrary v) =>
[(Int, Gen (f v))]
arbitraryF'
    shrinkF :: Arbitrary v => f v -> [f v]
    shrinkF f v
_ = []

{-| Derive an instance of 'ArbitraryF' for a type constructor of any
  first-order kind taking at least one argument. It is necessary that
  all types that are used by the data type definition are themselves
  instances of 'Arbitrary'. -}
makeArbitraryF :: Name -> Q [Dec]
makeArbitraryF :: Name -> Q [Dec]
makeArbitraryF Name
dt = do
  Just (DataInfo Cxt
_cxt Name
name [TyVarBndr flag]
args [Con]
constrs [DerivClause]
_deriving) <- Q Info -> Q (Maybe DataInfo)
abstractNewtypeQ forall a b. (a -> b) -> a -> b
$ Name -> Q Info
reify Name
dt
  let argNames :: Cxt
argNames = forall a b. (a -> b) -> [a] -> [b]
map (Name -> Type
VarT forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {flag}. TyVarBndr flag -> Name
tyVarBndrName) (forall a. [a] -> [a]
tail [TyVarBndr flag]
args)
      complType :: Type
complType = 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 = forall a b. (a -> b) -> [a] -> [b]
map (Name -> Cxt -> Type
mkClassP ''Arbitrary forall b c a. (b -> c) -> (a -> b) -> a -> c
. (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
  forall (m :: * -> *) a. Monad m => a -> m a
return [Cxt -> Type -> [Dec] -> Dec
mkInstanceD Cxt
preCond Type
classType [Dec
arbitraryDecl, Dec
shrinkDecl]]

{-|
  This function generates a declaration of the method 'arbitrary' for the given
  list of constructors using 'generateGenDecl'.
-}
generateArbitraryFDecl :: [Con] -> Q Dec
generateArbitraryFDecl :: [Con] -> Q Dec
generateArbitraryFDecl = Name -> [Con] -> Q Dec
generateGenDecl 'arbitraryF'

{-|
  This function generates a declaration of a generator having the given name using
  the given constructors, i.e., something like this:

  @
  \<name\> :: Gen \<type\>
  \<name\> = ...
  @

  where @\<type\>@ is the type of the given constructors. If the constructors do not belong
  to the same type this function fails. The generated function will generate only elements of
  this type using the given constructors. All argument types of these constructors are supposed
  to be instances of 'Arbitrary'.
-}

generateGenDecl :: Name -> [Con] -> Q Dec
generateGenDecl :: Name -> [Con] -> Q Dec
generateGenDecl Name
genName [Con]
constrs
    = do Exp
genBody <- forall (m :: * -> *). Quote m => [m Exp] -> m Exp
listE forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (forall {m :: * -> *}. Quote m => m Exp -> m Exp
addNum forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, Int) -> Q Exp
constrGen 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) []
         forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ Name -> [Clause] -> Dec
FunD Name
genName [Clause
genClause]
    where addNum :: m Exp -> m Exp
addNum m Exp
e = [| (1,$e) |]
          constrGen :: (Name,Int) -> ExpQ
          constrGen :: (Name, Int) -> Q Exp
constrGen (Name
constr, Int
n)
              = do [Name]
varNs <- Int -> String -> Q [Name]
newNames Int
n String
"x"
                   Name
newSizeN <- forall (m :: * -> *). Quote m => String -> m Name
newName String
"newSize"
                   let newSizeE :: Q Exp
newSizeE = forall (m :: * -> *). Quote m => Name -> m Exp
varE Name
newSizeN
                   let newSizeP :: Q Pat
newSizeP = forall (m :: * -> *). Quote m => Name -> m Pat
varP Name
newSizeN
                   let constrsE :: Q Exp
constrsE = forall (m :: * -> *). Quote m => Lit -> m Exp
litE forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> Lit
IntegerL forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Integral a => a -> Integer
toInteger forall a b. (a -> b) -> a -> b
$ Int
n
                   let binds :: [Q Stmt]
binds = (forall a b. (a -> b) -> [a] -> [b]
`map` [Name]
varNs) (\Name
var -> forall (m :: * -> *). Quote m => m Pat -> m Exp -> m Stmt
bindS
                                                     (forall (m :: * -> *). Quote m => Name -> m Pat
varP Name
var)
                                                     [| resize $newSizeE arbitrary |] )
                   let apps :: Q Exp
apps =  forall (m :: * -> *). Quote m => [m Exp] -> m Exp
appsE (forall (m :: * -> *). Quote m => Name -> m Exp
conE Name
constrforall a. a -> [a] -> [a]
: forall a b. (a -> b) -> [a] -> [b]
map forall (m :: * -> *). Quote m => Name -> m Exp
varE [Name]
varNs)
                   let build :: Q Exp
build = forall (m :: * -> *). Quote m => [m Stmt] -> m Exp
doE forall a b. (a -> b) -> a -> b
$
                               [Q Stmt]
binds forall a. [a] -> [a] -> [a]
++
                               [forall (m :: * -> *). Quote m => m Exp -> m Stmt
noBindS [|return $apps|]]
                   if Int
n 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) |]

{-|
  This function generates a declaration for the method 'shrinkF' using the given constructors.
  The constructors are supposed to belong to the same type.
-}
generateShrinkFDecl :: [Con] -> Q Dec
generateShrinkFDecl :: [Con] -> Q Dec
generateShrinkFDecl [Con]
constrs
    = let clauses :: [Q Clause]
clauses = forall a b. (a -> b) -> [a] -> [b]
map ((Name, Int) -> Q Clause
generateClauseforall b c a. (b -> c) -> (a -> b) -> a -> c
.Con -> (Name, Int)
abstractConType) [Con]
constrs
      in forall (m :: * -> *). Quote m => Name -> [m Clause] -> m 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 <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (\(Name
var,Name
resVar) -> forall (m :: * -> *). Quote m => m Pat -> m Exp -> m Stmt
bindS (forall (m :: * -> *). Quote m => Name -> m Pat
varP Name
resVar) [| $(varE var) : shrink $(varE var) |]) forall a b. (a -> b) -> a -> b
$ forall a b. [a] -> [b] -> [(a, b)]
zip [Name]
varNs [Name]
resVarNs
                 let ret :: Stmt
ret = Exp -> Stmt
NoBindS forall a b. (a -> b) -> a -> b
$ Exp -> Exp -> Exp
AppE (Name -> Exp
VarE 'return) (forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
foldl1 Exp -> Exp -> Exp
AppE ( Name -> Exp
ConE Name
constrforall a. a -> [a] -> [a]
: forall a b. (a -> b) -> [a] -> [b]
map Name -> Exp
VarE [Name]
resVarNs ))
                     stmtSeq :: [Stmt]
stmtSeq = [Stmt]
binds forall a. [a] -> [a] -> [a]
++ [Stmt
ret]
                     pat :: Pat
pat = Name -> Cxt -> [Pat] -> Pat
ConP Name
constr [] forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map Name -> Pat
VarP [Name]
varNs
                 forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ [Pat] -> Body -> [Dec] -> Clause
Clause [Pat
pat] (Exp -> Body
NormalB forall a b. (a -> b) -> a -> b
$ Exp -> Exp -> Exp
AppE (Name -> Exp
VarE 'tail) (Maybe ModName -> [Stmt] -> Exp
DoE forall a. Maybe a
Nothing [Stmt]
stmtSeq)) []