{-# LANGUAGE TemplateHaskell #-}

-- | Derive an 'Arbitrary' instance for a given data type.
module Test.Mutagen.TH.Arbitrary
  ( deriveArbitrary
  )
where

import Control.Monad.Extra (ifM)
import Data.List (partition)
import Language.Haskell.TH
  ( Lit (..)
  , Name
  , Q
  , Type (..)
  , isInstance
  , newName
  , varE
  )
import Language.Haskell.TH.Desugar
  ( DClause (..)
  , DCon (..)
  , DDec (..)
  , DExp (..)
  , DLetDec (..)
  , DMatch (..)
  , DPat (..)
  , DType (..)
  , dCaseE
  , dsExp
  )
import Test.Mutagen
  ( Arbitrary
  , Arbitrary1
  , Arbitrary2
  , Gen
  , arbitrary
  , liftArbitrary
  , liftArbitrary2
  , oneof
  , sized
  )
import Test.Mutagen.TH.Util
  ( applyTyVars
  , dConFields
  , dConFieldsTypes
  , dConName
  , dTyVarBndrName
  , isMaybeOf
  , mkApplicativeDExp
  , mkListDExp
  , reifyTypeDef
  , (.==.)
  )

{-------------------------------------------------------------------------------
-- * Deriving Arbitrary instances
-------------------------------------------------------------------------------}

-- | Derive an 'Arbitrary' instance for the given data type.
deriveArbitrary :: Name -> [Name] -> Q [DDec]
deriveArbitrary :: Name -> [Name] -> Q [DDec]
deriveArbitrary Name
typeName [Name]
ignoredCons = do
  -- Reify the type definition
  (dtvbs, dcons) <- Name -> Q ([DTyVarBndrVis], [DCon])
reifyTypeDef Name
typeName
  -- Apply the context type variables to the type name to get 'Type'-kinded
  -- target type to derive the instance for
  let targetType = Name -> [DTyVarBndrVis] -> DType
applyTyVars Name
typeName [DTyVarBndrVis]
dtvbs
  -- Keep only the constructors that are not ignored
  let wantedCons = (DCon -> Bool) -> [DCon] -> [DCon]
forall a. (a -> Bool) -> [a] -> [a]
filter (\DCon
con -> DCon -> Name
dConName DCon
con Name -> [Name] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` [Name]
ignoredCons) [DCon]
dcons
  -- Split terminal and recursive constructors
  let (recs, terms) =
        partition
          (\DCon
con -> DType
targetType DType -> [DType] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` DConFields -> [DType]
dConFieldsTypes (DCon -> DConFields
dConFields DCon
con))
          wantedCons
  -- Create some fresh TH variables
  gen_ <- newName "gen"
  size_ <- newName "size"
  -- Create generator expressions for each constructor
  recGens <- mapM (mkConGen targetType gen_ size_) recs
  termGens <- mapM (mkConGen targetType gen_ size_) terms
  -- Build the Arbitrary instance
  let insCxt = [Name -> DType
DConT ''Arbitrary DType -> DType -> DType
`DAppT` Name -> DType
DVarT (DTyVarBndrVis -> Name
dTyVarBndrName DTyVarBndrVis
tvb) | DTyVarBndrVis
tvb <- [DTyVarBndrVis]
dtvbs]
  let insTy = Name -> DType
DConT ''Arbitrary DType -> DType -> DType
`DAppT` DType
targetType
  let baseCase =
        DPat -> DExp -> DMatch
DMatch
          (Lit -> DPat
DLitP (Integer -> Lit
IntegerL Integer
0))
          (Name -> DExp
DVarE 'oneof DExp -> DExp -> DExp
`DAppE` [DExp] -> DExp
mkListDExp [DExp]
termGens)
  let recCase =
        DPat -> DExp -> DMatch
DMatch
          DPat
DWildP
          (Name -> DExp
DVarE 'oneof DExp -> DExp -> DExp
`DAppE` [DExp] -> DExp
mkListDExp ([DExp]
termGens [DExp] -> [DExp] -> [DExp]
forall a. Semigroup a => a -> a -> a
<> [DExp]
recGens))
  let genDecl =
        Name -> [DClause] -> DLetDec
DFunD
          Name
gen_
          [ [DPat] -> DExp -> DClause
DClause
              [Name -> DPat
DVarP Name
size_]
              ( DExp -> [DMatch] -> DExp
dCaseE
                  (Name -> DExp
DVarE Name
size_)
                  [DMatch
baseCase, DMatch
recCase]
              )
          ]
  let insGen = [DLetDec] -> DExp -> DExp
DLetE [DLetDec
genDecl] (Name -> DExp
DVarE 'sized DExp -> DExp -> DExp
`DAppE` Name -> DExp
DVarE Name
gen_)
  let insBody = [DLetDec -> DDec
DLetDec (Name -> [DClause] -> DLetDec
DFunD 'arbitrary [[DPat] -> DExp -> DClause
DClause [] DExp
insGen])]
  return [DInstanceD Nothing Nothing insCxt insTy insBody]

-- | Create the appropriate generator for a constructor.
mkConGen :: DType -> Name -> Name -> DCon -> Q DExp
mkConGen :: DType -> Name -> Name -> DCon -> Q DExp
mkConGen DType
targetType Name
gen_ Name
size_ (DCon [DTyVarBndrSpec]
_ [DType]
_ Name
cname DConFields
cfields DType
_) = do
  fieldGens <-
    (DType -> Q DExp) -> [DType] -> Q [DExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM
      (DType -> Name -> Name -> DType -> Q DExp
mkConFieldGen DType
targetType Name
gen_ Name
size_)
      (DConFields -> [DType]
dConFieldsTypes DConFields
cfields)
  return (mkApplicativeDExp cname fieldGens)

-- | Create the appropriate generator for a constructor field based on its type.
mkConFieldGen :: DType -> Name -> Name -> DType -> Q DExp
mkConFieldGen :: DType -> Name -> Name -> DType -> Q DExp
mkConFieldGen DType
targetType Name
gen_ Name
size_ DType
fieldTy =
  Exp -> Q DExp
forall (q :: * -> *). DsMonad q => Exp -> q DExp
dsExp (Exp -> Q DExp) -> Q Exp -> Q DExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< DType -> Q Exp
mkGen DType
fieldTy
  where
    mkGen :: DType -> Q Exp
mkGen DType
ty
      -- The field is (immediately) recursive
      -- ==> Use the reference to the recursive generator
      | DType
ty DType -> DType -> Bool
.==. DType
targetType =
          [e|$(Name -> Q Exp
forall (m :: * -> *). Quote m => Name -> m Exp
varE Name
gen_) (max 0 ($(Name -> Q Exp
forall (m :: * -> *). Quote m => Name -> m Exp
varE Name
size_) - 1))|]
      -- Some special cases in between
      -- ==> use special 'Gen' combinators accordingly
      | Just DType
a <- DType -> Maybe DType
isMaybeOf DType
ty =
          [e|sizedMaybe $(DType -> Q Exp
mkGen DType
a)|]
      -- Types of kind `Type -> Type`
      -- ==> Check for 'Arbitrary1' instance, otherwise fallback to 'arbitrary'
      | (DConT Name
f `DAppT` DType
a) <- DType
ty =
          Q Bool -> Q Exp -> Q Exp -> Q Exp
forall (m :: * -> *) a. Monad m => m Bool -> m a -> m a -> m a
ifM
            (Name -> [Type] -> Q Bool
isInstance ''Arbitrary1 [Name -> Type
ConT Name
f])
            [e|liftArbitrary $(DType -> Q Exp
mkGen DType
a)|]
            [e|arbitrary|]
      -- Types of kind `Type -> Type -> Type`
      -- ==> Check for 'Arbitrary2' instance, otherwise fallback to 'arbitrary'
      | (DConT Name
f `DAppT` DType
t1 `DAppT` DType
t2) <- DType
ty =
          Q Bool -> Q Exp -> Q Exp -> Q Exp
forall (m :: * -> *) a. Monad m => m Bool -> m a -> m a -> m a
ifM
            (Name -> [Type] -> Q Bool
isInstance ''Arbitrary2 [Name -> Type
ConT Name
f])
            [e|liftArbitrary2 $(DType -> Q Exp
mkGen DType
t1) $(DType -> Q Exp
mkGen DType
t2)|]
            [e|arbitrary|]
      -- The field type is something else
      -- ==> Assume it has an 'Arbitrary' instance and hope for the best
      | Bool
otherwise =
          [e|arbitrary|]

-- | Create a sized 'Maybe' generator that returns 'Nothing' when size is 0.
--
-- NOTE: we found this useful empirically for types where it's really hard to
-- randomly generate valid random values for. Defaulting to 'Nothing' at size 0
-- helps avoid generating invalid optionals.
sizedMaybe :: Gen a -> Gen (Maybe a)
sizedMaybe :: forall a. Gen a -> Gen (Maybe a)
sizedMaybe Gen a
gen = (Int -> Gen (Maybe a)) -> Gen (Maybe a)
forall a. (Int -> Gen a) -> Gen a
sized ((Int -> Gen (Maybe a)) -> Gen (Maybe a))
-> (Int -> Gen (Maybe a)) -> Gen (Maybe a)
forall a b. (a -> b) -> a -> b
$ \Int
size ->
  if Int
size Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
    then Maybe a -> Gen (Maybe a)
forall a. a -> Gen a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe a
forall a. Maybe a
Nothing
    else Gen a -> Gen (Maybe a)
forall a. Gen a -> Gen (Maybe a)
forall (f :: * -> *) a. Arbitrary1 f => Gen a -> Gen (f a)
liftArbitrary Gen a
gen