{-# LANGUAGE CPP #-}
{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE Trustworthy #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}

-- |
-- Module      :   Grisette.Core.Data.Class.Substitute
-- Copyright   :   (c) Sirui Lu 2021-2023
-- License     :   BSD-3-Clause (see the LICENSE file)
--
-- Maintainer  :   siruilu@cs.washington.edu
-- Stability   :   Experimental
-- Portability :   GHC only
module Grisette.Core.Data.Class.Substitute
  ( -- * Substituting symbolic constants
    SubstituteSym (..),
    SubstituteSym' (..),
  )
where

import Control.Monad.Except
import Control.Monad.Identity
import Control.Monad.Trans.Maybe
import qualified Control.Monad.Writer.Lazy as WriterLazy
import qualified Control.Monad.Writer.Strict as WriterStrict
import qualified Data.ByteString as B
import Data.Functor.Sum
import Data.Int
import Data.Word
import Generics.Deriving
import Generics.Deriving.Instances ()
import Grisette.IR.SymPrim.Data.Prim.InternedTerm.Term
import Grisette.IR.SymPrim.Data.Prim.InternedTerm.TermSubstitution
import Grisette.IR.SymPrim.Data.Prim.InternedTerm.TermUtils
import {-# SOURCE #-} Grisette.IR.SymPrim.Data.SymPrim

-- $setup
-- >>> import Grisette.Core
-- >>> import Grisette.IR.SymPrim

-- | Substitution of symbolic constants.
--
-- >>> a = "a" :: TypedSymbol Bool
-- >>> v = "x" &&~ "y" :: SymBool
-- >>> substituteSym a v (["a" &&~ "b", "a"] :: [SymBool])
-- [(&& (&& x y) b),(&& x y)]
--
-- __Note 1:__ This type class can be derived for algebraic data types.
-- You may need the @DerivingVia@ and @DerivingStrategies@ extensions.
--
-- > data X = ... deriving Generic deriving SubstituteSym via (Default X)
class SubstituteSym a where
  -- Substitute a symbolic constant to some symbolic value
  --
  -- >>> substituteSym "a" ("c" &&~ "d" :: Sym Bool) ["a" &&~ "b" :: Sym Bool, "a"]
  -- [(&& (&& c d) b),(&& c d)]
  substituteSym :: TypedSymbol b -> Sym b -> a -> a

-- | Auxiliary class for 'SubstituteSym' instance derivation
class SubstituteSym' a where
  -- | Auxiliary function for 'substituteSym' derivation
  substituteSym' :: TypedSymbol b -> Sym b -> a c -> a c

instance
  ( Generic a,
    SubstituteSym' (Rep a)
  ) =>
  SubstituteSym (Default a)
  where
  substituteSym :: forall b. TypedSymbol b -> Sym b -> Default a -> Default a
substituteSym TypedSymbol b
sym Sym b
val = forall a. a -> Default a
Default forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a x. Generic a => Rep a x -> a
to forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (a :: * -> *) b c.
SubstituteSym' a =>
TypedSymbol b -> Sym b -> a c -> a c
substituteSym' TypedSymbol b
sym Sym b
val forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a x. Generic a => a -> Rep a x
from forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Default a -> a
unDefault

instance SubstituteSym' U1 where
  substituteSym' :: forall b c. TypedSymbol b -> Sym b -> U1 c -> U1 c
substituteSym' TypedSymbol b
_ Sym b
_ = forall a. a -> a
id

instance SubstituteSym c => SubstituteSym' (K1 i c) where
  substituteSym' :: forall b c. TypedSymbol b -> Sym b -> K1 i c c -> K1 i c c
substituteSym' TypedSymbol b
sym Sym b
val (K1 c
v) = forall k i c (p :: k). c -> K1 i c p
K1 forall a b. (a -> b) -> a -> b
$ forall a b. SubstituteSym a => TypedSymbol b -> Sym b -> a -> a
substituteSym TypedSymbol b
sym Sym b
val c
v

instance SubstituteSym' a => SubstituteSym' (M1 i c a) where
  substituteSym' :: forall b c. TypedSymbol b -> Sym b -> M1 i c a c -> M1 i c a c
substituteSym' TypedSymbol b
sym Sym b
val (M1 a c
v) = forall k i (c :: Meta) (f :: k -> *) (p :: k). f p -> M1 i c f p
M1 forall a b. (a -> b) -> a -> b
$ forall (a :: * -> *) b c.
SubstituteSym' a =>
TypedSymbol b -> Sym b -> a c -> a c
substituteSym' TypedSymbol b
sym Sym b
val a c
v

instance (SubstituteSym' a, SubstituteSym' b) => SubstituteSym' (a :+: b) where
  substituteSym' :: forall b c. TypedSymbol b -> Sym b -> (:+:) a b c -> (:+:) a b c
substituteSym' TypedSymbol b
sym Sym b
val (L1 a c
l) = forall k (f :: k -> *) (g :: k -> *) (p :: k). f p -> (:+:) f g p
L1 forall a b. (a -> b) -> a -> b
$ forall (a :: * -> *) b c.
SubstituteSym' a =>
TypedSymbol b -> Sym b -> a c -> a c
substituteSym' TypedSymbol b
sym Sym b
val a c
l
  substituteSym' TypedSymbol b
sym Sym b
val (R1 b c
r) = forall k (f :: k -> *) (g :: k -> *) (p :: k). g p -> (:+:) f g p
R1 forall a b. (a -> b) -> a -> b
$ forall (a :: * -> *) b c.
SubstituteSym' a =>
TypedSymbol b -> Sym b -> a c -> a c
substituteSym' TypedSymbol b
sym Sym b
val b c
r

instance (SubstituteSym' a, SubstituteSym' b) => SubstituteSym' (a :*: b) where
  substituteSym' :: forall b c. TypedSymbol b -> Sym b -> (:*:) a b c -> (:*:) a b c
substituteSym' TypedSymbol b
sym Sym b
val (a c
a :*: b c
b) = forall (a :: * -> *) b c.
SubstituteSym' a =>
TypedSymbol b -> Sym b -> a c -> a c
substituteSym' TypedSymbol b
sym Sym b
val a c
a forall k (f :: k -> *) (g :: k -> *) (p :: k).
f p -> g p -> (:*:) f g p
:*: forall (a :: * -> *) b c.
SubstituteSym' a =>
TypedSymbol b -> Sym b -> a c -> a c
substituteSym' TypedSymbol b
sym Sym b
val b c
b

#define CONCRETE_SUBSTITUTESYM(type) \
instance  SubstituteSym type where \
  substituteSym _ _ = id

#if 1
CONCRETE_SUBSTITUTESYM(Bool)
CONCRETE_SUBSTITUTESYM(Integer)
CONCRETE_SUBSTITUTESYM(Char)
CONCRETE_SUBSTITUTESYM(Int)
CONCRETE_SUBSTITUTESYM(Int8)
CONCRETE_SUBSTITUTESYM(Int16)
CONCRETE_SUBSTITUTESYM(Int32)
CONCRETE_SUBSTITUTESYM(Int64)
CONCRETE_SUBSTITUTESYM(Word)
CONCRETE_SUBSTITUTESYM(Word8)
CONCRETE_SUBSTITUTESYM(Word16)
CONCRETE_SUBSTITUTESYM(Word32)
CONCRETE_SUBSTITUTESYM(Word64)
CONCRETE_SUBSTITUTESYM(B.ByteString)
#endif

instance SubstituteSym () where
  substituteSym :: forall b. TypedSymbol b -> Sym b -> () -> ()
substituteSym TypedSymbol b
_ Sym b
_ = forall a. a -> a
id

-- Either
deriving via
  (Default (Either a b))
  instance
    ( SubstituteSym a,
      SubstituteSym b
    ) =>
    SubstituteSym (Either a b)

-- Maybe
deriving via (Default (Maybe a)) instance (SubstituteSym a) => SubstituteSym (Maybe a)

-- List
deriving via (Default [a]) instance (SubstituteSym a) => SubstituteSym [a]

-- (,)
deriving via
  (Default (a, b))
  instance
    (SubstituteSym a, SubstituteSym b) =>
    SubstituteSym (a, b)

-- (,,)
deriving via
  (Default (a, b, c))
  instance
    ( SubstituteSym a,
      SubstituteSym b,
      SubstituteSym c
    ) =>
    SubstituteSym (a, b, c)

-- (,,,)
deriving via
  (Default (a, b, c, d))
  instance
    ( SubstituteSym a,
      SubstituteSym b,
      SubstituteSym c,
      SubstituteSym d
    ) =>
    SubstituteSym (a, b, c, d)

-- (,,,,)
deriving via
  (Default (a, b, c, d, e))
  instance
    ( SubstituteSym a,
      SubstituteSym b,
      SubstituteSym c,
      SubstituteSym d,
      SubstituteSym e
    ) =>
    SubstituteSym (a, b, c, d, e)

-- (,,,,,)
deriving via
  (Default (a, b, c, d, e, f))
  instance
    ( SubstituteSym a,
      SubstituteSym b,
      SubstituteSym c,
      SubstituteSym d,
      SubstituteSym e,
      SubstituteSym f
    ) =>
    SubstituteSym (a, b, c, d, e, f)

-- (,,,,,,)
deriving via
  (Default (a, b, c, d, e, f, g))
  instance
    ( SubstituteSym a,
      SubstituteSym b,
      SubstituteSym c,
      SubstituteSym d,
      SubstituteSym e,
      SubstituteSym f,
      SubstituteSym g
    ) =>
    SubstituteSym (a, b, c, d, e, f, g)

-- (,,,,,,,)
deriving via
  (Default (a, b, c, d, e, f, g, h))
  instance
    ( SubstituteSym a,
      SubstituteSym b,
      SubstituteSym c,
      SubstituteSym d,
      SubstituteSym e,
      SubstituteSym f,
      SubstituteSym g,
      SubstituteSym h
    ) =>
    SubstituteSym ((,,,,,,,) a b c d e f g h)

-- MaybeT
instance
  (SubstituteSym (m (Maybe a))) =>
  SubstituteSym (MaybeT m a)
  where
  substituteSym :: forall b. TypedSymbol b -> Sym b -> MaybeT m a -> MaybeT m a
substituteSym TypedSymbol b
sym Sym b
val (MaybeT m (Maybe a)
v) = forall (m :: * -> *) a. m (Maybe a) -> MaybeT m a
MaybeT forall a b. (a -> b) -> a -> b
$ forall a b. SubstituteSym a => TypedSymbol b -> Sym b -> a -> a
substituteSym TypedSymbol b
sym Sym b
val m (Maybe a)
v

-- ExceptT
instance
  (SubstituteSym (m (Either e a))) =>
  SubstituteSym (ExceptT e m a)
  where
  substituteSym :: forall b. TypedSymbol b -> Sym b -> ExceptT e m a -> ExceptT e m a
substituteSym TypedSymbol b
sym Sym b
val (ExceptT m (Either e a)
v) = forall e (m :: * -> *) a. m (Either e a) -> ExceptT e m a
ExceptT forall a b. (a -> b) -> a -> b
$ forall a b. SubstituteSym a => TypedSymbol b -> Sym b -> a -> a
substituteSym TypedSymbol b
sym Sym b
val m (Either e a)
v

-- Sum
deriving via
  (Default (Sum f g a))
  instance
    (SubstituteSym (f a), SubstituteSym (g a)) =>
    SubstituteSym (Sum f g a)

-- WriterT
instance
  (SubstituteSym (m (a, s))) =>
  SubstituteSym (WriterLazy.WriterT s m a)
  where
  substituteSym :: forall b. TypedSymbol b -> Sym b -> WriterT s m a -> WriterT s m a
substituteSym TypedSymbol b
sym Sym b
val (WriterLazy.WriterT m (a, s)
v) = forall w (m :: * -> *) a. m (a, w) -> WriterT w m a
WriterLazy.WriterT forall a b. (a -> b) -> a -> b
$ forall a b. SubstituteSym a => TypedSymbol b -> Sym b -> a -> a
substituteSym TypedSymbol b
sym Sym b
val m (a, s)
v

instance
  (SubstituteSym (m (a, s))) =>
  SubstituteSym (WriterStrict.WriterT s m a)
  where
  substituteSym :: forall b. TypedSymbol b -> Sym b -> WriterT s m a -> WriterT s m a
substituteSym TypedSymbol b
sym Sym b
val (WriterStrict.WriterT m (a, s)
v) = forall w (m :: * -> *) a. m (a, w) -> WriterT w m a
WriterStrict.WriterT forall a b. (a -> b) -> a -> b
$ forall a b. SubstituteSym a => TypedSymbol b -> Sym b -> a -> a
substituteSym TypedSymbol b
sym Sym b
val m (a, s)
v

-- Identity
instance SubstituteSym a => SubstituteSym (Identity a) where
  substituteSym :: forall b. TypedSymbol b -> Sym b -> Identity a -> Identity a
substituteSym TypedSymbol b
sym Sym b
val (Identity a
a) = forall a. a -> Identity a
Identity forall a b. (a -> b) -> a -> b
$ forall a b. SubstituteSym a => TypedSymbol b -> Sym b -> a -> a
substituteSym TypedSymbol b
sym Sym b
val a
a

-- IdentityT
instance SubstituteSym (m a) => SubstituteSym (IdentityT m a) where
  substituteSym :: forall b. TypedSymbol b -> Sym b -> IdentityT m a -> IdentityT m a
substituteSym TypedSymbol b
sym Sym b
val (IdentityT m a
a) = forall {k} (f :: k -> *) (a :: k). f a -> IdentityT f a
IdentityT forall a b. (a -> b) -> a -> b
$ forall a b. SubstituteSym a => TypedSymbol b -> Sym b -> a -> a
substituteSym TypedSymbol b
sym Sym b
val m a
a

instance SubstituteSym (Sym a) where
  substituteSym :: forall b. TypedSymbol b -> Sym b -> Sym a -> Sym a
substituteSym TypedSymbol b
sym (Sym Term b
val) (Sym Term a
x) =
    forall t a. Term t -> (SupportedPrim t => a) -> a
introSupportedPrimConstraint Term b
val forall a b. (a -> b) -> a -> b
$
      forall t a. Term t -> (SupportedPrim t => a) -> a
introSupportedPrimConstraint Term a
x forall a b. (a -> b) -> a -> b
$
        forall a. Term a -> Sym a
Sym forall a b. (a -> b) -> a -> b
$
          forall a b.
(SupportedPrim a, SupportedPrim b) =>
TypedSymbol a -> Term a -> Term b -> Term b
substTerm TypedSymbol b
sym Term b
val Term a
x