content: {-# LANGUAGE TemplateHaskell, CPP #-}
  -- |
  -- Module      : Test.LeanCheck.Derive
  -- Copyright   : (c) 2015-2025 Rudy Matela
  -- License     : 3-Clause BSD  (see the file LICENSE)
  -- Maintainer  : Rudy Matela <rudy@matela.com.br>
  --
  -- This module is part of LeanCheck,
  -- a simple enumerative property-based testing library.
  --
  -- Needs GHC and Template Haskell
  -- (tested on GHC 7.4, 7.6, 7.8, 7.10, 8.0, 8.2, 8.4, 8.6 and 8.8).
  --
  -- If LeanCheck does not compile under later GHCs, this module is probably the
  -- culprit.
  --
  -- If you rather do this through GHC Generics, please see:
  -- "Test.LeanCheck.Generic" (experimental).
  module Test.LeanCheck.Derive
    ( deriveListable
    , deriveListableIfNeeded
    , deriveListableCascading
    , deriveTiers
    , deriveList
    )
  where
  
  #ifdef __GLASGOW_HASKELL__
  
  import Language.Haskell.TH
  import Test.LeanCheck.Basic
  import Control.Monad (unless, filterM)
  import Data.List (delete)
  
  #if __GLASGOW_HASKELL__ < 706
  -- reportWarning was only introduced in GHC 7.6 / TH 2.8
  reportWarning :: String -> Q ()
  reportWarning  =  report False
  #endif
  
  -- | Derives a 'Listable' instance for a given type 'Name'.
  --
  -- Consider the following @Stack@ datatype:
  --
  -- > data Stack a  =  Stack a (Stack a) | Empty
  --
  -- Writing
  --
  -- > deriveListable ''Stack
  --
  -- will automatically derive the following 'Listable' instance:
  --
  -- > instance Listable a => Listable (Stack a) where
  -- >   tiers  =  cons2 Stack \/ cons0 Empty
  --
  -- __Warning:__ if the values in your type need to follow a data invariant, the
  --              derived instance won't respect it.  Use this only on "free"
  --              datatypes.
  --
  -- Needs the @TemplateHaskell@ extension.
  deriveListable :: Name -> DecsQ
  deriveListable  =  deriveListableX True False
  
  -- | Same as 'deriveListable' but does not warn when the requested instance
  --   already exists.  The function 'deriveListable' is preferable in most
  --   situations.
  deriveListableIfNeeded :: Name -> DecsQ
  deriveListableIfNeeded  =  deriveListableX False False
  
  -- | Derives a 'Listable' instance for a given type 'Name'
  --   cascading derivation of type arguments as well.
  --
  -- Consider the following series of datatypes:
  --
  -- > data Position  =  CEO | Manager | Programmer
  -- >
  -- > data Person  =  Person
  -- >              {  name :: String
  -- >              ,  age :: Int
  -- >              ,  position :: Position
  -- >              }
  -- >
  -- > data Company  =  Company
  -- >               {  name :: String
  -- >               ,  employees :: [Person]
  -- >               }
  --
  -- Writing
  --
  -- > deriveListableCascading ''Company
  --
  -- will automatically derive the following three 'Listable' instances:
  --
  -- > instance Listable Position where
  -- >   tiers  =  cons0 CEO \/ cons0 Manager \/ cons0 Programmer
  -- >
  -- > instance Listable Person where
  -- >   tiers  =  cons3 Person
  -- >
  -- > instance Listable Company where
  -- >   tiers  =  cons2 Company
  deriveListableCascading :: Name -> DecsQ
  deriveListableCascading  =  deriveListableX True True
  
  deriveListableX :: Bool -> Bool -> Name -> DecsQ
  deriveListableX warnExisting cascade t  =  do
    is <- t `isInstanceOf` ''Listable
    if is
    then do unless (not warnExisting) $
              reportWarning $ "Instance Listable " ++ show t
                           ++ " already exists, skipping derivation"
            return []
    else if cascade
         then reallyDeriveListableCascading t
         else reallyDeriveListable t
  
  reallyDeriveListable :: Name -> DecsQ
  reallyDeriveListable t  =  do
    (nt,vs) <- normalizeType t
  #if __GLASGOW_HASKELL__ >= 710
    cxt <- sequence [[t| Listable $(return v) |] | v <- vs]
  #else
    cxt <- sequence [classP ''Listable [return v] | v <- vs]
  #endif
  #if __GLASGOW_HASKELL__ >= 708
    cxt |=>| [d| instance Listable $(return nt)
                   where tiers  =  $(deriveTiers t) |]
  #else
    tiersE <- deriveTiers t
    return [ InstanceD
               cxt
               (AppT (ConT ''Listable) nt)
               [ValD (VarP 'tiers) (NormalB tiersE) []]
           ]
  #endif
  
  -- | Given a type 'Name', derives an expression to be placed as the result of
  --   'tiers':
  --
  -- > consN C1 \/ consN C2 \/ ... \/ consN CN
  --
  -- This function can be used in the definition of 'Listable' instances:
  --
  -- > instance Listable MyType where
  -- >   tiers  =  $(deriveTiers)
  deriveTiers :: Name -> ExpQ
  deriveTiers t  =  conse =<< typeConstructorsWithArityAndRecursiveness t
    where
    cone (n,arity,shouldReset) =  do
      (Just consN) <- lookupValueName $ "cons" ++ show arity
      -- the shouldReset switch is necessary to avoid derivations
      -- where we end with an empty tier at the head of the tiers list
      -- such as a type homomorphic to Either
      -- on earlier LeanCheck versions, we never reset
      if shouldReset
      then [| reset ($(varE consN) $(conE n)) |]
      else [|        $(varE consN) $(conE n)  |]
    conse  =  foldr1 (\e1 e2 -> [| $e1 \/ $e2 |]) . map cone . arityReset
    -- computes whether we should reset some constructors and which
    arityReset nars
      -- if there's a constructor with 0 arguments, we don't need to reset
      | or [arity == 0 | (_,arity,_) <- nars]  =  [(name,arity,False) | (name,arity,_) <- nars]
      -- otherwise we reset constructors that are not recursive to avoid infinite loops
      | otherwise  =  [(name, arity, arity > 0 && not isRecursive) | (name,arity,isRecursive) <- nars]
      -- by reset here of course we mean removing the preceding empty tiers
  
  -- | Given a type 'Name', derives an expression to be placed as the result of
  --   'list':
  --
  -- > concat $ consN C1 \/ consN C2 \/ ... \/ consN CN
  deriveList :: Name -> ExpQ
  deriveList t  =  [| concat $(deriveTiers t) |]
  
  -- Not only really derive Listable instances,
  -- but cascade through argument types.
  reallyDeriveListableCascading :: Name -> DecsQ
  reallyDeriveListableCascading t =
        return . concat
    =<< mapM reallyDeriveListable
    =<< filterM (fmap not . isTypeSynonym)
    =<< return . (t:) . delete t
    =<< t `typeConCascadingArgsThat` (`isntInstanceOf` ''Listable)
  
  -- * Template haskell utilities
  
  typeConArgs :: Name -> Q [Name]
  typeConArgs t  =  do
    is <- isTypeSynonym t
    if is
    then subtypeNames `fmap` typeSynonymType t
    else (nubMerges . map subtypeNames . concatMap snd) `fmap` typeConstructors t
  
  subtypeNames :: Type -> [Name]
  subtypeNames (AppT t1 t2)  =  subtypeNames t1 `nubMerge` subtypeNames t2
  subtypeNames (SigT t _)  =  subtypeNames t
  subtypeNames (VarT _)  =  []
  subtypeNames (ConT n)  =  [n]
  #if __GLASGOW_HASKELL__ >= 800
  -- subtypeNames (PromotedT n)  =  [n] ?
  subtypeNames (InfixT  t1 n t2)  =  subtypeNames t1 `nubMerge` subtypeNames t2
  subtypeNames (UInfixT t1 n t2)  =  subtypeNames t1 `nubMerge` subtypeNames t2
  subtypeNames (ParensT t)  =  subtypeNames t
  #endif
  subtypeNames _  =  []
  
  cascadingSubtypeNames :: [Type] -> Q [Name]
  cascadingSubtypeNames t  =  do
    nss <- mapM (`typeConCascadingArgsThat` (\t -> return $ t `notElem` ns)) ns
    return $ nubMerges (ns:nss)
    where
    ns  =  nubMerges $ map subtypeNames t
  
  typeConArgsThat :: Name -> (Name -> Q Bool) -> Q [Name]
  t `typeConArgsThat` p  =  filterM p =<< typeConArgs t
  
  typeConCascadingArgsThat :: Name -> (Name -> Q Bool) -> Q [Name]
  t `typeConCascadingArgsThat` p  =  do
    ts <- t `typeConArgsThat` p
    let p' t'  =  (t' `notElem` t:ts &&) `fmap` p t'
    tss <- mapM (`typeConCascadingArgsThat` p') ts
    return $ nubMerges (ts:tss)
  
  -- |
  -- Normalizes a type by applying it to necessary type variables
  -- making it accept zero type parameters.
  -- The normalized type is paired with a list of necessary type variables.
  --
  -- > > putStrLn $(stringE . show =<< normalizeType ''Int)
  -- > (ConT ''Int, [])
  --
  -- > > putStrLn $(stringE . show =<< normalizeType ''Maybe)
  -- > (AppT (ConT ''Maybe) (VarT ''a),[VarT ''a])
  --
  -- > > putStrLn $(stringE . show =<< normalizeType ''Either)
  -- > (AppT (AppT (ConT ''Either) (VarT ''a)) (VarT ''b),[VarT ''a,VarT ''b])
  --
  -- > > putStrLn $(stringE . show =<< normalizeType ''[])
  -- > (AppT (ConT ''[]) (VarT a),[VarT a])
  normalizeType :: Name -> Q (Type, [Type])
  normalizeType t  =  do
    ar <- typeArity t
    vs <- newVarTs ar
    return (foldl AppT (ConT t) vs, vs)
    where
      newNames :: [String] -> Q [Name]
      newNames  =  mapM newName
      newVarTs :: Int -> Q [Type]
      newVarTs n  =  map VarT
              `fmap` newNames (take n . map (:[]) $ cycle ['a'..'z'])
  
  -- |
  -- Normalizes a type by applying it to units to make it star-kinded.
  -- (cf. 'normalizeType')
  --
  -- > normalizeTypeUnits ''Int    === [t| Int |]
  -- > normalizeTypeUnits ''Maybe  === [t| Maybe () |]
  -- > normalizeTypeUnits ''Either === [t| Either () () |]
  normalizeTypeUnits :: Name -> Q Type
  normalizeTypeUnits t  =  do
    ar <- typeArity t
    return (foldl AppT (ConT t) (replicate ar (TupleT 0)))
  
  -- |
  -- Given a type name and a class name,
  -- returns whether the type is an instance of that class.
  -- The given type must be star-kinded (@ * @)
  -- and the given class double-star-kinded (@ * -> * @.
  --
  -- > > putStrLn $(stringE . show =<< ''Int `isInstanceOf` ''Num)
  -- > True
  --
  -- > > putStrLn $(stringE . show =<< ''Int `isInstanceOf` ''Fractional)
  -- > False
  isInstanceOf :: Name -> Name -> Q Bool
  isInstanceOf tn cl  =  do
    ty <- normalizeTypeUnits tn
    isInstance cl [ty]
  
  -- |
  -- The negation of 'isInstanceOf'.
  isntInstanceOf :: Name -> Name -> Q Bool
  isntInstanceOf tn  =  fmap not . isInstanceOf tn
  
  -- | Given a type name, return the number of arguments taken by that type.
  -- Examples in partially broken TH:
  --
  -- > > putStrLn $(stringE . show =<< typeArity ''Int)
  -- > 0
  --
  -- > > putStrLn $(stringE . show =<< typeArity ''Maybe)
  -- > 1
  --
  -- > > putStrLn $(stringE . show =<< typeArity ''Either)
  -- > 2
  --
  -- > > putStrLn $(stringE . show =<< typeArity ''[])
  -- > 1
  --
  -- > > putStrLn $(stringE . show =<< typeArity ''(,))
  -- > 2
  --
  -- > > putStrLn $(stringE . show =<< typeArity ''(,,))
  -- > 3
  --
  -- > > putStrLn $(stringE . show =<< typeArity ''String)
  -- > 0
  --
  -- This works for data and newtype declarations and
  -- it is useful when generating typeclass instances.
  typeArity :: Name -> Q Int
  typeArity t  =  fmap arity $ reify t
    where
    arity  =  length . args
  #if __GLASGOW_HASKELL__ < 800
    args (TyConI (DataD    _ _ ks   _ _))  =  ks
    args (TyConI (NewtypeD _ _ ks   _ _))  =  ks
  #else
    args (TyConI (DataD    _ _ ks _ _ _))  =  ks
    args (TyConI (NewtypeD _ _ ks _ _ _))  =  ks
  #endif
    args (TyConI (TySynD _ ks _))          =  ks
    args _  =  errorOn "typeArity"
            $  "neither newtype nor data nor type synonym: " ++ show t
  
  -- |
  -- Given a type 'Name',
  -- returns a list of its type constructor 'Name's
  -- paired with the type arguments they take.
  -- the type arguments they take.
  --
  -- > > :set -XTemplateHaskell
  -- > > putStrLn $(stringE . show =<< typeConstructors ''Bool)
  -- > [ ('False, [])
  -- > , ('True, [])
  -- > ]
  --
  -- > > putStrLn $(stringE . show =<< typeConstructors ''[])
  -- > [ ('[], [])
  -- > , ('(:), [VarT ''a, AppT ListT (VarT ''a)])
  -- > ]
  --
  -- > > putStrLn $(stringE . show =<< typeConstructors ''(,))
  -- > [('(,), [VarT (mkName "a"), VarT (mkName "b")])]
  --
  -- > > data Point  =  Pt Int Int
  -- > > putStrLn $(stringE . show =<< typeConstructors ''Point)
  -- > [('Pt,[ConT ''Int, ConT ''Int])]
  typeConstructors :: Name -> Q [(Name,[Type])]
  typeConstructors t  =  fmap (map normalize . cons) $ reify t
    where
  #if __GLASGOW_HASKELL__ < 800
    cons (TyConI (DataD    _ _ _   cs _))  =  cs
    cons (TyConI (NewtypeD _ _ _   c  _))  =  [c]
  #else
    cons (TyConI (DataD    _ _ _ _ cs _))  =  cs
    cons (TyConI (NewtypeD _ _ _ _ c  _))  =  [c]
  #endif
    cons _  =  []
    normalize (NormalC n ts)   =  (n,map snd ts)
    normalize (RecC    n ts)   =  (n,map trd ts)
    normalize (InfixC  t1 n t2)  =  (n,[snd t1,snd t2])
    normalize _  =  errorOn "typeConstructors"
                 $  "unexpected unhandled case when called with " ++ show t
    trd (x,y,z)  =  z
  
  -- |
  -- Given a type 'Name',
  -- returns a list of its type constructors 'Name's
  -- tupled with its arity and whether the constructor is recursive.
  --
  -- > > :set -XTemplateHaskell
  -- > > data Lst a  =  a :- Lst a | Nil  deriving Show
  -- > > putStrLn $(stringE . show =<< typeConstructorsWithArityAndRecursiveness ''Lst)
  -- > [ ('(:-),2,True)
  -- > , ('Nil,0,False)
  -- > ]
  --
  -- The recursive item does not work properly for some built-in types such as list
  -- as they have special representations within TH.
  --
  -- > > putStrLn $(stringE . show =<< typeConstructorsWithArityAndRecursiveness ''[])
  -- > [ ('[],0,False)
  -- > , ('(:),2,False)
  -- > ]
  typeConstructorsWithArityAndRecursiveness :: Name -> Q [(Name,Int,Bool)]
  typeConstructorsWithArityAndRecursiveness t  =  do
    cs <- typeConstructors t
    mapM ar cs
    where
    ar (n,ts)  =  do
      ns <- cascadingSubtypeNames ts
      return (n, length ts, t `elem` ns)
  
  -- |
  -- Is the given 'Name' a type synonym?
  --
  -- > > putStrLn $(stringE . show =<< isTypeSynonym 'show)
  -- > False
  --
  -- > > putStrLn $(stringE . show =<< isTypeSynonym ''Char)
  -- > False
  --
  -- > > putStrLn $(stringE . show =<< isTypeSynonym ''String)
  -- > True
  isTypeSynonym :: Name -> Q Bool
  isTypeSynonym  =  fmap is . reify
    where
    is (TyConI (TySynD _ _ _))  =  True
    is _                        =  False
  
  -- |
  -- Resolves a type synonym.
  --
  -- > > putStrLn $(stringE . show =<< typeSynonymType ''String)
  -- > AppT ListT (ConT ''Char)
  typeSynonymType :: Name -> Q Type
  typeSynonymType t  =  fmap typ $ reify t
    where
    typ (TyConI (TySynD _ _ t'))  =  t'
    typ _  =  errorOn "typeSynonymType" $ "not a type synonym: " ++ show t
  
  -- Append to instance contexts in a declaration.
  --
  -- > sequence [[|Eq b|],[|Eq c|]] |=>| [t|instance Eq a => Cl (Ty a) where f=g|]
  -- > == [t| instance (Eq a, Eq b, Eq c) => Cl (Ty a) where f  =  g |]
  (|=>|) :: Cxt -> DecsQ -> DecsQ
  c |=>| qds  =  map (=>++ c) `fmap` qds
    where
  #if __GLASGOW_HASKELL__ < 800
    (InstanceD   c ts ds) =>++ c'  =  InstanceD   (c++c') ts ds
  #else
    (InstanceD o c ts ds) =>++ c'  =  InstanceD o (c++c') ts ds
  #endif
    d                     =>++ _   =  d
  
  -- > nubMerge xs ys == nub (merge xs ys)
  -- > nubMerge xs ys == nub (sort (xs ++ ys))
  nubMerge :: Ord a => [a] -> [a] -> [a]
  nubMerge [] ys  =  ys
  nubMerge xs []  =  xs
  nubMerge (x:xs) (y:ys) | x < y      =  x :    xs  `nubMerge` (y:ys)
                         | x > y      =  y : (x:xs) `nubMerge`    ys
                         | otherwise  =  x :    xs  `nubMerge`    ys
  
  nubMerges :: Ord a => [[a]] -> [a]
  nubMerges  =  foldr nubMerge []
  
  #else
  -- When using Hugs or other compiler without Template Haskell
  
  errorNotGHC :: String -> a
  errorNotGHC fn  =  errorOn fn "only defined when using GHC"
  
  deriveListable :: a
  deriveListable  =  errorNotGHC "deriveListable"
  
  deriveListableIfNeeded :: a
  deriveListableIfNeeded  =  errorNotGHC "deriveListableIfNeeded"
  
  deriveListableCascading :: a
  deriveListableCascading  =  errorNotGHC "deriveListableCascading"
  
  deriveTiers :: a
  deriveTiers  =  errorNotGHC "deriveTiers"
  
  deriveList :: a
  deriveList  =  errorNotGHC "deriveList"
  
  -- closing #ifdef __GLASGOW_HASKELL__
  #endif
  
  errorOn :: String -> String -> a
  errorOn fn msg  =  error $ "Test.LeanCheck.Derive." ++ fn ++ ": " ++ msg
  