content: {-# LANGUAGE DefaultSignatures #-}
  {-# LANGUAGE LambdaCase #-}
  {-# LANGUAGE UndecidableInstances #-}
  {-# LANGUAGE ViewPatterns #-}
  
  {- |
  This module provides the class 'Codec' which takes care of marshalling data to and from external SMT-Solvers.
  
  A generic default implementation with 'GCodec' is possible.
  
  ==== __Example__
  
  @
  data V3 a = V3 a a a deriving Generic
  instance Codec a => Codec (V3 a)
  
  constantV3 :: V3 (Expr RealSort)
  constantV3 = encode $ V3 7 69 42
  @
  -}
  module Language.Hasmtlib.Codec
  (
    -- * Class
    Codec(..)
  
    -- * Generics
  , GCodec(..)
  , DefaultDecoded
  )
  where
  
  import Prelude hiding (not, (&&), (||), all, and)
  import Language.Hasmtlib.Type.Bitvec
  import Language.Hasmtlib.Type.Expr (Expr(..), SMTVar(..))
  import Language.Hasmtlib.Type.Solution
  import Language.Hasmtlib.Type.ArrayMap
  import Language.Hasmtlib.Type.SMTSort
  import Language.Hasmtlib.Type.Value
  import Language.Hasmtlib.Boolean
  import Data.Kind
  import Data.Proxy
  import Data.Coerce
  import qualified Data.List as List
  import Data.Bits hiding (And, Xor, xor)
  import Data.Map (Map)
  import Data.Sequence (Seq)
  import Data.IntMap as IM hiding (foldl)
  import Data.Dependent.Map as DMap
  import Data.Tree (Tree)
  import Data.Array (Array, Ix)
  import Data.Word
  import Data.Int
  import qualified Data.Text as Text
  import Data.Monoid (Sum, Product, First, Last, Dual)
  import qualified Data.Vector.Sized as V
  import Control.Monad
  import Control.Lens hiding (from, to)
  import GHC.Generics
  import GHC.TypeLits
  
  -- | Computes a default 'Decoded' 'Type' by distributing 'Decoded' over it's type arguments.
  type family DefaultDecoded a :: Type where
    DefaultDecoded (t a b c d e f g h) = t (Decoded a) (Decoded b) (Decoded c) (Decoded d) (Decoded e) (Decoded f) (Decoded g) (Decoded h)
    DefaultDecoded (t a b c d e f g) = t (Decoded a) (Decoded b) (Decoded c) (Decoded d) (Decoded e) (Decoded f) (Decoded g)
    DefaultDecoded (t a b c d e f) = t (Decoded a) (Decoded b) (Decoded c) (Decoded d) (Decoded e) (Decoded f)
    DefaultDecoded (t a b c d e) = t (Decoded a) (Decoded b) (Decoded c) (Decoded d) (Decoded e)
    DefaultDecoded (t a b c d) = t (Decoded a) (Decoded b) (Decoded c) (Decoded d)
    DefaultDecoded (t a b c) = t (Decoded a) (Decoded b) (Decoded c)
    DefaultDecoded (t a b) = t (Decoded a) (Decoded b)
    DefaultDecoded (t a) = t (Decoded a)
    DefaultDecoded x = TypeError (
            Text "DefaultDecoded (" :<>: ShowType x :<>: Text ") is not allowed."
      :$$:  Text "Try providing the associated Type Decoded (" :<>: ShowType x :<>: Text ") manually:"
      :$$:  Text "instance Codec (" :<>: ShowType x :<>: Text ") where "
      :$$:  Text "   type Decoded (" :<>: ShowType x :<>: Text ") = ... "
      )
  
  -- | Lift values to SMT-Values or decode them.
  --
  --   You can derive an instance of this class if your type is 'Generic'.
  class Codec a where
    -- | Result of decoding @a@.
    type Decoded a :: Type
    type Decoded a = DefaultDecoded a
  
    -- | Decode a value using given solution.
    decode :: Solution -> a -> Maybe (Decoded a)
    default decode :: (Generic a, Generic (Decoded a), GCodec (Rep a), GDecoded (Rep a) x ~ Rep (Decoded a) x) => Solution -> a -> Maybe (Decoded a)
    decode sol x = do
      gdecodedx <- gdecode sol $ from x
      Just $ to gdecodedx
  
    -- | Encode a value as constant.
    encode :: Decoded a -> a
    default encode :: (Generic a, Generic (Decoded a), GCodec (Rep a), GDecoded (Rep a) x ~ Rep (Decoded a) x) => Decoded a -> a
    encode = to . gencode . from
  
  -- | Decode and evaluate expressions
  instance KnownSMTSort t => Codec (Expr t) where
    type Decoded (Expr t) = HaskellType t
    decode sol (Var var)  = do
      let sungSort = sortSing @t
      (IntValueMap m) <- case sungSort of
        SBvSort enc n -> case bvEncSing' enc of
        -- Solution contains all BV as unsigned, if we have a Signed one we check the Unsigned ones and flip BvEnc
          SUnsigned -> DMap.lookup sungSort sol
          SSigned -> DMap.lookup (SBvSort (Proxy @Unsigned) n) sol <&>
            \case (IntValueMap ubvs) -> IntValueMap $ fmap (\case (BvValue ubv) -> BvValue $ asSigned ubv) ubvs
        _ -> DMap.lookup sungSort sol
      val <- IM.lookup (coerce var) m
      return $ unwrapValue val
    decode _ (Constant v)         = Just $ unwrapValue v
    decode sol (Plus x y)         = (+)   <$> decode sol x <*> decode sol y
    decode sol (Minus x y)        = (-)   <$> decode sol x <*> decode sol y
    decode sol (Neg x)            = fmap negate  (decode sol x)
    decode sol (Mul x y)          = (*)   <$> decode sol x <*> decode sol y
    decode sol (Abs x)            = fmap abs     (decode sol x)
    decode sol (Mod x y)          = mod   <$> decode sol x <*> decode sol y
    decode sol (Rem x y)          = rem   <$> decode sol x <*> decode sol y
    decode sol (IDiv x y)         = div   <$> decode sol x <*> decode sol y
    decode sol (Div x y)          = (/)   <$> decode sol x <*> decode sol y
    decode sol (LTH x y)          = (<)   <$> decode sol x <*> decode sol y
    decode sol (LTHE x y)         = (<=)  <$> decode sol x <*> decode sol y
    decode sol (EQU xs)           = do
      xs' <- decode sol (V.toList xs)
      case xs' of
        []   -> return true
        (x:xs'') -> return $ all (x ==) xs''
    decode sol (Distinct xs)      = do
      xs' <- decode sol (V.toList xs)
      let xss = List.filter ((==2) . length) $ List.permutations xs'
      return $ all (\case (a:b:_) -> a /= b ; _ -> true) xss
    decode sol (GTHE x y)         = (>=)  <$> decode sol x <*> decode sol y
    decode sol (GTH x y)          = (>)   <$> decode sol x <*> decode sol y
    decode sol (Not x)            = fmap   not  (decode sol x)
    decode sol (And x y)          = (&&) <$> decode sol x <*> decode sol y
    decode sol (Or x y)           = (||) <$> decode sol x <*> decode sol y
    decode sol (Impl x y)         = (==>) <$> decode sol x <*> decode sol y
    decode sol (Xor x y)          = xor   <$> decode sol x <*> decode sol y
    decode _ Pi                   = Just $ toRational pi
    decode sol (Sqrt x)           = fmap (toRational . sqrt . fromRational @Double)  (decode sol x)
    decode sol (Exp x)            = fmap (toRational . exp . fromRational @Double)   (decode sol x)
    decode sol (Sin x)            = fmap (toRational . sin . fromRational @Double)   (decode sol x)
    decode sol (Cos x)            = fmap (toRational . cos . fromRational @Double)   (decode sol x)
    decode sol (Tan x)            = fmap (toRational . tan . fromRational @Double)   (decode sol x)
    decode sol (Asin x)           = fmap (toRational . asin . fromRational @Double)  (decode sol x)
    decode sol (Acos x)           = fmap (toRational . acos . fromRational @Double)  (decode sol x)
    decode sol (Atan x)           = fmap (toRational . atan . fromRational @Double)  (decode sol x)
    decode sol (ToReal x)         = fmap realToFrac (decode sol x)
    decode sol (ToInt x)          = fmap truncate   (decode sol x)
    decode sol (IsInt x)          = fmap ((0 ==) . snd . properFraction) (decode sol x)
    decode sol (Ite p t f)        = (\p' t' f' -> if p' then t' else f') <$> decode sol p <*> decode sol t <*> decode sol f
    decode sol (BvNand x y)       = nand <$> sequenceA [decode sol x, decode sol y]
    decode sol (BvNor x y)        = nor  <$> sequenceA [decode sol x, decode sol y]
    decode sol (BvShL x y)        = do
      x' <- decode sol x
      y' <- decode sol y
      return $ shiftL x' $ fromIntegral (toInteger y')
    decode sol (BvLShR x y)       = do
      x' <- decode sol x
      y' <- decode sol y
      return $ shiftR x' $ fromIntegral (toInteger y')
    decode sol (BvAShR x y)       = do
      x' <- decode sol x
      y' <- decode sol y
      return $ shiftR x' $ fromIntegral (toInteger y')
    decode sol (BvConcat x y)     = bitvecConcat <$> decode sol x <*> decode sol y
    decode sol (BvRotL i x)       = rotateL <$> decode sol x <*> pure (fromIntegral i)
    decode sol (BvRotR i x)       = rotateR <$> decode sol x <*> pure (fromIntegral i)
    decode sol (ArrSelect i arr)  = arrSelect <$> decode sol i <*> decode sol arr
    decode sol (ArrStore i x arr) = arrStore <$> decode sol i <*> decode sol x <*> decode sol arr
    decode sol (StrConcat x y)         = (<>) <$> decode sol x <*> decode sol y
    decode sol (StrLength x)           = toInteger . Text.length <$> decode sol x
    decode sol (StrAt x i)             = (\x' i' -> Text.singleton $ Text.index x' (fromInteger i')) <$> decode sol x <*> decode sol i
    decode sol (StrSubstring x i j)    = (\x' (fromInteger -> i') (fromInteger -> j') -> Text.take (j' - i') $ Text.drop i' x') <$> decode sol x <*> decode sol i <*> decode sol j
    decode sol (StrPrefixOf x y)       = Text.isPrefixOf <$> decode sol x <*> decode sol y
    decode sol (StrSuffixOf x y)       = Text.isSuffixOf <$> decode sol x <*> decode sol y
    decode sol (StrContains x y)       = flip Text.isInfixOf <$> decode sol x <*> decode sol y
    decode sol (StrIndexOf x y i)      = join $ (\x' y' (fromInteger -> i') -> Text.findIndex ((y' ==) . Text.singleton) (Text.drop i' x') >>= Just . toInteger) <$> decode sol x <*> decode sol y <*> decode sol i
    decode sol (StrReplace src target replacement) = (\src' target' replacement' -> replaceOne target' replacement' src') <$> decode sol target <*> decode sol src <*> decode sol replacement
      where
        replaceOne pattern substitution text
          | Text.null back = text
          | otherwise = Text.concat [front, substitution, Text.drop (Text.length pattern) back]
            where
              (front, back) = Text.breakOn pattern text
    decode sol (StrReplaceAll src target replacement) = (\src' target' replacement' -> Text.replace target' replacement' src') <$> decode sol target <*> decode sol src <*> decode sol replacement
    decode _ (ForAll _ _)         = Nothing
    decode _ (Exists _ _)         = Nothing
    encode = Constant . wrapValue
  
  instance Codec () where type Decoded () = ()
  instance (Codec a, Codec b) => Codec (a,b)
  instance (Codec a, Codec b, Codec c) => Codec (a,b,c)
  instance (Codec a, Codec b, Codec c, Codec d) => Codec (a,b,c,d)
  instance (Codec a, Codec b, Codec c, Codec d, Codec e) => Codec (a,b,c,d,e)
  instance (Codec a, Codec b, Codec c, Codec d, Codec e, Codec f) => Codec (a,b,c,d,e,f)
  instance (Codec a, Codec b, Codec c, Codec d, Codec e, Codec f, Codec g) => Codec (a,b,c,d,e,f,g)
  instance (Codec a, Codec b, Codec c, Codec d, Codec e, Codec f, Codec g, Codec h) => Codec (a,b,c,d,e,f,g,h)
  instance Codec a => Codec [a]
  instance Codec a => Codec (Maybe a)
  instance Codec a => Codec (Tree a)
  instance (Codec a, Codec b) => Codec (Either a b)
  instance Codec a => Codec (Sum a)
  instance Codec a => Codec (Product a)
  instance Codec a => Codec (First a)
  instance Codec a => Codec (Last a)
  instance Codec a => Codec (Dual a)
  instance Codec a => Codec (Identity a)
  
  instance Codec a => Codec (IntMap a) where
    decode = traverse . decode
    encode = fmap encode
  
  instance Codec a => Codec (Seq a) where
    decode = traverse . decode
    encode = fmap encode
  
  instance Codec a => Codec (Map k a) where
    type Decoded (Map k a) = Map k (Decoded a)
    decode = traverse . decode
    encode = fmap encode
  
  instance (Ix i, Codec e) => Codec (Array i e) where
    type Decoded (Array i e) = Array i (Decoded e)
    decode = traverse . decode
    encode = fmap encode
  
  instance Codec Int where
    type Decoded Int = Int
    decode _ = Just
    encode = id
  
  instance Codec Integer where
    type Decoded Integer = Integer
    decode _ = Just
    encode = id
  
  instance Codec Natural where
    type Decoded Natural = Natural
    decode _ = Just
    encode = id
  
  instance Codec Word where
    type Decoded Word = Word
    decode _ = Just
    encode = id
  
  instance Codec Word8 where
    type Decoded Word8 = Word8
    decode _ = Just
    encode = id
  
  instance Codec Word16 where
    type Decoded Word16 = Word16
    decode _ = Just
    encode = id
  
  instance Codec Word32 where
    type Decoded Word32 = Word32
    decode _ = Just
    encode = id
  
  instance Codec Word64 where
    type Decoded Word64 = Word64
    decode _ = Just
    encode = id
  
  instance Codec Int8 where
    type Decoded Int8 = Int8
    decode _ = Just
    encode = id
  
  instance Codec Int16 where
    type Decoded Int16 = Int16
    decode _ = Just
    encode = id
  
  instance Codec Int32 where
    type Decoded Int32 = Int32
    decode _ = Just
    encode = id
  
  instance Codec Int64 where
    type Decoded Int64 = Int64
    decode _ = Just
    encode = id
  
  instance Codec Char where
    type Decoded Char = Char
    decode _ = Just
    encode = id
  
  instance Codec Float where
    type Decoded Float = Float
    decode _ = Just
    encode = id
  
  instance Codec Double where
    type Decoded Double = Double
    decode _ = Just
    encode = id
  
  instance Codec Ordering where
    type Decoded Ordering = Ordering
    decode _ = Just
    encode = id
  
  instance Codec Bool where
    type Decoded Bool = Bool
    decode _ = Just
    encode = id
  
  class GCodec f where
    type GDecoded f :: Type -> Type
    gdecode :: Solution -> f a -> Maybe (GDecoded f a)
    gencode :: GDecoded f a -> f a
  
  instance GCodec U1 where
    type GDecoded U1 = U1
    gdecode _ U1     = Just U1
    gencode          = id
  
  instance GCodec V1 where
    type GDecoded V1 = V1
    gdecode _        = Just
    gencode          = id
  
  instance (GCodec f, GCodec g) => GCodec (f :*: g) where
    type GDecoded (f :*: g) = (GDecoded f :*: GDecoded g)
    gdecode sol (a :*: b)   = (:*:) <$> gdecode sol a <*> gdecode sol b
    gencode (a :*: b)       = gencode a :*: gencode b
  
  instance (GCodec f, GCodec g) => GCodec (f :+: g) where
    type GDecoded (f :+: g) = (GDecoded f :+: GDecoded g)
    gdecode sol (L1 a)      = L1 <$> gdecode sol a
    gdecode sol (R1 a)      = R1 <$> gdecode sol a
    gencode (L1 a)          = L1 $ gencode a
    gencode (R1 a)          = R1 $ gencode a
  
  instance GCodec f => GCodec (M1 i c f) where
    type GDecoded (M1 i c f) = (M1 i c (GDecoded f))
    gdecode sol (M1 x)       = M1 <$> gdecode sol x
    gencode (M1 x)           = M1 $ gencode x
  
  instance Codec a => GCodec (K1 i a) where
    type GDecoded (K1 i a) = K1 i (Decoded a)
    gdecode sol (K1 a)     = K1 <$> decode sol a
    gencode (K1 a)         = K1 $ encode a
  