{-# LANGUAGE TemplateHaskell #-}

module Test.QuickCheck.TH.Generators.Internal.BuildArbitrary where

import           Language.Haskell.TH

import           Safe


-- | automatically build functions named `buildArbn` where n is an integer
-- greater than 0. $(buildArbAny 3) creates a function buildArb3 which takes
-- a constructor that takes 3 parameters, and returns an arbitrary instances
-- of that constructor. It assumes that the constructors type is an instance
-- of Arbitrary.
--
-- buildArb1 :: Arbitrary a
--           => (a -> b)
--           -> Gen b
-- buildArb1 f = f <$> arbitrary
--
-- buildArb2 :: (Arbitrary a, Arbitrary a1)
--           => (a1 -> a -> b)
--           -> Gen b
-- buildArb2 f = f <$> arbitrary <*> arbitrary

buildArbAny :: Int -> Q [Dec]
buildArbAny :: Int -> Q [Dec]
buildArbAny Int
l
  | Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
1     = [Dec] -> Q [Dec]
forall (m :: * -> *) a. Monad m => a -> m a
return []
  | Bool
otherwise = do
    -- name of the function to be returned
    let buildArbName :: Name
buildArbName = String -> Name
mkName (String
"buildArb" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
l)

    -- names of parameters used in the buildArbn function
    [Name]
arbParameterNames <- (String -> Q Name) -> [String] -> Q [Name]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM String -> Q Name
newName ([String] -> Q [Name]) -> [String] -> Q [Name]
forall a b. (a -> b) -> a -> b
$ ((String
"a" String -> String -> String
forall a. [a] -> [a] -> [a]
++) (String -> String) -> (Int -> String) -> Int -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> String
forall a. Show a => a -> String
show) (Int -> String) -> [Int] -> [String]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Int
0..(Int
lInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)] -- one for each parameter of f
    let mFirstArbParameterName :: Maybe Name
mFirstArbParameterName = [Name] -> Maybe Name
forall a. [a] -> Maybe a
headMay [Name]
arbParameterNames
        mRestArbParameterNames :: Maybe [Name]
mRestArbParameterNames = [Name] -> Maybe [Name]
forall a. [a] -> Maybe [a]
tailMay [Name]
arbParameterNames
    Name
bName <- String -> Q Name
newName String
"b"  -- the return type of buildArbn
    Name
fName <- String -> Q Name
newName String
"f"  -- the sum type constructor

    -- types used in the buildArbn function
    Maybe Name
mArbitraryTypeName <- String -> Q (Maybe Name)
lookupTypeName String
"Arbitrary"
    Maybe Name
mGenTypeName       <- String -> Q (Maybe Name)
lookupTypeName String
"Gen"

    -- functions used in the buildArbn function
    Maybe Name
mFmapName        <- String -> Q (Maybe Name)
lookupValueName String
"<$>"
    Maybe Name
mApName          <- String -> Q (Maybe Name)
lookupValueName String
"<*>"
    Maybe Name
mArbitraryValue  <- String -> Q (Maybe Name)
lookupValueName String
"arbitrary"

    -- check if all the types and functions were found
    case (,,,,,,) (Name
 -> [Name]
 -> Name
 -> Name
 -> Name
 -> Name
 -> Name
 -> (Name, [Name], Name, Name, Name, Name, Name))
-> Maybe Name
-> Maybe
     ([Name]
      -> Name
      -> Name
      -> Name
      -> Name
      -> Name
      -> (Name, [Name], Name, Name, Name, Name, Name))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe Name
mFirstArbParameterName Maybe
  ([Name]
   -> Name
   -> Name
   -> Name
   -> Name
   -> Name
   -> (Name, [Name], Name, Name, Name, Name, Name))
-> Maybe [Name]
-> Maybe
     (Name
      -> Name
      -> Name
      -> Name
      -> Name
      -> (Name, [Name], Name, Name, Name, Name, Name))
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Maybe [Name]
mRestArbParameterNames Maybe
  (Name
   -> Name
   -> Name
   -> Name
   -> Name
   -> (Name, [Name], Name, Name, Name, Name, Name))
-> Maybe Name
-> Maybe
     (Name
      -> Name
      -> Name
      -> Name
      -> (Name, [Name], Name, Name, Name, Name, Name))
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Maybe Name
mArbitraryTypeName Maybe
  (Name
   -> Name
   -> Name
   -> Name
   -> (Name, [Name], Name, Name, Name, Name, Name))
-> Maybe Name
-> Maybe
     (Name
      -> Name -> Name -> (Name, [Name], Name, Name, Name, Name, Name))
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Maybe Name
mGenTypeName Maybe
  (Name
   -> Name -> Name -> (Name, [Name], Name, Name, Name, Name, Name))
-> Maybe Name
-> Maybe
     (Name -> Name -> (Name, [Name], Name, Name, Name, Name, Name))
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Maybe Name
mFmapName Maybe
  (Name -> Name -> (Name, [Name], Name, Name, Name, Name, Name))
-> Maybe Name
-> Maybe (Name -> (Name, [Name], Name, Name, Name, Name, Name))
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Maybe Name
mApName Maybe (Name -> (Name, [Name], Name, Name, Name, Name, Name))
-> Maybe Name -> Maybe (Name, [Name], Name, Name, Name, Name, Name)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Maybe Name
mArbitraryValue of
      Maybe (Name, [Name], Name, Name, Name, Name, Name)
Nothing -> [Dec] -> Q [Dec]
forall (m :: * -> *) a. Monad m => a -> m a
return []
      Just (Name
firstArbParameterName,[Name]
restArbParameterNames,Name
arbitraryTypeName,Name
genTypeName,Name
fmapName,Name
apName,Name
arbitraryValue) -> do
        -- all of the variables in the function to be created, input and output
        let plainTVs :: [TyVarBndr]
plainTVs = Name -> TyVarBndr
PlainTV (Name -> TyVarBndr) -> [Name] -> [TyVarBndr]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ([Name]
arbParameterNames [Name] -> [Name] -> [Name]
forall a. [a] -> [a] -> [a]
++ [Name
bName])
            -- Arbitrary type instance required for all vars in arbParameterNames
            typeClassRequirements :: [Type]
typeClassRequirements = (Type -> Type -> Type
AppT (Name -> Type
ConT Name
arbitraryTypeName) ) (Type -> Type) -> [Type] -> [Type]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Name -> Type
VarT (Name -> Type) -> [Name] -> [Type]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Name]
arbParameterNames)
            genB :: Type
genB     = Type -> Type -> Type
AppT (Name -> Type
ConT Name
genTypeName) (Name -> Type
VarT Name
bName)

            -- last parameter and return type
            -- (a -> Gen b)
            aToGenB :: Type
aToGenB = (Type -> Type -> Type
AppT
                        (Type -> Type -> Type
AppT
                          Type
ArrowT
                          (Name -> Type
VarT Name
firstArbParameterName))
                        (Name -> Type
VarT Name
bName))
            -- fold function
            buildFunctionArgument :: Type -> Name -> Type
buildFunctionArgument Type
old Name
new = (Type -> Type -> Type
AppT
                                              (Type -> Type -> Type
AppT
                                                Type
ArrowT
                                                (Name -> Type
VarT Name
new))
                                              Type
old)
            -- build the rest of the parameters
            preFunctionArgument :: Type
preFunctionArgument = Type -> Type -> Type
AppT Type
ArrowT (Type -> Type) -> Type -> Type
forall a b. (a -> b) -> a -> b
$ (Type -> Name -> Type) -> Type -> [Name] -> Type
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Type -> Name -> Type
buildFunctionArgument
                                                      Type
aToGenB
                                                      [Name]
restArbParameterNames
            -- TH encoding for function arguments
            functionArgument :: Type
functionArgument = Type -> Type -> Type
AppT Type
preFunctionArgument Type
genB

            -- build the function body
            -- f <$> arbitrary
            fFmapArbitrary :: Exp
fFmapArbitrary = (Maybe Exp -> Exp -> Maybe Exp -> Exp
InfixE (Exp -> Maybe Exp
forall a. a -> Maybe a
Just (Name -> Exp
VarE Name
fName)) (Name -> Exp
VarE Name
fmapName) (Exp -> Maybe Exp
forall a. a -> Maybe a
Just (Name -> Exp
VarE Name
arbitraryValue)))
            arbRs :: [Name]
arbRs = Int -> Name -> [Name]
forall a. Int -> a -> [a]
replicate ([Name] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Name]
arbParameterNames Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Name
arbitraryValue

            -- build rest of body by folding <*> arbitrary
            preFunctionBody :: Exp
preFunctionBody = (Exp -> Name -> Exp) -> Exp -> [Name] -> Exp
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (\Exp
old Name
new -> Maybe Exp -> Exp -> Maybe Exp -> Exp
InfixE (Exp -> Maybe Exp
forall a. a -> Maybe a
Just Exp
old) (Name -> Exp
VarE Name
apName) (Exp -> Maybe Exp
forall a. a -> Maybe a
Just (Name -> Exp
VarE Name
new)))
                                    Exp
fFmapArbitrary
                                    [Name]
arbRs

            -- TH encoding for function body
            functionBody :: Dec
functionBody = Name -> [Clause] -> Dec
FunD Name
buildArbName [[Pat] -> Body -> [Dec] -> Clause
Clause [Name -> Pat
VarP Name
fName] (Exp -> Body
NormalB Exp
preFunctionBody) []]

        [Dec] -> Q [Dec]
forall (m :: * -> *) a. Monad m => a -> m a
return
          [Name -> Type -> Dec
SigD
            Name
buildArbName
            ([TyVarBndr] -> [Type] -> Type -> Type
ForallT
              [TyVarBndr]
plainTVs
              [Type]
typeClassRequirements
              Type
functionArgument
            )
          , Dec
functionBody
          ]