content: {-# LANGUAGE CPP #-}
  {-# LANGUAGE TemplateHaskell #-}
  {-# OPTIONS_GHC -Wno-unrecognised-pragmas #-}
  
  {-# HLINT ignore "Unused LANGUAGE pragma" #-}
  
  -- |
  -- Module      :   Grisette.Internal.TH.Util
  -- Copyright   :   (c) Sirui Lu 2024
  -- License     :   BSD-3-Clause (see the LICENSE file)
  --
  -- Maintainer  :   siruilu@cs.washington.edu
  -- Stability   :   Experimental
  -- Portability :   GHC only
  module Grisette.Internal.TH.Util
    ( constructorInfoToType,
      tvIsMode,
      tvIsNat,
      tvIsStar,
      tvIsStarToStar,
      substDataType,
      reifyDatatypeWithFreshNames,
      singleParamClassParamKind,
      binaryClassParamKind,
      getTypeWithMaybeSubst,
      dropLastTypeParam,
      dropNTypeParam,
      classParamKinds,
      allSameKind,
      classNumParam,
      kindNumParam,
      concatPreds,
      putHaddock,
      allUsedNamesMaybe,
      allUsedNames,
      isNonUnitTupleString,
      isNonUnitTuple,
      integerE,
      mangleName,
      dataTypeHasExistential,
    )
  where
  
  #if MIN_VERSION_template_haskell(2,18,0)
  import Language.Haskell.TH.Syntax
    ( DocLoc (DeclDoc),
      ModName (ModName),
      Name (Name),
      NameFlavour (NameG, NameQ, NameS),
      addModFinalizer,
      putDoc,
    )
  #else
  import Language.Haskell.TH.Syntax
    ( ModName (ModName),
      Name (Name),
      NameFlavour (NameG, NameQ, NameS),
    )
  #endif
  
  import Control.Monad (when)
  import Data.Char (isAlphaNum, ord)
  import qualified Data.Map as M
  import qualified Data.Set as S
  import GHC.TypeNats (Nat)
  import Grisette.Internal.Unified.EvalModeTag (EvalModeTag)
  import Language.Haskell.TH
    ( Dec (ClassD),
      Exp
        ( AppE,
          AppTypeE,
          ConE,
          CondE,
          InfixE,
          LamE,
          ListE,
          LitE,
          ParensE,
          SigE,
          TupE,
          UInfixE,
          VarE
        ),
      Info (ClassI),
      Kind,
      Pred,
      Q,
      Type (AppT, ArrowT, ConT, ForallT, StarT, VarT),
      integerL,
      litE,
      nameBase,
      newName,
      pprint,
      reify,
      varT,
    )
  import Language.Haskell.TH.Datatype
    ( ConstructorInfo (constructorContext, constructorFields, constructorVars),
      DatatypeInfo (datatypeCons, datatypeInstTypes, datatypeVars),
      TypeSubstitution (applySubstitution),
      datatypeType,
      reifyDatatype,
      tvName,
    )
  import Language.Haskell.TH.Datatype.TyVarBndr
    ( Specificity (SpecifiedSpec),
      TyVarBndrUnit,
      TyVarBndr_,
      mapTVFlag,
      mapTVName,
      tvKind,
    )
  
  -- | Convert a 'ConstructorInfo' to a 'Type' of the constructor.
  constructorInfoToType :: DatatypeInfo -> ConstructorInfo -> Q Type
  constructorInfoToType dataType info = do
    let binders =
          mapTVFlag (const SpecifiedSpec)
            <$> datatypeVars dataType ++ constructorVars info
    let ctx = constructorContext info
    let fields = constructorFields info
    let tyBody =
          foldr (AppT . AppT ArrowT) (datatypeType dataType) fields
    if null binders then return tyBody else return $ ForallT binders ctx tyBody
  
  -- | Check if a type variable is of kind 'EvalModeTag'.
  tvIsMode :: TyVarBndr_ flag -> Bool
  tvIsMode = (== ConT ''EvalModeTag) . tvKind
  
  -- | Check if a type variable is of kind 'Nat'.
  tvIsNat :: TyVarBndr_ flag -> Bool
  tvIsNat = (== ConT ''Nat) . tvKind
  
  -- | Check if a type variable is of kind 'Data.Kind.Type'.
  tvIsStar :: TyVarBndr_ flag -> Bool
  tvIsStar = (== StarT) . tvKind
  
  -- | Check if a type variable is of kind 'Data.Kind.Type -> Data.Kind.Type'.
  tvIsStarToStar :: TyVarBndr_ flag -> Bool
  tvIsStarToStar = (== (AppT (AppT ArrowT StarT) StarT)) . tvKind
  
  -- | Substitute the type variables in a 'DatatypeInfo' with the given
  -- substitution map.
  substDataType :: DatatypeInfo -> M.Map Name Type -> DatatypeInfo
  substDataType d substMap =
    d
      { datatypeInstTypes = applySubstitution substMap <$> datatypeInstTypes d,
        datatypeCons = applySubstitution substMap <$> datatypeCons d
      }
  
  -- | Convert a 'DatatypeInfo' to a 'DatatypeInfo' with fresh type variable
  -- names.
  datatypeToFreshNames :: DatatypeInfo -> Q DatatypeInfo
  datatypeToFreshNames d = do
    let vars = datatypeVars d
    let names = tvName <$> vars
    freshNames <- traverse (newName . show) names
    let newDTVars = zipWith (\v n -> mapTVName (const n) v) vars freshNames
    let substMap = M.fromList $ zip names (VarT <$> freshNames)
    return $ substDataType d {datatypeVars = newDTVars} substMap
  
  -- | Reify a datatype with fresh type variable names.
  reifyDatatypeWithFreshNames :: Name -> Q DatatypeInfo
  reifyDatatypeWithFreshNames name = do
    d <- reifyDatatype name
    datatypeToFreshNames d
  
  -- | Check if all type variables have the same kind.
  allSameKind :: [TyVarBndrUnit] -> Bool
  allSameKind [] = True
  allSameKind (x : xs) = all ((== tvKind x) . tvKind) xs
  
  -- | Get the kinds of the type parameters of a class.
  classParamKinds :: Name -> Q [Kind]
  classParamKinds className = do
    cls <- reify className
    case cls of
      ClassI (ClassD _ _ bndrs _ _) _ -> return $ tvKind <$> bndrs
      _ ->
        fail $
          "symmetricClassParamKind:" <> show className <> " is not a class"
  
  -- | Get the number of type parameters of a class.
  classNumParam :: Name -> Q Int
  classNumParam className = do
    cls <- reify className
    case cls of
      ClassI (ClassD _ _ bndrs _ _) _ -> return $ length bndrs
      _ ->
        fail $
          "classNumParam:" <> show className <> " is not a class"
  
  -- | Get the kind of the single type parameter of a class.
  singleParamClassParamKind :: Name -> Q Kind
  singleParamClassParamKind className = do
    cls <- reify className
    case cls of
      ClassI (ClassD _ _ bndrs _ _) _ ->
        case bndrs of
          [x] -> return $ tvKind x
          _ ->
            fail $
              "singleParamClassParamKind: only support classes with one type "
                <> "parameter, but "
                <> show className
                <> " has "
                <> show (length bndrs)
      _ ->
        fail $
          "singleParamClassParamKind:" <> show className <> " is not a class"
  
  -- | Get the kind of the binary type parameter of a class.
  binaryClassParamKind :: Name -> Q Kind
  binaryClassParamKind className = do
    cls <- reify className
    case cls of
      ClassI (ClassD _ _ bndrs _ _) _ ->
        case bndrs of
          [x, y] -> do
            when (tvKind x /= tvKind y) $
              fail "binaryClassParamKind: type parameters have different kinds"
            return $ tvKind x
          _ ->
            fail $
              "binaryClassParamKind: only support classes with two type "
                <> "parameters, but "
                <> show className
                <> " has "
                <> show (length bndrs)
      _ ->
        fail $
          "binaryClassParamKind:" <> show className <> " is not a class"
  
  -- | Get a type with a possible substitution.
  getTypeWithMaybeSubst :: TyVarBndrUnit -> Maybe Type -> Q Type
  getTypeWithMaybeSubst tv Nothing = varT $ tvName tv
  getTypeWithMaybeSubst _ (Just t) = return t
  
  -- | Drop the last instantiated type parameter of a type.
  dropLastTypeParam :: Type -> Q Type
  dropLastTypeParam (AppT c _) = return c
  dropLastTypeParam v =
    fail $
      "dropLastTypeParam: have no type parameters: "
        <> pprint v
        <> " / "
        <> show v
  
  -- | Drop the last N instantiated type parameters of a type.
  dropNTypeParam :: Int -> Type -> Q Type
  dropNTypeParam 0 t = return t
  dropNTypeParam n t = dropLastTypeParam t >>= dropNTypeParam (n - 1)
  
  -- | Get the number of type parameters of a kind.
  kindNumParam :: Kind -> Q Int
  kindNumParam (AppT (AppT ArrowT _) k) = (1 +) <$> kindNumParam k
  kindNumParam _ = return 0
  
  -- | Concatenate two 'Maybe [Pred]'.
  concatPreds :: Maybe [Pred] -> Maybe [Pred] -> Maybe [Pred]
  concatPreds Nothing Nothing = Nothing
  concatPreds (Just ps) Nothing = Just ps
  concatPreds Nothing (Just ps) = Just ps
  concatPreds (Just ps1) (Just ps2) = Just $ ps1 ++ ps2
  
  #if MIN_VERSION_template_haskell(2,18,0)
  -- | Put a haddock comment on a declaration.
  putHaddock :: Name -> String -> Q ()
  putHaddock name = addModFinalizer . putDoc (DeclDoc name) 
  #else
  -- | Put a haddock comment on a declaration.
  -- (No-op because compiling with GHC < 9.2)
  putHaddock :: Name -> String -> Q ()
  putHaddock _ _ = return ()
  #endif
  
  -- | Get the names used in an expression.
  allUsedNamesMaybe :: Maybe Exp -> S.Set Name
  allUsedNamesMaybe Nothing = S.empty
  allUsedNamesMaybe (Just exp) = allUsedNames exp
  
  -- | Get the names used in an expression.
  allUsedNames :: Exp -> S.Set Name
  allUsedNames (VarE nm) = S.singleton nm
  allUsedNames (ConE n) = S.singleton n
  allUsedNames (LitE _) = S.empty
  allUsedNames (AppE e1 e2) = allUsedNames e1 `S.union` allUsedNames e2
  allUsedNames (AppTypeE e1 _) = allUsedNames e1
  allUsedNames (InfixE l e r) =
    allUsedNamesMaybe l `S.union` allUsedNames e `S.union` allUsedNamesMaybe r
  allUsedNames (UInfixE l e r) =
    allUsedNames l `S.union` allUsedNames e `S.union` allUsedNames r
  allUsedNames (ParensE e) = allUsedNames e
  allUsedNames (LamE _ e) = allUsedNames e
  allUsedNames (TupE es) = mconcat $ allUsedNamesMaybe <$> es
  allUsedNames (CondE e1 e2 e3) =
    allUsedNames e1 `S.union` allUsedNames e2 `S.union` allUsedNames e3
  allUsedNames (ListE es) = mconcat $ allUsedNames <$> es
  allUsedNames (SigE e _) = allUsedNames e
  allUsedNames exp = error $ "allUsedNames: unsupported expression: " <> show exp
  
  -- | Check if a string is the data constructor name of a non-unit tuple.
  isNonUnitTupleString :: String -> Bool
  isNonUnitTupleString ('(' : ',' : _) = True
  isNonUnitTupleString _ = False
  
  -- | Check if a name is the data constructor name of a non-unit tuple.
  isNonUnitTuple :: Name -> Bool
  isNonUnitTuple nm =
    isNonUnitTupleString $ nameBase nm
  
  -- | Convert an integer to an 'Exp'.
  integerE :: (Integral a) => a -> Q Exp
  integerE = litE . integerL . fromIntegral
  
  -- | Mangle a name string to contain only alphanumeric characters and
  -- underscores.
  mangleName :: Name -> String
  mangleName nm@(Name _ flavor) =
    case flavor of
      NameS -> mangleBaseName $ nameBase nm
      NameQ mod -> mangleModName mod <> "_" <> mangleBaseName (nameBase nm)
      NameG _ _ mod -> mangleModName mod <> "_" <> mangleBaseName (nameBase nm)
      _ -> error $ "mangleName: unsupported name flavor: " <> show flavor
    where
      mangleModName (ModName m) = mangleBaseName m
      mangleBaseName l = "Mangled" ++ go l
      go [] = []
      go (c : cs)
        | isAlphaNum c || c == '_' = c : go cs
        | otherwise = "_" <> show (ord c) <> go cs
  
  -- | Check if a data type has existential variables in constructors.
  dataTypeHasExistential :: Name -> Q Bool
  dataTypeHasExistential typName = do
    d <- reifyDatatype typName
    return $ not $ all (null . constructorVars) $ datatypeCons d
  