{-# LANGUAGE CPP #-}
{-# LANGUAGE GADTs #-}
{-|
Module:      Data.Functor.Deriving.Internal
Copyright:   (C) 2015-2017 Ryan Scott
License:     BSD-style (see the file LICENSE)
Maintainer:  Ryan Scott
Portability: Template Haskell

The machinery needed to derive 'Foldable', 'Functor', and 'Traversable' instances.

For more info on how deriving @Functor@ works, see
<https://ghc.haskell.org/trac/ghc/wiki/Commentary/Compiler/DeriveFunctor this GHC wiki page>.
-}
module Data.Functor.Deriving.Internal (
      -- * 'Foldable'
      deriveFoldable
    , makeFoldMap
    , makeFoldr
    , makeFold
    , makeFoldl
      -- * 'Functor'
    , deriveFunctor
    , makeFmap
      -- * 'Traversable'
    , deriveTraversable
    , makeTraverse
    , makeSequenceA
    , makeMapM
    , makeSequence
    ) where

import           Control.Monad (guard, zipWithM)

import           Data.Deriving.Internal
import           Data.Either (rights)
import           Data.List
import qualified Data.Map as Map (keys, lookup)
import           Data.Maybe

import           Language.Haskell.TH.Lib
import           Language.Haskell.TH.Syntax

-- | Generates a 'Foldable' instance declaration for the given data type or data
-- family instance.
deriveFoldable :: Name -> Q [Dec]
deriveFoldable = deriveFunctorClass Foldable

-- | Generates a lambda expression which behaves like 'foldMap' (without requiring a
-- 'Foldable' instance).
makeFoldMap :: Name -> Q Exp
makeFoldMap = makeFunctorFun FoldMap

-- | Generates a lambda expression which behaves like 'foldr' (without requiring a
-- 'Foldable' instance).
makeFoldr :: Name -> Q Exp
makeFoldr = makeFunctorFun Foldr

-- | Generates a lambda expression which behaves like 'fold' (without requiring a
-- 'Foldable' instance).
makeFold :: Name -> Q Exp
makeFold name = makeFoldMap name `appE` varE idValName

-- | Generates a lambda expression which behaves like 'foldl' (without requiring a
-- 'Foldable' instance).
makeFoldl :: Name -> Q Exp
makeFoldl name = do
  f <- newName "f"
  z <- newName "z"
  t <- newName "t"
  lamE [varP f, varP z, varP t] $
    appsE [ varE appEndoValName
          , appsE [ varE getDualValName
                  , appsE [ makeFoldMap name, foldFun f, varE t]
                  ]
          , varE z
          ]
  where
    foldFun :: Name -> Q Exp
    foldFun n = infixApp (conE dualDataName)
                         (varE composeValName)
                         (infixApp (conE endoDataName)
                                   (varE composeValName)
                                   (varE flipValName `appE` varE n)
                         )

-- | Generates a 'Functor' instance declaration for the given data type or data
-- family instance.
deriveFunctor :: Name -> Q [Dec]
deriveFunctor = deriveFunctorClass Functor

-- | Generates a lambda expression which behaves like 'fmap' (without requiring a
-- 'Functor' instance).
makeFmap :: Name -> Q Exp
makeFmap = makeFunctorFun Fmap

-- | Generates a 'Traversable' instance declaration for the given data type or data
-- family instance.
deriveTraversable :: Name -> Q [Dec]
deriveTraversable = deriveFunctorClass Traversable

-- | Generates a lambda expression which behaves like 'traverse' (without requiring a
-- 'Traversable' instance).
makeTraverse :: Name -> Q Exp
makeTraverse = makeFunctorFun Traverse

-- | Generates a lambda expression which behaves like 'sequenceA' (without requiring a
-- 'Traversable' instance).
makeSequenceA :: Name -> Q Exp
makeSequenceA name = makeTraverse name `appE` varE idValName

-- | Generates a lambda expression which behaves like 'mapM' (without requiring a
-- 'Traversable' instance).
makeMapM :: Name -> Q Exp
makeMapM name = do
  f <- newName "f"
  lam1E (varP f) . infixApp (varE unwrapMonadValName) (varE composeValName) $
                   makeTraverse name `appE` wrapMonadExp f
  where
    wrapMonadExp :: Name -> Q Exp
    wrapMonadExp n = infixApp (conE wrapMonadDataName) (varE composeValName) (varE n)

-- | Generates a lambda expression which behaves like 'sequence' (without requiring a
-- 'Traversable' instance).
makeSequence :: Name -> Q Exp
makeSequence name = makeMapM name `appE` varE idValName

-------------------------------------------------------------------------------
-- Code generation
-------------------------------------------------------------------------------

-- | Derive a class instance declaration (depending on the FunctorClass argument's value).
deriveFunctorClass :: FunctorClass -> Name -> Q [Dec]
deriveFunctorClass fc name = withType name fromCons where
  fromCons :: Name -> Cxt -> [TyVarBndr] -> [Con] -> Maybe [Type] -> Q [Dec]
  fromCons name' ctxt tvbs cons mbTys = (:[]) `fmap` do
    (instanceCxt, instanceType)
        <- buildTypeInstance fc name' ctxt tvbs mbTys
    instanceD (return instanceCxt)
              (return instanceType)
              (functorFunDecs fc cons)

-- | Generates a declaration defining the primary function(s) corresponding to a
-- particular class (fmap for Functor, foldr and foldMap for Foldable, and
-- traverse for Traversable).
--
-- For why both foldr and foldMap are derived for Foldable, see Trac #7436.
functorFunDecs :: FunctorClass -> [Con] -> [Q Dec]
functorFunDecs fc cons = map makeFunD $ functorClassToFuns fc where
  makeFunD :: FunctorFun -> Q Dec
  makeFunD ff =
    funD (functorFunName ff)
         [ clause []
                  (normalB $ makeFunctorFunForCons ff cons)
                  []
         ]

-- | Generates a lambda expression which behaves like the FunctorFun argument.
makeFunctorFun :: FunctorFun -> Name -> Q Exp
makeFunctorFun ff name = withType name fromCons where
  fromCons :: Name -> Cxt -> [TyVarBndr] -> [Con] -> Maybe [Type] -> Q Exp
  fromCons name' ctxt tvbs cons mbTys =
    -- We force buildTypeInstance here since it performs some checks for whether
    -- or not the provided datatype can actually have fmap/foldr/traverse/etc.
    -- implemented for it, and produces errors if it can't.
    buildTypeInstance (functorFunToClass ff) name' ctxt tvbs mbTys
      `seq` makeFunctorFunForCons ff cons

-- | Generates a lambda expression for the given constructors.
-- All constructors must be from the same type.
makeFunctorFunForCons :: FunctorFun -> [Con] -> Q Exp
makeFunctorFunForCons ff cons = do
  argNames <- mapM newName $ catMaybes [ Just "f"
                                       , guard (ff == Foldr) >> Just "z"
                                       , Just "value"
                                       ]
  let mapFun:others = argNames
      z     = head others -- If we're deriving foldr, this will be well defined
                          -- and useful. Otherwise, it'll be ignored.
      value = last others
  lamE (map varP argNames)
      . appsE
      $ [ varE $ functorFunConstName ff
        , if null cons
             then appE (varE errorValName)
                       (stringE $ "Void " ++ nameBase (functorFunName ff))
             else caseE (varE value)
                        (map (makeFunctorFunForCon ff z mapFun) cons)
        ] ++ map varE argNames

-- | Generates a lambda expression for a single constructor.
makeFunctorFunForCon :: FunctorFun -> Name -> Name -> Con -> Q Match
makeFunctorFunForCon ff z mapFun con = do
  let conName = constructorName con
  (ts, tvMap) <- reifyConTys1 (functorFunToClass ff) [mapFun] conName
  argNames    <- newNameList "_arg" $ length ts
  makeFunctorFunForArgs ff z tvMap conName ts argNames

-- | Generates a lambda expression for a single constructor's arguments.
makeFunctorFunForArgs :: FunctorFun
                      -> Name
                      -> TyVarMap1
                      -> Name
                      -> [Type]
                      -> [Name]
                      -> Q Match
makeFunctorFunForArgs ff z tvMap conName tys args =
  match (conP conName $ map varP args)
        (normalB $ functorFunCombine ff conName z args mappedArgs)
        []
  where
    mappedArgs :: Q [Either Exp Exp]
    mappedArgs = zipWithM (makeFunctorFunForArg ff tvMap conName) tys args

-- | Generates a lambda expression for a single argument of a constructor.
--  The returned value is 'Right' if its type mentions the last type
-- parameter. Otherwise, it is 'Left'.
makeFunctorFunForArg :: FunctorFun
                     -> TyVarMap1
                     -> Name
                     -> Type
                     -> Name
                     -> Q (Either Exp Exp)
makeFunctorFunForArg ff tvMap conName ty tyExpName =
  makeFunctorFunForType ff tvMap conName True ty `appEitherE` varE tyExpName

-- | Generates a lambda expression for a specific type. The returned value is
-- 'Right' if its type mentions the last type parameter. Otherwise,
-- it is 'Left'.
makeFunctorFunForType :: FunctorFun
                      -> TyVarMap1
                      -> Name
                      -> Bool
                      -> Type
                      -> Q (Either Exp Exp)
makeFunctorFunForType ff tvMap conName covariant (VarT tyName) =
  case Map.lookup tyName tvMap of
    Just (OneName mapName) ->
      fmap Right $ if covariant
                      then varE mapName
                      else contravarianceError conName
    -- Invariant: this should only happen when deriving fmap
    Nothing -> fmap Left $ functorFunTriv ff
makeFunctorFunForType ff tvMap conName covariant (SigT ty _) =
  makeFunctorFunForType ff tvMap conName covariant ty
makeFunctorFunForType ff tvMap conName covariant (ForallT _ _ ty) =
  makeFunctorFunForType ff tvMap conName covariant ty
makeFunctorFunForType ff tvMap conName covariant ty =
  let tyCon  :: Type
      tyArgs :: [Type]
      tyCon:tyArgs = unapplyTy ty

      numLastArgs :: Int
      numLastArgs = min 1 $ length tyArgs

      lhsArgs, rhsArgs :: [Type]
      (lhsArgs, rhsArgs) = splitAt (length tyArgs - numLastArgs) tyArgs

      tyVarNames :: [Name]
      tyVarNames = Map.keys tvMap

      mentionsTyArgs :: Bool
      mentionsTyArgs = any (`mentionsName` tyVarNames) tyArgs

      makeFunctorFunTuple :: ([Q Pat] -> Q Pat) -> (Int -> Name) -> Int
                          -> Q (Either Exp Exp)
      makeFunctorFunTuple mkTupP mkTupleDataName n = do
         args <- mapM newName $ catMaybes [ Just "x"
                                          , guard (ff == Foldr) >> Just "z"
                                          ]
         xs <- newNameList "_tup" n

         let x = head args
             z = last args
         fmap Right $ lamE (map varP args) $ caseE (varE x)
              [ match (mkTupP $ map varP xs)
                      (normalB $ functorFunCombine ff
                                                   (mkTupleDataName n)
                                                   z
                                                   xs
                                                   (zipWithM makeFunctorFunTupleField tyArgs xs)
                      )
                      []
              ]

      makeFunctorFunTupleField :: Type -> Name -> Q (Either Exp Exp)
      makeFunctorFunTupleField fieldTy fieldName =
        makeFunctorFunForType ff tvMap conName covariant fieldTy
          `appEitherE` varE fieldName

      fc :: FunctorClass
      fc = functorFunToClass ff

   in case tyCon of
     ArrowT
       | not (allowFunTys fc) -> noFunctionsError conName
       | mentionsTyArgs, [argTy, resTy] <- tyArgs ->
         do x <- newName "x"
            b <- newName "b"
            fmap Right . lamE [varP x, varP b] $
              covFunctorFun covariant resTy `appE` (varE x `appE`
                (covFunctorFun (not covariant) argTy `appE` varE b))
         where
           covFunctorFun :: Bool -> Type -> Q Exp
           covFunctorFun cov = fmap fromEither . makeFunctorFunForType ff tvMap conName cov
#if MIN_VERSION_template_haskell(2,6,0)
     UnboxedTupleT n
       | n > 0 && mentionsTyArgs -> makeFunctorFunTuple unboxedTupP unboxedTupleDataName n
#endif
     TupleT n
       | n > 0 && mentionsTyArgs -> makeFunctorFunTuple tupP tupleDataName n
     _ -> do
         itf <- isTyFamily tyCon
         if any (`mentionsName` tyVarNames) lhsArgs || (itf && mentionsTyArgs)
           then outOfPlaceTyVarError fc conName
           else if any (`mentionsName` tyVarNames) rhsArgs
                  then fmap Right . functorFunApp ff . appsE $
                         ( varE (functorFunName ff)
                         : map (fmap fromEither . makeFunctorFunForType ff tvMap conName covariant)
                                rhsArgs
                         )
                  else fmap Left $ functorFunTriv ff

-------------------------------------------------------------------------------
-- Class-specific constants
-------------------------------------------------------------------------------

-- | A representation of which class is being derived.
data FunctorClass = Functor | Foldable | Traversable

instance ClassRep FunctorClass where
    arity _ = 1

    allowExQuant Foldable = True
    allowExQuant _        = False

    fullClassName Functor     = functorTypeName
    fullClassName Foldable    = foldableTypeName
    fullClassName Traversable = traversableTypeName

    classConstraint fClass 1 = Just $ fullClassName fClass
    classConstraint  _      _ = Nothing

-- | A representation of which function is being generated.
data FunctorFun = Fmap | Foldr | FoldMap | Traverse
  deriving Eq

instance Show FunctorFun where
    showsPrec _ Fmap     = showString "fmap"
    showsPrec _ Foldr    = showString "foldr"
    showsPrec _ FoldMap  = showString "foldMap"
    showsPrec _ Traverse = showString "traverse"

functorFunConstName :: FunctorFun -> Name
functorFunConstName Fmap     = fmapConstValName
functorFunConstName Foldr    = foldrConstValName
functorFunConstName FoldMap  = foldMapConstValName
functorFunConstName Traverse = traverseConstValName

functorFunName :: FunctorFun -> Name
functorFunName Fmap     = fmapValName
functorFunName Foldr    = foldrValName
functorFunName FoldMap  = foldMapValName
functorFunName Traverse = traverseValName

functorClassToFuns :: FunctorClass -> [FunctorFun]
functorClassToFuns Functor     = [Fmap]
functorClassToFuns Foldable    = [Foldr, FoldMap]
functorClassToFuns Traversable = [Traverse]

functorFunToClass :: FunctorFun -> FunctorClass
functorFunToClass Fmap     = Functor
functorFunToClass Foldr    = Foldable
functorFunToClass FoldMap  = Foldable
functorFunToClass Traverse = Traversable

allowFunTys :: FunctorClass -> Bool
allowFunTys Functor = True
allowFunTys _       = False

-------------------------------------------------------------------------------
-- Assorted utilities
-------------------------------------------------------------------------------

-- See Trac #7436 for why explicit lambdas are used
functorFunTriv :: FunctorFun -> Q Exp
functorFunTriv Fmap = do
  x <- newName "x"
  lam1E (varP x) $ varE x
-- We filter out trivial expressions from derived foldr, foldMap, and traverse
-- implementations, so if we attempt to call functorFunTriv on one of those
-- methods, we've done something wrong.
functorFunTriv ff = return . error $ "functorFunTriv: " ++ show ff

functorFunApp :: FunctorFun -> Q Exp -> Q Exp
functorFunApp Foldr e = do
  x <- newName "x"
  z <- newName "z"
  lamE [varP x, varP z] $ appsE [e, varE z, varE x]
functorFunApp _ e = e

functorFunCombine :: FunctorFun
                  -> Name
                  -> Name
                  -> [Name]
                  -> Q [Either Exp Exp]
                  -> Q Exp
functorFunCombine Fmap     = fmapCombine
functorFunCombine Foldr    = foldrCombine
functorFunCombine FoldMap  = foldMapCombine
functorFunCombine Traverse = traverseCombine

fmapCombine :: Name
            -> Name
            -> [Name]
            -> Q [Either Exp Exp]
            -> Q Exp
fmapCombine conName _ _ = fmap (foldl' AppE (ConE conName) . fmap fromEither)

-- foldr, foldMap, and traverse are handled differently from fmap, since
-- they filter out subexpressions whose types do not mention the last
-- type parameter. See
-- https://ghc.haskell.org/trac/ghc/wiki/Commentary/Compiler/DeriveFunctor#AlternativestrategyforderivingFoldableandTraversable
-- for further discussion.

foldrCombine :: Name
             -> Name
             -> [Name]
             -> Q [Either Exp Exp]
             -> Q Exp
foldrCombine _ zName _ = fmap (foldr AppE (VarE zName) . rights)

foldMapCombine :: Name
               -> Name
               -> [Name]
               -> Q [Either Exp Exp]
               -> Q Exp
foldMapCombine _ _ _ = fmap (go . rights)
  where
    go :: [Exp] -> Exp
    go [] = VarE memptyValName
    go es = foldr1 (AppE . AppE (VarE mappendValName)) es

traverseCombine :: Name
                -> Name
                -> [Name]
                -> Q [Either Exp Exp]
                -> Q Exp
traverseCombine conName _ args essQ = do
    ess <- essQ

    let argTysTyVarInfo :: [Bool]
        argTysTyVarInfo = map isRight ess

        argsWithTyVar, argsWithoutTyVar :: [Name]
        (argsWithTyVar, argsWithoutTyVar) = partitionByList argTysTyVarInfo args

        conExpQ :: Q Exp
        conExpQ
          | null argsWithTyVar
          = appsE (conE conName:map varE argsWithoutTyVar)
          | otherwise = do
              bs <- newNameList "b" $ length args
              let bs'  = filterByList  argTysTyVarInfo bs
                  vars = filterByLists argTysTyVarInfo
                                       (map varE bs) (map varE args)
              lamE (map varP bs') (appsE (conE conName:vars))

    conExp <- conExpQ

    let go :: [Exp] -> Exp
        go []     = VarE pureValName `AppE` conExp
        go [e] = VarE fmapValName `AppE` conExp `AppE` e
        go (e1:e2:es) = foldl' (\se1 se2 -> InfixE (Just se1) (VarE apValName) (Just se2))
          (VarE liftA2ValName `AppE` conExp `AppE` e1 `AppE` e2) es

    return . go . rights $ ess