{-# 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,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
_ = []

{-| 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]
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]]

{-|
  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 <- [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) |]

{-|
  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 = (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)) []