{-# language CPP                   #-}
{-# language ExplicitNamespaces    #-}
{-# language MultiWayIf            #-}
{-# language TemplateHaskellQuotes #-}

-- | Main module of @kind-generics-th@.
-- Please refer to the @README@ file for documentation on how to use this package.
module Generics.Kind.TH
  ( deriveGenericK
  , deriveGenericKQuiet
  , preDeriveGenericK
  , postDeriveGenericK
  ) where

import           Control.Applicative
import           Control.Monad
import qualified Data.Kind                    as Kind
import           Data.List
import           Data.Maybe
import           Data.Type.Equality           (type (~~))
import           Fcf.Family.TH                (fcfify, isTypeFamilyOrSynonym, promoteNDFamily)
import           GHC.Generics                 as Generics hiding (conIsRecord, conName,
                                                           datatypeName)
import           Generics.Kind
import           Language.Haskell.TH          as TH
import           Language.Haskell.TH.Syntax   as TH
import           Language.Haskell.TH.Datatype as THAbs
import           Language.Haskell.TH.Datatype.TyVarBndr

#if MIN_VERSION_template_haskell(2,15,0)
import           GHC.Classes                  (IP)
#endif

-- | Given the 'Name' of a data type (or, the 'Name' of a constructor belonging
-- to a data type), generate 'GenericK' instances for that data type. You will
-- likely need to enable most of these language extensions in order for GHC to
-- accept the generated code:
--
-- * @DataKinds@
--
-- * @EmptyCase@ (if using an empty data type)
--
-- * @FlexibleInstances@
--
-- * @MultiParamTypeClasses@
--
-- * @PolyKinds@ (if using a poly-kinded data type)
--
-- * @TemplateHaskell@
--
-- * @TypeFamilies@
--
-- If the data type uses type families, 'deriveGenericK' warns that it
-- skips the 'GenericK' instances that require special support for it
--
-- - Use 'preDeriveGenericK' and 'postDeriveGenericK' to support type families.
-- - Use 'deriveGenericKQuiet' to silence the warnings.
deriveGenericK :: Name -> Q [Dec]
deriveGenericK = deriveGenericKWarnIf True

-- | Variant of 'deriveGenericK' that doesn't emit warnings.
deriveGenericKQuiet :: Name -> Q [Dec]
deriveGenericKQuiet = deriveGenericKWarnIf False

deriveGenericKWarnIf :: Bool -> Name -> Q [Dec]
deriveGenericKWarnIf warn name = uncurry (++) <$> deriveGenericK' (NoFamilies warn) name

-- | Generate 'GenericK' instances for data types that may mention
-- type families.
--
-- This 'preDeriveGenericK' is to be used in combination with
-- 'postDeriveGenericK'. These two functions let us stage the compilation of
-- the generated type instances, because GHC cannot compile them in a single
-- group.
--
-- - 'preDeriveGenericK' generates type instances to promote type families
--   that occur in the given data types (using 'fcfify'; see
--   <https://hackage.haskell.org/package/fcf-family fcf-family>).
--   The 'GenericK' instances are not produced at this stage,
--   they are accumulated in some internal global queue.
-- - 'postDeriveGenericK' produces all of the accumulated 'GenericK' instances.
--   It should be called in a slice separated from 'preDeriveGenericK'.
--   Multiple calls to `preDeriveGenericK` may precede 'postDeriveGenericK'.
--
-- @
-- 'preDeriveGenericK' ''MyT1
-- 'preDeriveGenericK' ''MyT2
-- 'preDeriveGenericK' ''MyT3
-- 'postDeriveGenericK'
-- @
--
-- You will need to enable the extensions @UndecidableInstances@ and @PolyKinds@
-- (even if your data types are not poly-kinded)
-- in addition to those mentioned in the documentation of 'deriveGenericK'.
preDeriveGenericK :: Name -> Q [Dec]
preDeriveGenericK n = do
  (pre, post) <- deriveGenericK' YesFamilies n
  pushGenericKQueue post
  pure pre

-- | See 'preDeriveGenericK'.
postDeriveGenericK :: Q [Dec]
postDeriveGenericK = takeGenericKQueue

-- | Flag to control support for type families, because that requires a
-- different API (preDeriveGenericK, postDeriveGenericK instead of
-- deriveGenericK).
data FamilyFriendliness
  = NoFamilies Bool -- ^ Whether to warn when a type family is detected.
  | YesFamilies

-- | Return a pair of:
--
-- - 'fcfify'-generated instances
-- - 'GenericK' instances
deriveGenericK' :: FamilyFriendliness -> Name -> Q ([Dec], [Dec])
deriveGenericK' familyFriendliness n = do
  DatatypeInfo{ datatypeName      = dataName
              , datatypeInstTypes = univVars
              , datatypeVariant   = variant
              , datatypeCons      = cons
              } <- reifyDatatype n
  cons' <- traverse resolveConSynonyms cons
  let deriveInsts :: [Type] -> [Type] -> Q [Dec]
      deriveInsts argsToKeep argsToDrop = do
        inst <- deriveGenericKFor argsToKeep argsToDrop
        case argsToKeep of
          [] -> pure [inst]
          (argToDrop:argsToKeep') -> do
            argToDrop' <- resolveTypeSynonyms argToDrop
            if |  -- Can the argument to drop be eta-reduced?
                  Just argNameToDrop <- distinctTyVarType (freeVariables argsToKeep')
                                                          argToDrop'
                  -- Check for dependent quantification, which we currently can't handle.
               ,  argNameToDrop `notElem`
                    freeVariables (map typeKind argsToDrop
                                ++ map tvKind (gatherExistentials cons'))
               -> do let allInnerTypes  = gatherConstraints cons' ++ gatherFields cons'
                     -- Check if the argument appears in a type family application.
                     inTyFamApp <- or <$> traverse (isInTypeFamilyApp argNameToDrop)
                                                   allInnerTypes
                     case familyFriendliness of
                       NoFamilies warn | inTyFamApp -> do
                         -- Found type family application when family suppport is disabled.
                         -- Emit a warning and don't generate GenericK instances for fewer argsToKeep.
                         when warn (reportWarning $ tyFamWarning n dataName argsToKeep argsToDrop)
                         pure [inst]
                       _ -> (inst:) <$> deriveInsts argsToKeep' (argToDrop':argsToDrop)
               |  otherwise
               -> pure [inst]

      -- Generate a single GenericK instance for a given set of data type
      -- arguments and indexed arguments.
      deriveGenericKFor :: [Type] -> [Type] -> Q Dec
      deriveGenericKFor argsToKeep argsToDrop = do
        let argNamesToDrop = map varTToName argsToDrop
            kind = foldr ((\x y -> ArrowT `AppT` x `AppT` y) . typeKind)
                         (ConT ''Kind.Type) argsToDrop
            dataApp = pure $ SigT (foldr (flip AppT) (ConT dataName) argsToKeep) kind
        instanceD (pure [])
                  (conT ''GenericK `appT` dataApp)
                  [ tySynInstDCompat ''RepK Nothing [dataApp] $
                      deriveRepK dataName argNamesToDrop variant cons'
                  , deriveFromK cons'
                  , deriveToK cons'
                  ]

  insts <- deriveInsts (reverse univVars) []
  fcfInsts <- takeFcfifyQueue
  pure (fcfInsts, insts)

-- | Warning to show when a type family is found by 'deriveGenericK'.
tyFamWarning :: Name -> Name -> [Type] -> [Type] -> String
tyFamWarning name dataName argsToKeep' argsToDrop' =
  let argsToKeep = getVarTName <$> reverse argsToKeep'
      argsToDrop = getVarTName <$> argsToDrop'
  in tyFamWarning' name dataName argsToKeep argsToDrop

-- | 'tyFamWarning' with variable names instead of Type, all in left-to-right order.
tyFamWarning' :: Name -> Name -> [String] -> [String] -> String
tyFamWarning' name dataName argsToKeep argsToDrop = unlines $
  ("Found type family in definition of "
    ++ quoteName name ++ ". Some instances have been skipped.") :
  map ("    " ++) (
    "Declared instances:" :
    showDeclaredInstances dataName argsToKeep argsToDrop ++
    "Skipped instances:" :
    showSkippedInstances dataName argsToKeep ++
    "To enable type family support and obtain those skipped instances:" :
    ("\t$(preDeriveGenericK " ++ quoteName name ++ ")") :
    ("\t$(postDeriveGenericK " ++ quoteName name ++ ")") :
    "To silence this warning:" :
    ("\t$(deriveGenericKQuiet " ++ quoteName name ++ ")") :
    [])

-- | This assumes most uses are going to be unqualified names.
quoteName :: Name -> String
quoteName name@(Name _ (NameG DataName _ _)) = "'" ++ nameBase name
quoteName name = "''" ++ nameBase name

showDeclaredInstances :: Name -> [String] -> [String] -> [String]
showDeclaredInstances name argsToKeep argsToDrop =
  (\args -> "\tinstance GenericK " ++ showConArgs name (argsToKeep ++ args)) <$> inits argsToDrop

showSkippedInstances :: Name -> [String] -> [String]
showSkippedInstances name argsToKeep =
  (\args -> "\tinstance GenericK " ++ showConArgs name args) <$> init (inits argsToKeep)

-- We manually pretty-print the types to drop module qualifiers.
showConArgs :: Name -> [String] -> String
showConArgs name [] = nameBase name
showConArgs name args = "(" ++ intercalate " " (nameBase name : args) ++ ")"

-- | Find type variable stored in types coming from 'datatypeInstTypes'
-- (should be of the form (v :: k))
getVarTName :: Type -> String
getVarTName (SigT t _) = getVarTName t
getVarTName (VarT name) = nameBase name
getVarTName _ = "_a"

-- | @'distinctTyVarType' tvSet ty@ returns @'Just' tvTy@ if @ty@:
--
-- a. Is a type variable named @tvTy@, and
-- b. @tvTy@ is not an element of @tvSet@.
--
-- Otherwise, returns 'Nothing'.
distinctTyVarType :: [Name] -> Type -> Maybe Name
distinctTyVarType tvSet ty = do
  tvTy <- varTToName_maybe ty
  guard $ tvTy `notElem` tvSet
  pure tvTy

deriveRepK :: Name -> [Name]
           -> DatatypeVariant -> [ConstructorInfo] -> Q Type
deriveRepK dataName univVarNames dataVariant cons = do
  cons' <- traverse constructor cons
  metaData $ foldBal (\x y -> InfixT x ''(:+:) y) (ConT ''V1) cons'
  where
    metaData :: Type -> Q Type
    metaData t = do
      m   <- maybe (fail "Cannot fetch module name!")  pure (nameModule dataName)
      pkg <- maybe (fail "Cannot fetch package name!") pure (namePackage dataName)
      conT ''D1
        `appT` (promotedT 'MetaData `appT`
                litT (strTyLit (nameBase dataName)) `appT`
                litT (strTyLit m) `appT`
                litT (strTyLit pkg) `appT`
                promoteBool (isNewtypeVariant dataVariant))
        `appT` pure t

    constructor :: ConstructorInfo -> Q Type
    constructor ConstructorInfo{ constructorName       = conName
                               , constructorVars       = exTvbs
                               , constructorContext    = conCtxt
                               , constructorFields     = fields
                               , constructorStrictness = fieldStricts
                               , constructorVariant    = conVariant
                               } = do
      mbFi <- reifyFixity conName
      conT ''C1
        `appT` (promotedT 'MetaCons `appT`
                litT (strTyLit (nameBase conName)) `appT`
                fixityIPromotedType mbFi conIsInfix `appT`
                promoteBool conIsRecord)
        `appT` do prod <- foldBal (\x y -> InfixT x ''(:*:) y) (ConT ''U1) <$> selectors
                  ctxtProd <- context prod
                  existentials ctxtProd
      where
        conIsRecord :: Bool
        conIsRecord =
          case conVariant of
            NormalConstructor   -> False
            InfixConstructor    -> False
            RecordConstructor{} -> True

        conIsInfix :: Bool
        conIsInfix =
          case conVariant of
            NormalConstructor   -> False
            InfixConstructor    -> True
            RecordConstructor{} -> False

        context :: Type -> Q Type
        context = ntext ''(:=>:) allTvbNames conCtxt

        cocontext :: [Name] -> Cxt -> Type -> Q Type
        cocontext = ntext '(:=>>:)

        ntext :: Name -> [Name] -> Cxt -> Type -> Q Type
        ntext (==>) tvbNames ctxt ty =
          case ctxt of
            [] -> pure ty -- Don't use (:=>:) if there are no constraints
            _  -> infixT (atomizeContext tvbNames ctxt) (==>) (pure ty)

        existentials :: Type -> Q Type
        existentials ty =
          foldl' (\x tvb -> conT ''Exists `appT` pure (tvKind tvb) `appT` x)
                 (pure ty) exTvbs

        selectors :: Q [Type]
        selectors =
          case conVariant of
            NormalConstructor         -> nonRecordCase
            InfixConstructor          -> nonRecordCase
            RecordConstructor records -> recordCase records
          where
            nonRecordCase :: Q [Type]
            nonRecordCase = mkCase (map (const Nothing) fields)

            recordCase :: [Name] -> Q [Type]
            recordCase records = mkCase (map Just records)

            mkCase :: [Maybe Name] -> Q [Type]
            mkCase mbRecords = do
              dcdStricts <- reifyConStrictness conName
              zipWith4M selector mbRecords fieldStricts dcdStricts fields

        selector :: Maybe Name -> FieldStrictness -> TH.DecidedStrictness -> Type -> Q Type
        selector mbRecord (FieldStrictness fu fs) ds field = do
          let mbSelNameT =
                case mbRecord of
                  Just record -> promotedT 'Just `appT` litT (strTyLit (nameBase record))
                  Nothing     -> promotedT 'Nothing
          conT ''S1
            `appT` (promotedT 'MetaSel `appT`
                    mbSelNameT `appT`
                    promoteSourceUnpackedness (generifyUnpackedness fu) `appT`
                    promoteSourceStrictness (generifyStrictness fs) `appT`
                    promoteDecidedStrictness (generifyDecidedStrictness ds))
            `appT` (conT ''Field `appT` prenex allTvbNames field)

        atomizeContext :: [Name] -> Cxt -> Q Type
        atomizeContext tvbNames =
          foldBal (\x y -> infixT x '(:&:) y)
                  (promotedT 'Kon `appT` tupleT 0)
          . map (atomize tvbNames)

#if MIN_VERSION_template_haskell(2,17,0)
        foralls :: [TyVarBndr Specificity] -> Q Type -> Q Type
#else
        foralls :: [TyVarBndr] -> Q Type -> Q Type
#endif
        foralls vs ty =
           foldr (\_ x -> promotedT 'ForAll `appT` x) ty vs

        prenex :: [Name] -> Type -> Q Type
        prenex tvbNames (ForallT vars ctxt ty) =
          let tvbNames' = reverse (map tvName vars) ++ tvbNames in
          (foralls vars . (cocontext tvbNames' ctxt =<<) . prenex tvbNames') ty
        prenex tvbNames ty = atomize tvbNames ty

        atomize :: [Name] -> Type -> Q Type
        atomize tvbNames = flip go []
          where
            -- Collect arguments in a list while descending to the left of AppT,
            -- in case this is a type family application.
            go :: Type -> [Q Type] -> Q Type
            -- Var case
            go ty@(VarT n) =
              case elemIndex n tvbNames of
                Just idx -> appsT $ enumerateTyVar idx
                Nothing  -> kon ty

            -- Either a type constructor or a type family
            go ty@(ConT n)         = \args -> do
              isTFS <- isTypeFamilyOrSynonym n
              if isTFS
              then do (fam, arity) <- promoteNDFamily n
                      (args1, args2) <- splitAt arity <$> sequence args
                      let saturated = all isKonApp args1
                      if saturated then kon ty args
                      else do
                        fcfify n >>= pushFcfifyQueue
                        PromotedT 'Eval
                          `AppT` (PromotedT 'Kon `AppT` fam `appAtom` consTupleAtom args1)
                          `appsT` (pure <$> args2)
              else kon ty args

            -- Kon cases
            go ty@PromotedT{}      = kon ty
            go ty@TupleT{}         = kon ty
            go ty@ArrowT           = kon ty
            go ty@ListT            = kon ty
            go ty@PromotedTupleT{} = kon ty
            go ty@PromotedNilT     = kon ty
            go ty@PromotedConsT    = kon ty
            go ty@StarT            = kon ty
            go ty@ConstraintT      = kon ty
            go ty@LitT{}           = kon ty
            go ty@WildCardT        = kon ty
            go ty@UnboxedTupleT{}  = kon ty
            go ty@UnboxedSumT{}    = kon ty
            go EqualityT           = kon (ConT ''(~~))
                                       -- EqualityT can refer to both homogeneous
                                       -- and heterogeneous equality, but TH always
                                       -- splices EqualityT back in as if it were
                                       -- homogeneous. To be on the safe side, always
                                       -- conservatively assume that the equality it
                                       -- heterogeneous, since it is more permissive.
#if MIN_VERSION_template_haskell(2,17,0)
            go ty@MulArrowT{}      = kon ty
#endif

            -- Recursive cases
            go (AppT ty1 ty2) = go ty1 . (go ty2 [] :)
            go (InfixT ty1 n ty2)  = go (ConT n `AppT` ty1 `AppT` ty2)
            go (UInfixT ty1 n ty2) = go (ConT n `AppT` ty1 `AppT` ty2)
#if MIN_VERSION_template_haskell(2,19,0)
            go (PromotedInfixT  ty1 n ty2) = go (ConT n `AppT` ty1 `AppT` ty2)
            go (PromotedUInfixT ty1 n ty2) = go (ConT n `AppT` ty1 `AppT` ty2)
#endif
            go (SigT ty _)         = go ty
            go (ParensT ty)        = fmap ParensT . go ty
#if MIN_VERSION_template_haskell(2,15,0)
            go (AppKindT ty _)       = go ty
            go (ImplicitParamT n ty) = go (ConT ''IP `AppT` LitT (StrTyLit n) `AppT` ty)
                                         -- Desugar (?n :: T) into (IP "n" T)
#endif

            -- Failure cases
            go ty@ForallT{}       = \_ -> can'tRepresent "rank-n type" ty
#if MIN_VERSION_template_haskell(2,16,0)
            go ty@ForallVisT{}    = \_ -> can'tRepresent "rank-n type" ty
#endif

            kon :: Type -> [Q Type] -> Q Type
            kon ty tys = do ty' <- promotedT 'Kon `appT` pure ty
                            appsT ty' tys

            appsT :: Type -> [Q Type] -> Q Type
            appsT ty1 [] = pure ty1
            appsT ty1 (ty2' : tys) = do ty2 <- ty2'
                                        case (ty1, ty2) of
                                          (PromotedT kon1 `AppT` tyArg1,
                                           PromotedT kon2 `AppT` tyArg2)
                                                 |  kon1 == 'Kon, kon2 == 'Kon
                                                 -> kon (AppT tyArg1 tyArg2) tys
                                          (_, _) -> appsT (ty1 `appAtom` ty2) tys

            can'tRepresent :: String -> Type -> Q a
            can'tRepresent thing ty = fail $ "Unsupported " ++ thing ++ ": " ++ pprint ty

        allTvbNames :: [Name]
        allTvbNames = map tvName exTvbs ++ univVarNames

    fixityIPromotedType :: Maybe TH.Fixity -> Bool -> Q Type
    fixityIPromotedType mbFi True =
               promotedT 'InfixI
        `appT` promoteAssociativity (fdToAssociativity fd)
        `appT` litT (numTyLit (toInteger n))
      where
        Fixity n fd = fromMaybe defaultFixity mbFi
    fixityIPromotedType _ False = promotedT 'PrefixI

isKonApp :: Type -> Bool
isKonApp (PromotedT kon `AppT` _) = kon == 'Kon
isKonApp _ = False

appAtom :: Type -> Type -> Type
appAtom t t' = InfixT t '(:@:) t'

consTupleAtom :: [Type] -> Type
consTupleAtom [] = PromotedT 'Kon `AppT` PromotedT '()
consTupleAtom (t : ts) =
  (PromotedT 'Kon `AppT` PromotedT '(,)) `appAtom` t `appAtom` consTupleAtom ts

deriveFromK :: [ConstructorInfo] -> Q Dec
deriveFromK cons = do
  x <- newName "x"
  funD 'fromK
       [clause [varP x]
               (normalB $ conE 'M1 `appE` caseE (varE x) cases)
               []]
  where
    cases :: [Q Match]
    cases = zipWith (fromCon (length cons)) [1..] cons

    fromCon :: Int -- Total number of constructors
            -> Int -- Constructor index
            -> ConstructorInfo -> Q Match
    fromCon n i ConstructorInfo{ constructorName    = conName
                               , constructorVars    = exTvbs
                               , constructorContext = conCtxt
                               , constructorFields  = fields
                               } = do
      fNames <- newNameList "f" $ length fields
      match (conP conName (map varP fNames))
            (normalB $ lrE i n $ conE 'M1 `appE`
              do prod <- foldBal (\x y -> infixE (Just x) (conE '(:*:)) (Just y))
                           (conE 'U1)
                           (zipWith fromField fNames fields)
                 ctxtProd <- context prod
                 existentials ctxtProd)
            []
      where
        fromField :: Name -> Type -> Q Exp
        fromField fName fty = conE 'M1 `appE` (conE 'Field `appE` prenex fty (varE fName))

        prenex :: Type -> Q Exp -> Q Exp
        prenex (ForallT vars ctxt ty) e =
          foldr (\_ -> appE (conE 'ForAllI)) (cocontext ctxt =<< prenex ty e) vars
        prenex _ e = e

        context :: Exp -> Q Exp
        context = ntext 'SuchThat conCtxt

        cocontext :: Cxt -> Exp -> Q Exp
        cocontext = ntext 'SuchThatI

        ntext :: Name -> Cxt -> Exp -> Q Exp
        ntext suchThat ctxt e =
          case ctxt of
            [] -> pure e
            _  -> conE suchThat `appE` pure e

        existentials :: Exp -> Q Exp
        existentials e = foldl' (\x _ -> conE 'Exists `appE` x) (pure e) exTvbs

deriveToK :: [ConstructorInfo] -> Q Dec
deriveToK cons = do
  x <- newName "x"
  funD 'toK
       [clause [conP 'M1 [varP x]]
               (normalB $ caseE (varE x) cases)
               []]
  where
    cases :: [Q Match]
    cases = zipWith (toCon (length cons)) [1..] cons

    toCon :: Int -- Total number of constructors
          -> Int -- Constructor index
          -> ConstructorInfo -> Q Match
    toCon n i ConstructorInfo{ constructorName    = conName
                             , constructorVars    = exTvbs
                             , constructorContext = conCtxt
                             , constructorFields  = fields
                             } = do
      fNames <- newNameList "f" $ length fields
      match (lrP i n $ conP 'M1
              [ do prod <- foldBal (\x y -> infixP x '(:*:) y)
                                  (conP 'U1 [])
                                  (map (\x -> conP 'M1 [conP 'Field [varP x]]) fNames)
                   ctxtProd <- context prod
                   existentials ctxtProd
              ] )
            (normalB $ foldl' appE (conE conName) (zipWith toField fNames fields))
            []
        where
          toField :: Name -> Type -> Q Exp
          toField fName ty = prenex ty (varE fName)

          prenex :: Type -> Q Exp -> Q Exp
          prenex (ForallT vars ctxt ty) e =
            prenex ty (cocontext ctxt =<< foldl (\x _ -> varE 'unwrapI `appE` (varE 'toWrappedI `appE` x)) e vars)
          prenex _ e = e

          context :: Pat -> Q Pat
          context = ntext (conP 'SuchThat . (:[])) conCtxt

          cocontext :: Cxt -> Exp -> Q Exp
          cocontext = ntext (varE 'unSuchThatI `appE`)

          ntext :: (Q a -> Q a) -> Cxt -> a -> Q a
          ntext suchThat ctxt p =
            case ctxt of
              [] -> pure p
              _  -> suchThat (pure p)

          existentials :: Pat -> Q Pat
          existentials p = foldl' (\x _ -> conP 'Exists [x]) (pure p) exTvbs

-- | If a Type is a SigT, returns its kind signature. Otherwise, return Type.
typeKind :: Type -> Kind
typeKind (SigT _ k) = k
typeKind _          = ConT ''Kind.Type

fdToAssociativity :: FixityDirection -> Associativity
fdToAssociativity InfixL = LeftAssociative
fdToAssociativity InfixR = RightAssociative
fdToAssociativity InfixN = NotAssociative

generifyUnpackedness :: Unpackedness -> Generics.SourceUnpackedness
generifyUnpackedness UnspecifiedUnpackedness = Generics.NoSourceUnpackedness
generifyUnpackedness NoUnpack                = Generics.SourceNoUnpack
generifyUnpackedness Unpack                  = Generics.SourceUnpack

generifyStrictness :: Strictness -> Generics.SourceStrictness
generifyStrictness UnspecifiedStrictness = Generics.NoSourceStrictness
generifyStrictness Lazy                  = Generics.SourceLazy
generifyStrictness THAbs.Strict          = Generics.SourceStrict

generifyDecidedStrictness :: TH.DecidedStrictness -> Generics.DecidedStrictness
generifyDecidedStrictness TH.DecidedLazy   = Generics.DecidedLazy
generifyDecidedStrictness TH.DecidedStrict = Generics.DecidedStrict
generifyDecidedStrictness TH.DecidedUnpack = Generics.DecidedUnpack

promoteSourceUnpackedness :: Generics.SourceUnpackedness -> Q Type
promoteSourceUnpackedness Generics.NoSourceUnpackedness = promotedT 'Generics.NoSourceUnpackedness
promoteSourceUnpackedness Generics.SourceNoUnpack       = promotedT 'Generics.SourceNoUnpack
promoteSourceUnpackedness Generics.SourceUnpack         = promotedT 'Generics.SourceUnpack

promoteSourceStrictness :: Generics.SourceStrictness -> Q Type
promoteSourceStrictness Generics.NoSourceStrictness = promotedT 'Generics.NoSourceStrictness
promoteSourceStrictness Generics.SourceLazy         = promotedT 'Generics.SourceLazy
promoteSourceStrictness Generics.SourceStrict       = promotedT 'Generics.SourceStrict

promoteDecidedStrictness :: Generics.DecidedStrictness -> Q Type
promoteDecidedStrictness Generics.DecidedLazy   = promotedT 'Generics.DecidedLazy
promoteDecidedStrictness Generics.DecidedStrict = promotedT 'Generics.DecidedStrict
promoteDecidedStrictness Generics.DecidedUnpack = promotedT 'Generics.DecidedUnpack

promoteAssociativity :: Associativity -> Q Type
promoteAssociativity LeftAssociative  = promotedT 'LeftAssociative
promoteAssociativity RightAssociative = promotedT 'RightAssociative
promoteAssociativity NotAssociative   = promotedT 'NotAssociative

promoteBool :: Bool -> Q Type
promoteBool True  = promotedT 'True
promoteBool False = promotedT 'False

enumerateTyVar :: Int -> Type
-- Special-case the first 10, if only to generate more compact code
enumerateTyVar 0 = ConT ''Var0
enumerateTyVar 1 = ConT ''Var1
enumerateTyVar 2 = ConT ''Var2
enumerateTyVar 3 = ConT ''Var3
enumerateTyVar 4 = ConT ''Var4
enumerateTyVar 5 = ConT ''Var5
enumerateTyVar 6 = ConT ''Var6
enumerateTyVar 7 = ConT ''Var7
enumerateTyVar 8 = ConT ''Var8
enumerateTyVar 9 = ConT ''Var9
enumerateTyVar n = PromotedT 'Var `AppT` nTimes n (AppT (PromotedT 'VS)) (PromotedT 'VZ)

-- | Variant of foldr for producing balanced lists
foldBal :: (a -> a -> a) -> a -> [a] -> a
foldBal _  x []  = x
foldBal _  _ [y] = y
foldBal op x l   = let (a,b) = splitAt (length l `div` 2) l
                   in foldBal op x a `op` foldBal op x b

lrP :: Int -- Constructor index
    -> Int -- Total number of constructors
    -> Q Pat -> Q Pat
lrP i n p
  | n == 0       = fail "lrP: impossible"
  | n == 1       = p
  | i <= div n 2 = conP 'L1 [lrP i     (div n 2) p]
  | otherwise    = conP 'R1 [lrP (i-m) (n-m)     p]
                     where m = div n 2

lrE :: Int -- Constructor index
    -> Int -- Total number of constructors
    -> Q Exp -> Q Exp
lrE i n e
  | n == 0       = fail "lrE: impossible"
  | n == 1       = e
  | i <= div n 2 = conE 'L1 `appE` lrE i     (div n 2) e
  | otherwise    = conE 'R1 `appE` lrE (i-m) (n-m)     e
                     where m = div n 2

isNewtypeVariant :: DatatypeVariant -> Bool
isNewtypeVariant Datatype        = False
isNewtypeVariant Newtype         = True
isNewtypeVariant DataInstance    = False
isNewtypeVariant NewtypeInstance = True

-- | Extract 'Just' the 'Name' from a type variable. If the argument 'Type' is
-- not a type variable, return 'Nothing'.
varTToName_maybe :: Type -> Maybe Name
varTToName_maybe (VarT n)   = Just n
varTToName_maybe (SigT t _) = varTToName_maybe t
varTToName_maybe _          = Nothing

-- | Extract the 'Name' from a type variable. If the argument 'Type' is not a
-- type variable, throw an error.
varTToName :: Type -> Name
varTToName = fromMaybe (error "Not a type variable!") . varTToName_maybe

zipWith4M :: Monad m => (a -> b -> c -> d -> m e)
          -> [a] -> [b] -> [c] -> [d] -> m [e]
zipWith4M _ []     _      _      _      = pure []
zipWith4M _ _      []     _      _      = pure []
zipWith4M _ _      _      []     _      = pure []
zipWith4M _ _      _      _      []     = pure []
zipWith4M f (x:xs) (y:ys) (z:zs) (a:as)
  = do r  <- f x y z a
       rs <- zipWith4M f xs ys zs as
       pure $ r:rs

-- | Compose a function with itself n times.  (nth rather than twice)
nTimes :: Int -> (a -> a) -> (a -> a)
nTimes 0 _ = id
nTimes 1 f = f
nTimes n f = f . nTimes (n-1) f

-- | Generate a list of fresh names with a common prefix, and numbered suffixes.
newNameList :: String -> Int -> Q [Name]
newNameList prefix n = traverse (newName . (prefix ++) . show) [1..n]

gatherExistentials :: [ConstructorInfo] -> [TyVarBndrUnit]
gatherExistentials = concatMap constructorVars

gatherConstraints :: [ConstructorInfo] -> [Pred]
gatherConstraints = concatMap constructorContext

gatherFields :: [ConstructorInfo] -> [Type]
gatherFields = concatMap constructorFields

-- | Detect if a name occurs as an argument to some type family.
isInTypeFamilyApp :: Name -> Type -> Q Bool
isInTypeFamilyApp name = go
  where
    go :: Type -> Q Bool
    go ty@AppT{}          = case splitAppTs ty of
                              (tyFun, tyArgs)
                                |  ConT tcName <- tyFun
                                -> goTyConApp tcName tyArgs
                                |  otherwise
                                -> or <$> traverse go (tyFun:tyArgs)
    go (InfixT ty1 n ty2) = go (ConT n `AppT` ty1 `AppT` ty2)
    go (SigT ty ki)       = liftA2 (||) (go ty) (go ki)
    go (ParensT ty)       = go ty
    go _                  = pure False

    goTyConApp :: Name -> [Type] -> Q Bool
    goTyConApp tcName tcArgs = do
      info <- reify tcName
      case info of
        FamilyI (OpenTypeFamilyD (TypeFamilyHead _ bndrs _ _)) _
          -> withinFirstArgs bndrs
        FamilyI (ClosedTypeFamilyD (TypeFamilyHead _ bndrs _ _) _) _
          -> withinFirstArgs bndrs
        _ -> pure False
      where
        withinFirstArgs :: [a] -> Q Bool
        withinFirstArgs bndrs =
          let firstArgs = take (length bndrs) tcArgs
          in pure $ name `elem` freeVariables firstArgs

-- | Split a chain of 'AppT's to a linear chain of arguments.
splitAppTs :: Type -> (Type, [Type])
splitAppTs ty = split ty ty []
  where
    split :: Type -> Type -> [Type] -> (Type, [Type])
    split _      (AppT ty1 ty2)     args = split ty1 ty1 (ty2:args)
    split origTy (InfixT ty1 n ty2) args = split origTy (ConT n `AppT` ty1 `AppT` ty2) args
    split origTy (SigT ty' _)       args = split origTy ty' args
    split origTy (ParensT ty')      args = split origTy ty' args
    split origTy _                  args = (origTy, args)

-- | Resolve all of the type synonyms in a 'ConstructorInfo'.
resolveConSynonyms :: ConstructorInfo -> Q ConstructorInfo
resolveConSynonyms con@ConstructorInfo{ constructorVars    = vars
                                      , constructorContext = context
                                      , constructorFields  = fields
                                      } = do
  vars'    <- traverse (\tvb ->
                         elimTV (\_n -> pure tvb)
                                (\n k -> kindedTV n <$> resolveTypeSynonyms k) tvb) vars
  context' <- traverse resolveTypeSynonyms context
  fields'  <- traverse resolveTypeSynonyms fields
  pure $ con{ constructorVars = vars'
            , constructorContext = context'
            , constructorFields  = fields'
            }

-- | Store 'GenericK' instances to be produced after having typechecked
-- 'fcfify'-generated instances.
newtype GenericKQueue = GenericKQueue [Dec]

pushGenericKQueue :: [Dec] -> Q ()
pushGenericKQueue d = do
  GenericKQueue decs <- fromMaybe (GenericKQueue []) <$> TH.getQ
  TH.putQ (GenericKQueue (d ++ decs))

takeGenericKQueue :: Q [Dec]
takeGenericKQueue = do
  GenericKQueue decs <- fromMaybe (GenericKQueue []) <$> TH.getQ
  TH.putQ (GenericKQueue [])
  pure decs

-- | Store 'fcfify'-generated instances for the current data type.
-- This could also be done with StateT in deriveRepK but that's
-- a more invasive change.
newtype FcfifyQueue = FcfifyQueue [Dec]

pushFcfifyQueue :: [Dec] -> Q ()
pushFcfifyQueue d = do
  FcfifyQueue decs <- fromMaybe (FcfifyQueue []) <$> TH.getQ
  TH.putQ (FcfifyQueue (d ++ decs))

takeFcfifyQueue :: Q [Dec]
takeFcfifyQueue = do
  FcfifyQueue decs <- fromMaybe (FcfifyQueue []) <$> TH.getQ
  TH.putQ (FcfifyQueue [])
  pure decs