grisette-0.11.0.0: src/Grisette/Internal/Unified/UnifiedFun.hs

{-# LANGUAGE CPP #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE QuantifiedConstraints #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilyDependencies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}

-- |
-- Module      :   Grisette.Internal.Unified.UnifiedFun
-- 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.Unified.UnifiedFun
  ( UnifiedFunConstraint,
    UnifiedFun (..),
    unifiedFunInstanceName,
    genUnifiedFunInstance,
    GetFun2,
    GetFun3,
    GetFun4,
    GetFun5,
    GetFun6,
    GetFun7,
    GetFun8,
  )
where

#if MIN_VERSION_base(4,20,0)
#else
import Data.Foldable (Foldable (foldl'))
#endif

import Control.DeepSeq (NFData)
import Data.Binary (Binary)
import Data.Bytes.Serial (Serial)
import Data.Hashable (Hashable)
import qualified Data.Kind
import Data.Serialize (Serialize)
import Data.Typeable (Typeable)
import GHC.TypeLits (KnownNat, Nat, type (<=))
import Grisette.Internal.Core.Data.Class.EvalSym (EvalSym)
import Grisette.Internal.Core.Data.Class.ExtractSym (ExtractSym)
import Grisette.Internal.Core.Data.Class.Function (Apply (FunType), Function)
import Grisette.Internal.Core.Data.Class.Mergeable (Mergeable)
import Grisette.Internal.Core.Data.Class.PPrint (PPrint)
import Grisette.Internal.Core.Data.Class.SubstSym (SubstSym)
import Grisette.Internal.Core.Data.Class.ToCon (ToCon)
import Grisette.Internal.Core.Data.Class.ToSym (ToSym)
import Grisette.Internal.SymPrim.AlgReal (AlgReal)
import Grisette.Internal.SymPrim.BV (IntN, WordN)
import Grisette.Internal.SymPrim.FP (FP, ValidFP)
import Grisette.Internal.SymPrim.SymAlgReal (SymAlgReal)
import Grisette.Internal.SymPrim.SymBV (SymIntN, SymWordN)
import Grisette.Internal.SymPrim.SymBool (SymBool)
import Grisette.Internal.SymPrim.SymFP (SymFP)
import Grisette.Internal.SymPrim.SymInteger (SymInteger)
import Grisette.Internal.SymPrim.SymTabularFun (type (=~>))
import Grisette.Internal.SymPrim.TabularFun (type (=->))
import Grisette.Internal.Unified.EvalModeTag (EvalModeTag (C, S))
import Grisette.Internal.Unified.Theories
  ( TheoryToUnify (UAlgReal, UBool, UFP, UFun, UIntN, UInteger, UWordN),
  )
import Grisette.Internal.Unified.UnifiedAlgReal (GetAlgReal)
import Grisette.Internal.Unified.UnifiedBV (UnifiedBVImpl (GetIntN, GetWordN))
import Grisette.Internal.Unified.UnifiedBool (UnifiedBool (GetBool))
import Grisette.Internal.Unified.UnifiedFP (UnifiedFPImpl (GetFP))
import Grisette.Internal.Unified.UnifiedInteger (GetInteger)
import Language.Haskell.TH
  ( DecsQ,
    Pred,
    Q,
    TyLit (NumTyLit),
    Type (AppT, ConT, ForallT, LitT, VarT),
    appT,
    classD,
    conT,
    instanceD,
    mkName,
    newName,
    promotedT,
    varT,
  )
import qualified Language.Haskell.TH
import Language.Haskell.TH.Datatype.TyVarBndr
  ( kindedTV,
    mapTVFlag,
    specifiedSpec,
    tvName,
  )
import Language.Haskell.TH.Syntax (Lift)

#if MIN_VERSION_template_haskell(2,21,0)
type TyVarBndrVis = Language.Haskell.TH.TyVarBndrVis
#elif MIN_VERSION_template_haskell(2,17,0)
type TyVarBndrVis = Language.Haskell.TH.TyVarBndr ()
#else
type TyVarBndrVis = Language.Haskell.TH.TyVarBndr
#endif

-- | Provide unified function types.
class UnifiedFun (mode :: EvalModeTag) where
  -- | Get a unified function type. Resolves to t'Grisette.SymPrim.=->' in 'C'
  -- mode, and t'Grisette.SymPrim.=~>' in 'S' mode.
  type
    GetFun mode =
      (fun :: Data.Kind.Type -> Data.Kind.Type -> Data.Kind.Type) | fun -> mode

instance UnifiedFun 'C where
  type GetFun 'C = (=->)

instance UnifiedFun 'S where
  type GetFun 'S = (=~>)

-- | The unified function type with 2 arguments.
type GetFun2 mode a b = GetFun mode a b

-- | The unified function type with 3 arguments.
type GetFun3 mode a b c = GetFun mode a (GetFun mode b c)

-- | The unified function type with 4 arguments.
type GetFun4 mode a b c d = GetFun mode a (GetFun mode b (GetFun mode c d))

-- | The unified function type with 5 arguments.
type GetFun5 mode a b c d e =
  GetFun mode a (GetFun mode b (GetFun mode c (GetFun mode d e)))

-- | The unified function type with 6 arguments.
type GetFun6 mode a b c d e f =
  GetFun
    mode
    a
    (GetFun mode b (GetFun mode c (GetFun mode d (GetFun mode e f))))

-- | The unified function type with 7 arguments.
type GetFun7 mode a b c d e f g =
  GetFun
    mode
    a
    ( GetFun
        mode
        b
        (GetFun mode c (GetFun mode d (GetFun mode e (GetFun mode f g))))
    )

-- | The unified function type with 8 arguments.
type GetFun8 mode a b c d e f g h =
  GetFun
    mode
    a
    ( GetFun
        mode
        b
        ( GetFun
            mode
            c
            (GetFun mode d (GetFun mode e (GetFun mode f (GetFun mode g h))))
        )
    )

-- | The constraint for a unified function.
type UnifiedFunConstraint mode a b ca cb sa sb =
  ( Show (GetFun mode a b),
    Binary (GetFun mode a b),
    Serial (GetFun mode a b),
    Serialize (GetFun mode a b),
    NFData (GetFun mode a b),
    Eq (GetFun mode a b),
    EvalSym (GetFun mode a b),
    ExtractSym (GetFun mode a b),
    Mergeable (GetFun mode a b),
    PPrint (GetFun mode a b),
    SubstSym (GetFun mode a b),
    Hashable (GetFun mode a b),
    Lift (GetFun mode a b),
    Typeable (GetFun mode a b),
    ToCon (GetFun mode a b) (ca =-> cb),
    ToCon (sa =~> sb) (GetFun mode a b),
    ToSym (GetFun mode a b) (sa =~> sb),
    ToSym (ca =-> cb) (GetFun mode a b),
    Function (GetFun mode a b) a b,
    Apply (GetFun mode a b),
    FunType (GetFun mode a b) ~ (a -> b)
  )

genInnerUnifiedFunInstance ::
  String ->
  TyVarBndrVis ->
  [Pred] ->
  [TyVarBndrVis] ->
  [(Type, Type, Type)] ->
  DecsQ
genInnerUnifiedFunInstance nm mode preds bndrs tys = do
  x <- classD (goPred tys) (mkName nm) (mode : bndrs) [] []
  dc <-
    instanceD
      (return preds)
      (applyTypeList (promotedT 'C : additionalTypes))
      []
  ds <-
    instanceD
      (return preds)
      (applyTypeList (promotedT 'S : additionalTypes))
      []
  return [x, dc, ds]
  where
    additionalTypes = (varT . tvName) <$> bndrs
    applyTypeList = foldl appT (conT (mkName nm))
    goPred :: [(Type, Type, Type)] -> Q [Pred]
    goPred [] = fail "Empty list of function types, at least 2."
    goPred [_] = return []
    goPred (x : xs) = do
      p1 <- pred x xs
      pr <- goPred xs
      return $ p1 : pr
    listTys :: [(Type, Type, Type)] -> Q (Type, Type, Type)
    listTys [] = fail "Should not happen"
    listTys [(u, c, s)] = return (u, c, s)
    listTys ((u, c, s) : xs) = do
      (u', c', s') <- listTys xs
      return
        ( AppT (AppT (AppT (ConT ''GetFun) (VarT $ tvName mode)) u) u',
          AppT (AppT (ConT ''(=->)) c) c',
          AppT (AppT (ConT ''(=~>)) s) s'
        )
    pred (ua, ca, sa) l = do
      (ub, cb, sb) <- listTys l
      [t|
        UnifiedFunConstraint
          $(return (VarT $ tvName mode))
          $(return ua)
          $(return ub)
          $(return ca)
          $(return cb)
          $(return sa)
          $(return sb)
        |]

genOuterUnifiedFunInstance ::
  String -> String -> TyVarBndrVis -> [Pred] -> [TyVarBndrVis] -> DecsQ
genOuterUnifiedFunInstance nm innerName mode preds bndrs = do
  let bndrs' = mapTVFlag (const specifiedSpec) <$> bndrs
  x <-
    classD
      ( return
          [ ForallT bndrs' preds $
              foldl' AppT (ConT $ mkName innerName) $
                VarT . tvName <$> mode : bndrs
          ]
      )
      (mkName nm)
      [mode]
      []
      []
  dc <-
    instanceD
      (return [])
      (appT (conT $ mkName nm) (promotedT 'C))
      []
  ds <-
    instanceD
      (return [])
      (appT (conT $ mkName nm) (promotedT 'S))
      []
  return [x, dc, ds]

-- | Generate unified function instance names.
unifiedFunInstanceName :: String -> [TheoryToUnify] -> String
unifiedFunInstanceName prefix theories =
  prefix ++ "Fun" ++ (concatMap show theories)

-- | Generate unified function instances.
genUnifiedFunInstance :: String -> [TheoryToUnify] -> DecsQ
genUnifiedFunInstance prefix theories = do
  modeName <- newName "mode"
  let modeType = VarT modeName
  allArgs <- traverse (genArgs modeType) theories
  let baseName = unifiedFunInstanceName prefix theories
  rinner <-
    genInnerUnifiedFunInstance
      baseName
      (kindedTV modeName (ConT ''EvalModeTag))
      (concatMap (\(_, p, _, _, _) -> p) allArgs)
      (concatMap (\(t, _, _, _, _) -> t) allArgs)
      ((\(_, _, u, c, s) -> (u, c, s)) <$> allArgs)
  router <-
    if all (\(bndr, _, _, _, _) -> null bndr) allArgs
      then return []
      else
        genOuterUnifiedFunInstance
          ("All" ++ baseName)
          baseName
          (kindedTV modeName (ConT ''EvalModeTag))
          (concatMap (\(_, p, _, _, _) -> p) allArgs)
          (concatMap (\(t, _, _, _, _) -> t) allArgs)
  return $ rinner ++ router
  where
    genArgs ::
      Type -> TheoryToUnify -> Q ([TyVarBndrVis], [Pred], Type, Type, Type)
    genArgs mode UBool =
      return
        ( [],
          [],
          AppT (ConT ''GetBool) mode,
          ConT ''Bool,
          ConT ''SymBool
        )
    genArgs mode UIntN = do
      n <- newName "n"
      let nType = VarT n
      return
        ( [kindedTV n (ConT ''Nat)],
          [ AppT (ConT ''KnownNat) nType,
            AppT (AppT (ConT ''(<=)) (LitT $ NumTyLit 1)) nType
          ],
          AppT (AppT (ConT ''GetIntN) mode) nType,
          AppT (ConT ''IntN) nType,
          AppT (ConT ''SymIntN) nType
        )
    genArgs mode UWordN = do
      n <- newName "n"
      let nType = VarT n
      return
        ( [kindedTV n (ConT ''Nat)],
          [ AppT (ConT ''KnownNat) nType,
            AppT (AppT (ConT ''(<=)) (LitT $ NumTyLit 1)) nType
          ],
          AppT (AppT (ConT ''GetWordN) mode) nType,
          AppT (ConT ''WordN) nType,
          AppT (ConT ''SymWordN) nType
        )
    genArgs mode UInteger =
      return
        ( [],
          [],
          AppT (ConT ''GetInteger) mode,
          ConT ''Integer,
          ConT ''SymInteger
        )
    genArgs mode UAlgReal =
      return
        ( [],
          [],
          AppT (ConT ''GetAlgReal) mode,
          ConT ''AlgReal,
          ConT ''SymAlgReal
        )
    genArgs mode UFP = do
      eb <- newName "eb"
      sb <- newName "sb"
      let ebType = VarT eb
      let sbType = VarT sb
      return
        ( [kindedTV eb (ConT ''Nat), kindedTV sb (ConT ''Nat)],
          [AppT (AppT (ConT ''ValidFP) ebType) sbType],
          AppT (AppT (AppT (ConT ''GetFP) mode) ebType) sbType,
          AppT (AppT (ConT ''FP) ebType) sbType,
          AppT (AppT (ConT ''SymFP) ebType) sbType
        )
    genArgs _ UFun {} = fail "UFun cannot be nested."