content: {-# LANGUAGE CPP #-}
  {-# LANGUAGE DataKinds #-}
  {-# LANGUAGE LambdaCase #-}
  {-# LANGUAGE RecordWildCards #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE TemplateHaskell #-}
  {-# LANGUAGE TypeOperators #-}
  
  -- |
  -- Module      :   Grisette.Internal.TH.Derivation.Common
  -- 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.Derivation.Common
    ( CheckArgsResult (..),
      checkArgs,
      ctxForVar,
      EvalModeConfig (..),
      DeriveConfig (..),
      extraEvalModeConstraint,
      extraBitSizeConstraint,
      extraFpBitSizeConstraint,
      extraExtraMergeableConstraint,
      extraConstraint,
      specializeResult,
      evalModeSpecializeList,
      isVarUsedInFields,
      freshenCheckArgsResult,
    )
  where
  
  import Control.Monad (foldM, unless, when, zipWithM)
  import Data.Bifunctor (first)
  import qualified Data.Map as M
  import Data.Maybe (catMaybes, mapMaybe)
  import qualified Data.Set as S
  import GHC.TypeLits (KnownNat, Nat, type (<=))
  import Grisette.Internal.Internal.Decl.Core.Data.Class.Mergeable
    ( Mergeable,
      Mergeable1,
      Mergeable2,
    )
  import Grisette.Internal.SymPrim.FP (ValidFP)
  import Grisette.Internal.Unified.EvalModeTag (EvalModeTag (C, S))
  import Grisette.Internal.Unified.Util (DecideEvalMode)
  import Language.Haskell.TH
    ( Kind,
      Name,
      Pred,
      Q,
      Type (AppT, ArrowT, ConT, PromotedT, StarT, VarT),
      conT,
      nameBase,
      newName,
    )
  import Language.Haskell.TH.Datatype
    ( ConstructorInfo (constructorFields, constructorName, constructorVars),
      DatatypeInfo (datatypeCons, datatypeVars),
      TypeSubstitution (applySubstitution, freeVariables),
      reifyDatatype,
      tvName,
    )
  import Language.Haskell.TH.Datatype.TyVarBndr (mapTVName, tvKind)
  
  -- | Result of 'checkArgs' for a data type.
  data CheckArgsResult = CheckArgsResult
    { constructors :: [ConstructorInfo],
      keptVars :: [(Type, Kind)],
      argVars :: [(Type, Kind)]
    }
  
  -- | Specialize the evaluation mode tags for the t'CheckArgsResult'.
  specializeResult :: [(Int, EvalModeTag)] -> CheckArgsResult -> Q CheckArgsResult
  specializeResult evalModeConfigs result = do
    let modeToName C = 'C
        modeToName S = 'S
    map <-
      foldM
        ( \lst (n, tag) -> do
            let (_, knd) = lst !! n
            return $
              take n lst
                ++ [(PromotedT $ modeToName tag, knd)]
                ++ drop (n + 1) lst
        )
        (keptVars result)
        evalModeConfigs
    return $ result {keptVars = map}
  
  freshenConstructorInfo :: ConstructorInfo -> Q ConstructorInfo
  freshenConstructorInfo conInfo = do
    let vars = constructorVars conInfo
    newNames <- traverse (newName . nameBase . tvName) vars
    let newVars = zipWith (mapTVName . const) newNames vars
    let substMap = M.fromList $ zip (tvName <$> vars) $ VarT <$> newNames
    return $ applySubstitution substMap conInfo {constructorVars = newVars}
  
  -- | Freshen the type variables in the t'CheckArgsResult'.
  freshenCheckArgsResult :: Bool -> CheckArgsResult -> Q CheckArgsResult
  freshenCheckArgsResult freshenNats result = do
    let genNewName :: (Type, Kind) -> Q (Maybe Name)
        genNewName (VarT _, knd) =
          if not freshenNats && knd == ConT ''Nat
            then return Nothing
            else Just <$> newName "a"
        genNewName _ = return Nothing
    keptNewNames <- traverse genNewName (keptVars result)
    argNewNames <- traverse genNewName (argVars result)
  
    let substMap =
          M.fromList
            $ mapMaybe
              ( \(newName, oldVar) ->
                  case (newName, oldVar) of
                    (Just newName, (VarT oldName, _)) ->
                      Just (oldName, VarT newName)
                    _ -> Nothing
              )
            $ zip
              (keptNewNames ++ argNewNames)
              (keptVars result ++ argVars result)
    constructors <-
      mapM freshenConstructorInfo $
        applySubstitution substMap $
          constructors result
    let newKeptVars = first (applySubstitution substMap) <$> (keptVars result)
    let newArgVars = first (applySubstitution substMap) <$> (argVars result)
    return $
      result
        { constructors = constructors,
          keptVars = newKeptVars,
          argVars = newArgVars
        }
  
  -- | Check if the number of type parameters is valid for a data type, and return
  -- new names for the type variables, split into kept and arg parts.
  checkArgs ::
    String ->
    Int ->
    Name ->
    Bool ->
    Int ->
    Q CheckArgsResult
  checkArgs clsName maxArgNum typName allowExistential n = do
    when (n < 0) $
      fail $
        unlines
          [ "Cannot derive "
              ++ clsName
              ++ " instance with negative type parameters",
            "\tRequested: " ++ show n,
            "\tHint: Use a non-negative number of type parameters"
          ]
    when (n > maxArgNum) $
      fail $
        unlines
          [ "Cannot derive "
              <> clsName
              <> " instance with more than "
              <> show maxArgNum
              <> " type parameters",
            "\tRequested: " <> show n
          ]
    d <- reifyDatatype typName
    let dvars = datatypeVars d
    when (length dvars < n) $
      fail $
        unlines
          [ "Cannot derive "
              <> clsName
              <> show n
              <> " instance for the type "
              <> show typName,
            "\tReason: The type "
              <> show typName
              <> " has only "
              <> show (length dvars)
              <> " type variables."
          ]
    let keptVars =
          (\bndr -> (VarT $ tvName bndr, tvKind bndr))
            <$> take (length dvars - n) dvars
    let argVars =
          (\bndr -> (VarT $ tvName bndr, tvKind bndr))
            <$> drop (length dvars - n) dvars
    let constructors = datatypeCons d
    unless allowExistential $
      mapM_
        ( \c ->
            when (constructorVars c /= []) $
              fail $
                unlines
                  [ "Cannot derive "
                      <> clsName
                      <> show n
                      <> " instance for the type "
                      <> show typName,
                    "\tReason: The constructor "
                      <> nameBase (constructorName c)
                      <> " has existential variables"
                  ]
        )
        constructors
    mapM_
      ( \c -> do
          let fields = constructorFields c
          let existentialVars = tvName <$> constructorVars c
          let fieldReferencedVars = freeVariables fields
          let notReferencedVars =
                S.fromList existentialVars S.\\ S.fromList fieldReferencedVars
          unless (null notReferencedVars) $
            fail $
              unlines
                [ "Cannot derive "
                    <> clsName
                    <> show n
                    <> " instance for the type "
                    <> show typName,
                  "Reason: Ambiguous existential variable in the constructor: "
                    <> nameBase (constructorName c)
                    <> ", this is not supported. Please consider binding the "
                    <> "existential variable to a field. You can use Proxy type to "
                    <> "do this."
                ]
      )
      constructors
    return $ CheckArgsResult {..}
  
  isVarUsedInConstructorFields :: [ConstructorInfo] -> Name -> Bool
  isVarUsedInConstructorFields constructors var =
    let allFields = concatMap constructorFields constructors
        allFieldsFreeVars = S.fromList $ freeVariables allFields
     in S.member var allFieldsFreeVars
  
  -- | Check if a variable is used in the fields of a constructor.
  isVarUsedInFields :: CheckArgsResult -> Name -> Bool
  isVarUsedInFields CheckArgsResult {..} =
    isVarUsedInConstructorFields constructors
  
  -- | Generate a context for a variable in a GADT.
  ctxForVar :: [Type] -> Type -> Kind -> Q (Maybe Pred)
  ctxForVar instanceExps ty knd = case knd of
    StarT ->
      Just
        <$> [t|$(return $ head instanceExps) $(return ty)|]
    AppT (AppT ArrowT StarT) StarT ->
      Just
        <$> [t|$(return $ instanceExps !! 1) $(return ty)|]
    AppT (AppT (AppT ArrowT StarT) StarT) StarT ->
      Just
        <$> [t|$(return $ instanceExps !! 2) $(return ty)|]
    AppT (AppT (AppT (AppT ArrowT StarT) StarT) StarT) StarT ->
      Just
        <$> [t|$(return $ instanceExps !! 3) $(return ty)|]
    AppT (AppT (AppT (AppT ArrowT StarT) StarT) StarT) _ ->
      fail $ "Unsupported kind: " <> show knd
    _ -> return Nothing
  
  -- | Configuration for constraints for evaluation modes tag.
  --
  -- * 'EvalModeConstraints' specifies a list of constraints for the tag, for
  --   example, we may use 'Grisette.Unified.EvalModeBase' and
  --   'Grisette.Unified.EvalModeBV' to specify that the evaluation mode must
  --   support both base (boolean and data types) and bit vectors. This should be
  --   used when the data type uses bit vectors.
  --
  -- * 'EvalModeSpecified' specifies a that an evaluation mode tag should be
  --   specialized to a specific tag for all the instances.
  data EvalModeConfig
    = EvalModeConstraints [Name]
    | EvalModeSpecified EvalModeTag
  
  -- | Configuration for deriving instances for a data type.
  data DeriveConfig = DeriveConfig
    { evalModeConfig :: [(Int, EvalModeConfig)],
      bitSizePositions :: [Int],
      fpBitSizePositions :: [(Int, Int)],
      unconstrainedPositions :: [Int],
      needExtraMergeableUnderEvalMode :: Bool,
      needExtraMergeableWithConcretizedEvalMode :: Bool,
      useNoStrategy :: Bool,
      useSerialForCerealAndBinary :: Bool
    }
  
  -- | Get all the evaluation modes to specialize in the t'DeriveConfig'.
  evalModeSpecializeList :: DeriveConfig -> [(Int, EvalModeTag)]
  evalModeSpecializeList DeriveConfig {..} =
    mapMaybe
      ( \(n, cfg) ->
          case cfg of
            EvalModeConstraints _ -> Nothing
            EvalModeSpecified tag -> Just (n, tag)
      )
      evalModeConfig
  
  instance Semigroup DeriveConfig where
    l <> r =
      DeriveConfig
        { evalModeConfig = evalModeConfig l <> evalModeConfig r,
          bitSizePositions = bitSizePositions l <> bitSizePositions r,
          fpBitSizePositions = fpBitSizePositions l <> fpBitSizePositions r,
          unconstrainedPositions = unconstrainedPositions l <> unconstrainedPositions r,
          needExtraMergeableUnderEvalMode =
            needExtraMergeableUnderEvalMode l
              || needExtraMergeableUnderEvalMode r,
          needExtraMergeableWithConcretizedEvalMode =
            needExtraMergeableWithConcretizedEvalMode l
              || needExtraMergeableWithConcretizedEvalMode r,
          useNoStrategy = useNoStrategy l || useNoStrategy r,
          useSerialForCerealAndBinary =
            useSerialForCerealAndBinary l && useSerialForCerealAndBinary r
        }
  
  instance Monoid DeriveConfig where
    mempty = DeriveConfig [] [] [] [] False False False True
    mappend = (<>)
  
  -- | Generate extra constraints for evaluation modes.
  extraEvalModeConstraint ::
    Name -> Name -> [(Type, Kind)] -> (Int, EvalModeConfig) -> Q [Pred]
  extraEvalModeConstraint
    tyName
    instanceName
    args
    (n, EvalModeConstraints names)
      | n >= length args = return []
      | otherwise = do
          let (arg, argKind) = args !! n
          when (argKind /= ConT ''EvalModeTag) $
            fail $
              "Cannot introduce EvalMode constraint for the "
                <> show n
                <> "th argument of "
                <> show tyName
                <> " when deriving the "
                <> show instanceName
                <> " instance because it is not an EvalModeTag."
          traverse (\nm -> [t|$(conT nm) $(return arg)|]) names
  extraEvalModeConstraint _ _ _ (_, EvalModeSpecified _) = return []
  
  -- | Generate extra constraints for bit vectors.
  extraBitSizeConstraint :: Name -> Name -> [(Type, Kind)] -> Int -> Q [Pred]
  extraBitSizeConstraint tyName instanceName args n
    | n >= length args = return []
    | otherwise = do
        let (arg, argKind) = args !! n
        when (argKind /= ConT ''Nat) $
          fail $
            "Cannot introduce BitSize constraint for the "
              <> show n
              <> "th argument of "
              <> show tyName
              <> " when deriving the "
              <> show instanceName
              <> " instance because it is not a Nat."
        predKnown <- [t|KnownNat $(return arg)|]
        predPositive <- [t|1 <= $(return arg)|]
        return [predKnown, predPositive]
  
  -- | Generate extra constraints for floating point exponents and significands.
  extraFpBitSizeConstraint ::
    Name -> Name -> [(Type, Kind)] -> (Int, Int) -> Q [Pred]
  extraFpBitSizeConstraint tyName instanceName args (eb, sb)
    | eb >= length args || sb >= length args = return []
    | otherwise = do
        let (argEb, argEbKind) = args !! eb
        let (argSb, argSbKind) = args !! sb
        when (argEbKind /= ConT ''Nat || argSbKind /= ConT ''Nat) $
          fail $
            "Cannot introduce ValidFP constraint for the "
              <> show eb
              <> "th and "
              <> show sb
              <> "th arguments of "
              <> show tyName
              <> " when deriving the "
              <> show instanceName
              <> " instance because they are not Nats."
        pred <- [t|ValidFP $(return argEb) $(return argSb)|]
        return [pred]
  
  -- | Generate extra constraints for 'Mergeable' instances.
  extraExtraMergeableConstraint ::
    DeriveConfig -> [ConstructorInfo] -> [(Type, Kind)] -> Q [Pred]
  extraExtraMergeableConstraint deriveConfig constructors args = do
    let isTypeUsedInFields' (VarT nm) =
          isVarUsedInConstructorFields constructors nm
        isTypeUsedInFields' _ = False
    catMaybes
      <$> zipWithM
        ( \position (arg, knd) ->
            if isTypeUsedInFields' arg
              && notElem position (unconstrainedPositions deriveConfig)
              then
                ctxForVar
                  [ ConT ''Mergeable,
                    ConT ''Mergeable1,
                    ConT ''Mergeable2
                  ]
                  arg
                  knd
              else return Nothing
        )
        [0 ..]
        args
  
  -- | Generate extra constraints for a data type.
  extraConstraint ::
    DeriveConfig ->
    Name ->
    Name ->
    [(Type, Kind)] ->
    [(Type, Kind)] ->
    [ConstructorInfo] ->
    Q [Pred]
  extraConstraint
    deriveConfig@DeriveConfig {..}
    tyName
    instanceName
    extraArgs
    keptArgs
    constructors = do
      evalModePreds <-
        traverse
          (extraEvalModeConstraint tyName instanceName keptArgs)
          evalModeConfig
      extraArgEvalModePreds <-
        if null evalModeConfig
          then
            traverse
              ( \(arg, kind) ->
                  if kind == ConT ''EvalModeTag
                    then (: []) <$> [t|DecideEvalMode $(return arg)|]
                    else return []
              )
              extraArgs
          else return []
      bitSizePreds <-
        traverse
          (extraBitSizeConstraint tyName instanceName keptArgs)
          bitSizePositions
      fpBitSizePreds <-
        traverse
          (extraFpBitSizeConstraint tyName instanceName keptArgs)
          fpBitSizePositions
      extraMergeablePreds <-
        if needExtraMergeableUnderEvalMode
          && ( any
                 ( \case
                     (_, EvalModeConstraints _) -> True
                     (_, EvalModeSpecified _) -> False
                 )
                 evalModeConfig
                 || needExtraMergeableWithConcretizedEvalMode
             )
          then extraExtraMergeableConstraint deriveConfig constructors keptArgs
          else return []
      return $
        concat (extraArgEvalModePreds ++ evalModePreds)
          ++ if null constructors
            then []
            else extraMergeablePreds ++ concat (bitSizePreds ++ fpBitSizePreds)
  