    {-# LANGUAGE CPP #-}
    {-# LANGUAGE GADTs #-}

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

    Exports functions to mechanically derive 'Eq', 'Eq1', and 'Eq2' instances.

    Note: this is an internal module, and as such, the API presented here is not
    guaranteed to be stable, even between minor releases of this library.
    -}
    module Data.Eq.Deriving.Internal (
          -- * 'Eq'
          deriveEq
        , makeEq
        , makeNotEq
          -- * 'Eq1'
        , deriveEq1
    #if defined(NEW_FUNCTOR_CLASSES)
        , makeLiftEq
    #endif
        , makeEq1
    #if defined(NEW_FUNCTOR_CLASSES)
          -- * 'Eq2'
        , deriveEq2
        , makeLiftEq2
        , makeEq2
    #endif
        ) where

    import           Data.Deriving.Internal
    import           Data.List (foldl1', partition)
    import qualified Data.Map as Map

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

    -- | Generates an 'Eq' instance declaration for the given data type or data
    -- family instance.
    deriveEq :: Name -> Q [Dec]
    deriveEq = deriveEqClass Eq

    -- | Generates a lambda expression which behaves like '(==)' (without
    -- requiring an 'Eq' instance).
    makeEq :: Name -> Q Exp
    makeEq = makeEqClass Eq

    -- | Generates a lambda expression which behaves like '(/=)' (without
    -- requiring an 'Eq' instance).
    makeNotEq :: Name -> Q Exp
    makeNotEq name = do
        x1 <- newName "x1"
        x2 <- newName "x2"
        lamE [varP x1, varP x2] $ varE notValName `appE`
            (makeEq name `appE` varE x1 `appE` varE x2)

    -- | Generates an 'Eq1' instance declaration for the given data type or data
    -- family instance.
    deriveEq1 :: Name -> Q [Dec]
    deriveEq1 = deriveEqClass Eq1

    #if defined(NEW_FUNCTOR_CLASSES)
    -- | Generates a lambda expression which behaves like 'liftEq' (without
    -- requiring an 'Eq1' instance).
    --
    -- This function is not available with @transformers-0.4@.
    makeLiftEq :: Name -> Q Exp
    makeLiftEq = makeEqClass Eq1

    -- | Generates a lambda expression which behaves like 'eq1' (without
    -- requiring an 'Eq1' instance).
    makeEq1 :: Name -> Q Exp
    makeEq1 name = makeLiftEq name `appE` varE eqValName
    #else
    -- | Generates a lambda expression which behaves like 'eq1' (without
    -- requiring an 'Eq1' instance).
    makeEq1 :: Name -> Q Exp
    makeEq1 = makeEqClass Eq1
    #endif

    #if defined(NEW_FUNCTOR_CLASSES)
    -- | Generates an 'Eq2' instance declaration for the given data type or data
    -- family instance.
    --
    -- This function is not available with @transformers-0.4@.
    deriveEq2 :: Name -> Q [Dec]
    deriveEq2 = deriveEqClass Eq2

    -- | Generates a lambda expression which behaves like 'liftEq2' (without
    -- requiring an 'Eq2' instance).
    --
    -- This function is not available with @transformers-0.4@.
    makeLiftEq2 :: Name -> Q Exp
    makeLiftEq2 = makeEqClass Eq2

    -- | Generates a lambda expression which behaves like 'eq2' (without
    -- requiring an 'Eq2' instance).
    --
    -- This function is not available with @transformers-0.4@.
    makeEq2 :: Name -> Q Exp
    makeEq2 name = makeLiftEq name `appE` varE eqValName `appE` varE eqValName
    #endif

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

    -- | Derive an Eq(1)(2) instance declaration (depending on the EqClass
    -- argument's value).
    deriveEqClass :: EqClass -> Name -> Q [Dec]
    deriveEqClass eClass name = do
      info <- reifyDatatype name
      case info of
        DatatypeInfo { datatypeContext   = ctxt
                     , datatypeName      = parentName
    #if MIN_VERSION_th_abstraction(0,3,0)
                     , datatypeInstTypes = instTypes
    #else
                     , datatypeVars      = instTypes
    #endif
                     , datatypeVariant   = variant
                     , datatypeCons      = cons
                     } -> do
          (instanceCxt, instanceType)
              <- buildTypeInstance eClass parentName ctxt instTypes variant
          (:[]) `fmap` instanceD (return instanceCxt)
                                 (return instanceType)
                                 (eqDecs eClass instTypes cons)

    -- | Generates a declaration defining the primary function corresponding to a
    -- particular class ((==) for Eq, liftEq for Eq1, and
    -- liftEq2 for Eq2).
    eqDecs :: EqClass -> [Type] -> [ConstructorInfo] -> [Q Dec]
    eqDecs eClass instTypes cons =
        [ funD (eqName eClass)
               [ clause []
                        (normalB $ makeEqForCons eClass instTypes cons)
                        []
               ]
        ]

    -- | Generates a lambda expression which behaves like (==) (for Eq),
    -- liftEq (for Eq1), or liftEq2 (for Eq2).
    makeEqClass :: EqClass -> Name -> Q Exp
    makeEqClass eClass name = do
      info <- reifyDatatype name
      case info of
        DatatypeInfo { datatypeContext   = ctxt
                     , datatypeName      = parentName
    #if MIN_VERSION_th_abstraction(0,3,0)
                     , datatypeInstTypes = instTypes
    #else
                     , datatypeVars      = instTypes
    #endif
                     , datatypeVariant   = variant
                     , datatypeCons      = cons
                     } -> do
          -- We force buildTypeInstance here since it performs some checks for whether
          -- or not the provided datatype can actually have (==)/liftEq/etc.
          -- implemented for it, and produces errors if it can't.
          buildTypeInstance eClass parentName ctxt instTypes variant
            >> makeEqForCons eClass instTypes cons

    -- | Generates a lambda expression for (==)/liftEq/etc. for the
    -- given constructors. All constructors must be from the same type.
    makeEqForCons :: EqClass -> [Type] -> [ConstructorInfo] -> Q Exp
    makeEqForCons eClass instTypes cons = do
        value1 <- newName "value1"
        value2 <- newName "value2"
        eqDefn <- newName "eqDefn"
        eqs    <- newNameList "eq" $ arity eClass

        let lastTyVars = map varTToName $ drop (length instTypes - fromEnum eClass) instTypes
            tvMap      = Map.fromList $ zipWith (\x y -> (x, OneName y)) lastTyVars eqs

        lamE (map varP $
    #if defined(NEW_FUNCTOR_CLASSES)
                         eqs ++
    #endif
                         [value1, value2]
             ) . appsE
             $ [ varE $ eqConstName eClass
               , letE [ funD eqDefn $ map (makeCaseForCon eClass tvMap) patMatchCons
                                   ++ fallThroughCase
                      ] $ varE eqDefn `appE` varE value1 `appE` varE value2
               ]
    #if defined(NEW_FUNCTOR_CLASSES)
                 ++ map varE eqs
    #endif
                 ++ [varE value1, varE value2]
      where
        nullaryCons, nonNullaryCons :: [ConstructorInfo]
        (nullaryCons, nonNullaryCons) = partition isNullaryCon cons

        tagMatchCons, patMatchCons :: [ConstructorInfo]
        (tagMatchCons, patMatchCons)
          | length nullaryCons > 10 = (nullaryCons, nonNullaryCons)
          | otherwise               = ([],          cons)

        fallThroughCase :: [Q Clause]
        fallThroughCase
          | null tagMatchCons = case patMatchCons of
              []  -> [makeFallThroughCaseTrue]  -- No constructors: _ == _ = True
              [_] -> []                         -- One constructor: no fall-through case
              _   -> [makeFallThroughCaseFalse] -- Two or more constructors:
                                                --   _ == _ = False
          | otherwise = [makeTagCase]

    makeTagCase :: Q Clause
    makeTagCase = do
        a     <- newName "a"
        aHash <- newName "a#"
        b     <- newName "b"
        bHash <- newName "b#"
        clause (map varP [a,b])
               (normalB $ untagExpr [(a, aHash), (b, bHash)] $
                   primOpAppExpr (varE aHash) eqIntHashValName (varE bHash)) []

    makeFallThroughCaseFalse, makeFallThroughCaseTrue :: Q Clause
    makeFallThroughCaseFalse = makeFallThroughCase falseDataName
    makeFallThroughCaseTrue  = makeFallThroughCase trueDataName

    makeFallThroughCase :: Name -> Q Clause
    makeFallThroughCase dataName = clause [wildP, wildP] (normalB $ conE dataName) []

    makeCaseForCon :: EqClass -> TyVarMap1 -> ConstructorInfo -> Q Clause
    makeCaseForCon eClass tvMap
      (ConstructorInfo { constructorName = conName, constructorFields = ts }) = do
        ts' <- mapM resolveTypeSynonyms ts
        let tsLen = length ts'
        as <- newNameList "a" tsLen
        bs <- newNameList "b" tsLen
        clause [conP conName (map varP as), conP conName (map varP bs)]
               (normalB $ makeCaseForArgs eClass tvMap conName ts' as bs)
               []

    makeCaseForArgs :: EqClass
                    -> TyVarMap1
                    -> Name
                    -> [Type]
                    -> [Name]
                    -> [Name]
                    -> Q Exp
    makeCaseForArgs _ _ _ [] [] [] = conE trueDataName
    makeCaseForArgs eClass tvMap conName tys as bs =
        foldl1' (\q e -> infixApp q (varE andValName) e)
                (zipWith3 (makeCaseForArg eClass tvMap conName) tys as bs)

    makeCaseForArg :: EqClass
                   -> TyVarMap1
                   -> Name
                   -> Type
                   -> Name
                   -> Name
                   -> Q Exp
    makeCaseForArg _ _ _ (ConT tyName) a b = primEqExpr
      where
        aExpr, bExpr :: Q Exp
        aExpr = varE a
        bExpr = varE b

        makePrimEqExpr :: Name -> Q Exp
        makePrimEqExpr n = primOpAppExpr aExpr n bExpr

        primEqExpr :: Q Exp
        primEqExpr =
          case Map.lookup tyName primOrdFunTbl of
            Just (_, _, eq, _, _) -> makePrimEqExpr eq
            Nothing               -> infixApp aExpr (varE eqValName) bExpr
    makeCaseForArg eClass tvMap conName ty a b =
        makeCaseForType eClass tvMap conName ty `appE` varE a `appE` varE b

    makeCaseForType :: EqClass
                    -> TyVarMap1
                    -> Name
                    -> Type
                    -> Q Exp
    #if defined(NEW_FUNCTOR_CLASSES)
    makeCaseForType _ tvMap _ (VarT tyName) =
        varE $ case Map.lookup tyName tvMap of
          Just (OneName eq) -> eq
          Nothing           -> eqValName
    #else
    makeCaseForType _ _ _ VarT{} = varE eqValName
    #endif
    makeCaseForType eClass tvMap conName (SigT ty _)      = makeCaseForType eClass tvMap conName ty
    makeCaseForType eClass tvMap conName (ForallT _ _ ty) = makeCaseForType eClass tvMap conName ty
    #if defined(NEW_FUNCTOR_CLASSES)
    makeCaseForType eClass tvMap conName ty = do
        let tyCon :: Type
            tyArgs :: [Type]
            tyCon:tyArgs = unapplyTy ty

            numLastArgs :: Int
            numLastArgs = min (arity eClass) (length tyArgs)

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

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

        itf <- isTyFamily tyCon
        if any (`mentionsName` tyVarNames) lhsArgs
              || itf && any (`mentionsName` tyVarNames) tyArgs
           then outOfPlaceTyVarError eClass conName
           else if any (`mentionsName` tyVarNames) rhsArgs
                   then appsE $ [ varE . eqName $ toEnum numLastArgs]
                                ++ map (makeCaseForType eClass tvMap conName) rhsArgs
                   else varE eqValName
    #else
    makeCaseForType eClass tvMap conName ty = do
      let varNames = Map.keys tvMap

      a' <- newName "a'"
      b' <- newName "b'"
      case varNames of
        [] -> varE eqValName
        varName:_ ->
          if mentionsName ty varNames
             then lamE (map varP [a',b']) $ varE eq1ValName
                    `appE` (makeFmapApplyNeg eClass conName ty varName `appE` varE a')
                    `appE` (makeFmapApplyNeg eClass conName ty varName `appE` varE b')
             else varE eqValName
    #endif

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

    -- | A representation of which @Eq@ variant is being derived.
    data EqClass = Eq
                 | Eq1
    #if defined(NEW_FUNCTOR_CLASSES)
                 | Eq2
    #endif
      deriving (Bounded, Enum)

    instance ClassRep EqClass where
        arity = fromEnum

        allowExQuant _ = True

        fullClassName Eq  = eqTypeName
        fullClassName Eq1 = eq1TypeName
    #if defined(NEW_FUNCTOR_CLASSES)
        fullClassName Eq2 = eq2TypeName
    #endif

        classConstraint eClass i
          | eMin <= i && i <= eMax = Just $ fullClassName (toEnum i :: EqClass)
          | otherwise              = Nothing
          where
            eMin, eMax :: Int
            eMin = fromEnum (minBound :: EqClass)
            eMax = fromEnum eClass

    eqConstName :: EqClass -> Name
    eqConstName Eq  = eqConstValName
    #if defined(NEW_FUNCTOR_CLASSES)
    eqConstName Eq1 = liftEqConstValName
    eqConstName Eq2 = liftEq2ConstValName
    #else
    eqConstName Eq1 = eq1ConstValName
    #endif

    eqName :: EqClass -> Name
    eqName Eq  = eqValName
    #if defined(NEW_FUNCTOR_CLASSES)
    eqName Eq1 = liftEqValName
    eqName Eq2 = liftEq2ValName
    #else
    eqName Eq1 = eq1ValName
    #endif