{-# LANGUAGE ScopedTypeVariables #-}
module Language.Haskell.Liquid.GHC.Plugin.Serialisation (
      -- * Serialising and deserialising things from/to specs.
        serialiseLiquidLib
      , deserialiseLiquidLib
      , deserialiseLiquidLibFromEPS

      ) where

import qualified Data.Array                               as Array
import           Data.Foldable                            ( asum )

import           Control.Monad

import qualified Data.Binary                             as B
import qualified Data.Binary.Builder                     as Builder
import qualified Data.Binary.Put                         as B
import qualified Data.ByteString.Lazy                    as B
import           Data.Data (Data)
import           Data.Generics (ext0, gmapAccumT)
import           Data.HashMap.Strict                     as M
import           Data.Maybe                               ( listToMaybe )
import           Data.Word                               (Word8)

import qualified Liquid.GHC.API as GHC
import           Language.Haskell.Liquid.GHC.Plugin.Types (LiquidLib)
import           Language.Haskell.Liquid.Types.Names


--
-- Serialising and deserialising Specs
--

getLiquidLibBytesFromEPS
  :: GHC.Module
  -> GHC.ExternalPackageState
  -> Maybe LiquidLibBytes
getLiquidLibBytesFromEPS :: Module -> ExternalPackageState -> Maybe LiquidLibBytes
getLiquidLibBytesFromEPS Module
thisModule ExternalPackageState
eps = Maybe LiquidLibBytes
extractFromEps
  where
    extractFromEps :: Maybe LiquidLibBytes
    extractFromEps :: Maybe LiquidLibBytes
extractFromEps = [LiquidLibBytes] -> Maybe LiquidLibBytes
forall a. [a] -> Maybe a
listToMaybe ([LiquidLibBytes] -> Maybe LiquidLibBytes)
-> [LiquidLibBytes] -> Maybe LiquidLibBytes
forall a b. (a -> b) -> a -> b
$ ([Word8] -> LiquidLibBytes)
-> AnnEnv -> CoreAnnTarget -> [LiquidLibBytes]
forall a.
Typeable a =>
([Word8] -> a) -> AnnEnv -> CoreAnnTarget -> [a]
GHC.findAnns [Word8] -> LiquidLibBytes
LiquidLibBytes (ExternalPackageState -> AnnEnv
GHC.eps_ann_env ExternalPackageState
eps) (Module -> CoreAnnTarget
forall name. Module -> AnnTarget name
GHC.ModuleTarget Module
thisModule)

getLiquidLibBytes :: GHC.Module
                        -> GHC.ExternalPackageState
                        -> GHC.HomePackageTable
                        -> Maybe LiquidLibBytes
getLiquidLibBytes :: Module
-> ExternalPackageState -> HomePackageTable -> Maybe LiquidLibBytes
getLiquidLibBytes Module
thisModule ExternalPackageState
eps HomePackageTable
hpt =
    [Maybe LiquidLibBytes] -> Maybe LiquidLibBytes
forall (t :: * -> *) (f :: * -> *) a.
(Foldable t, Alternative f) =>
t (f a) -> f a
asum [Maybe LiquidLibBytes
extractFromHpt, Module -> ExternalPackageState -> Maybe LiquidLibBytes
getLiquidLibBytesFromEPS Module
thisModule ExternalPackageState
eps]
  where
    extractFromHpt :: Maybe LiquidLibBytes
    extractFromHpt :: Maybe LiquidLibBytes
extractFromHpt = do
      modInfo <- HomePackageTable -> ModuleName -> Maybe HomeModInfo
GHC.lookupHpt HomePackageTable
hpt (Module -> ModuleName
forall unit. GenModule unit -> ModuleName
GHC.moduleName Module
thisModule)
      guard (thisModule == (GHC.mi_module . GHC.hm_iface $ modInfo))
      xs <- mapM (GHC.fromSerialized LiquidLibBytes . GHC.ifAnnotatedValue) (GHC.mi_anns . GHC.hm_iface $ modInfo)
      listToMaybe xs

newtype LiquidLibBytes = LiquidLibBytes { LiquidLibBytes -> [Word8]
unLiquidLibBytes :: [Word8] }

-- | Serialise a 'LiquidLib', removing the termination checks from the target.
serialiseLiquidLib :: LiquidLib -> GHC.Module -> IO GHC.Annotation
serialiseLiquidLib :: LiquidLib -> Module -> IO Annotation
serialiseLiquidLib LiquidLib
lib Module
thisModule = do
    bs <- LiquidLib -> IO ByteString
encodeLiquidLib LiquidLib
lib
    return $ GHC.Annotation
      (GHC.ModuleTarget thisModule)
      (GHC.toSerialized unLiquidLibBytes (LiquidLibBytes $ B.unpack bs))

deserialiseLiquidLib
  :: GHC.Module
  -> GHC.ExternalPackageState
  -> GHC.HomePackageTable
  -> GHC.NameCache
  -> IO (Maybe LiquidLib)
deserialiseLiquidLib :: Module
-> ExternalPackageState
-> HomePackageTable
-> NameCache
-> IO (Maybe LiquidLib)
deserialiseLiquidLib Module
thisModule ExternalPackageState
eps HomePackageTable
hpt NameCache
nameCache = do
    let mlibbs :: Maybe LiquidLibBytes
mlibbs = Module
-> ExternalPackageState -> HomePackageTable -> Maybe LiquidLibBytes
getLiquidLibBytes Module
thisModule ExternalPackageState
eps HomePackageTable
hpt
    case Maybe LiquidLibBytes
mlibbs of
      Just (LiquidLibBytes [Word8]
ws) -> do
        let bs :: ByteString
bs = [Word8] -> ByteString
B.pack [Word8]
ws
        LiquidLib -> Maybe LiquidLib
forall a. a -> Maybe a
Just (LiquidLib -> Maybe LiquidLib)
-> IO LiquidLib -> IO (Maybe LiquidLib)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> NameCache -> ByteString -> IO LiquidLib
decodeLiquidLib NameCache
nameCache ByteString
bs
      Maybe LiquidLibBytes
_ -> Maybe LiquidLib -> IO (Maybe LiquidLib)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe LiquidLib
forall a. Maybe a
Nothing

deserialiseLiquidLibFromEPS
  :: GHC.Module
  -> GHC.ExternalPackageState
  -> GHC.NameCache
  -> IO (Maybe LiquidLib)
deserialiseLiquidLibFromEPS :: Module -> ExternalPackageState -> NameCache -> IO (Maybe LiquidLib)
deserialiseLiquidLibFromEPS Module
thisModule ExternalPackageState
eps NameCache
nameCache = do
    let mlibbs :: Maybe LiquidLibBytes
mlibbs = Module -> ExternalPackageState -> Maybe LiquidLibBytes
getLiquidLibBytesFromEPS Module
thisModule ExternalPackageState
eps
    case Maybe LiquidLibBytes
mlibbs of
      Just (LiquidLibBytes [Word8]
ws) -> do
        let bs :: ByteString
bs = [Word8] -> ByteString
B.pack [Word8]
ws
        LiquidLib -> Maybe LiquidLib
forall a. a -> Maybe a
Just (LiquidLib -> Maybe LiquidLib)
-> IO LiquidLib -> IO (Maybe LiquidLib)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> NameCache -> ByteString -> IO LiquidLib
decodeLiquidLib NameCache
nameCache ByteString
bs
      Maybe LiquidLibBytes
_ -> Maybe LiquidLib -> IO (Maybe LiquidLib)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe LiquidLib
forall a. Maybe a
Nothing

encodeLiquidLib :: LiquidLib -> IO B.ByteString
encodeLiquidLib :: LiquidLib -> IO ByteString
encodeLiquidLib LiquidLib
lib0 = do
    let (LiquidLib
lib1, [LHResolvedName]
ns) = LiquidLib -> (LiquidLib, [LHResolvedName])
forall a. Data a => a -> (a, [LHResolvedName])
collectLHNames LiquidLib
lib0
    bh <- Int -> IO BinHandle
GHC.openBinMem (Int
1024Int -> Int -> Int
forall a. Num a => a -> a -> a
*Int
1024)
    GHC.putWithUserData GHC.QuietBinIFace bh ns
    GHC.withBinBuffer bh $ \ByteString
bs ->
      ByteString -> IO ByteString
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (ByteString -> IO ByteString) -> ByteString -> IO ByteString
forall a b. (a -> b) -> a -> b
$ Builder -> ByteString
Builder.toLazyByteString (Builder -> ByteString) -> Builder -> ByteString
forall a b. (a -> b) -> a -> b
$ PutM () -> Builder
forall a. PutM a -> Builder
B.execPut (LiquidLib -> PutM ()
forall t. Binary t => t -> PutM ()
B.put LiquidLib
lib1) Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> ByteString -> Builder
Builder.fromByteString ByteString
bs

decodeLiquidLib :: GHC.NameCache -> B.ByteString -> IO LiquidLib
decodeLiquidLib :: NameCache -> ByteString -> IO LiquidLib
decodeLiquidLib NameCache
nameCache ByteString
bs0 = do
    case ByteString
-> Either
     (ByteString, ByteOffset, String)
     (ByteString, ByteOffset, LiquidLib)
forall a.
Binary a =>
ByteString
-> Either
     (ByteString, ByteOffset, String) (ByteString, ByteOffset, a)
B.decodeOrFail ByteString
bs0 of
      Left (ByteString
_, ByteOffset
_, String
err) -> String -> IO LiquidLib
forall a. HasCallStack => String -> a
error (String -> IO LiquidLib) -> String -> IO LiquidLib
forall a b. (a -> b) -> a -> b
$ String
"decodeLiquidLib: decodeOrFail: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
err
      Right (ByteString
bs1, ByteOffset
_, LiquidLib
lib) -> do
        bh <- ByteString -> IO BinHandle
GHC.unsafeUnpackBinBuffer (ByteString -> IO BinHandle) -> ByteString -> IO BinHandle
forall a b. (a -> b) -> a -> b
$ ByteString -> ByteString
B.toStrict ByteString
bs1
        ns <- GHC.getWithUserData nameCache bh
        let n = Int -> Word
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Word) -> Int -> Word
forall a b. (a -> b) -> a -> b
$ [LHResolvedName] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [LHResolvedName]
ns
            arr = (Word, Word) -> [LHResolvedName] -> Array Word LHResolvedName
forall i e. Ix i => (i, i) -> [e] -> Array i e
Array.listArray (Word
0, Word
n Word -> Word -> Word
forall a. Num a => a -> a -> a
- Word
1) [LHResolvedName]
ns
        return $ mapLHNames (resolveLHNameIndex arr) lib
  where
    resolveLHNameIndex :: Array.Array Word LHResolvedName -> LHName -> LHName
    resolveLHNameIndex :: Array Word LHResolvedName -> LHName -> LHName
resolveLHNameIndex Array Word LHResolvedName
arr LHName
lhname =
      case HasCallStack => LHName -> LHResolvedName
LHName -> LHResolvedName
getLHNameResolved LHName
lhname of
        LHRIndex Word
i ->
          if Word
i Word -> Word -> Bool
forall a. Ord a => a -> a -> Bool
<= (Word, Word) -> Word
forall a b. (a, b) -> b
snd (Array Word LHResolvedName -> (Word, Word)
forall i e. Array i e -> (i, i)
Array.bounds Array Word LHResolvedName
arr) then
            LHResolvedName -> Symbol -> LHName
makeResolvedLHName (Array Word LHResolvedName
arr Array Word LHResolvedName -> Word -> LHResolvedName
forall i e. Ix i => Array i e -> i -> e
Array.! Word
i) (LHName -> Symbol
getLHNameSymbol LHName
lhname)
          else
            String -> LHName
forall a. HasCallStack => String -> a
error (String -> LHName) -> String -> LHName
forall a b. (a -> b) -> a -> b
$ String
"decodeLiquidLib: index out of bounds: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ (Word, (Word, Word)) -> String
forall a. Show a => a -> String
show (Word
i, Array Word LHResolvedName -> (Word, Word)
forall i e. Array i e -> (i, i)
Array.bounds Array Word LHResolvedName
arr)
        LHResolvedName
_ ->
          LHName
lhname

newtype AccF a b = AccF { forall a b. AccF a b -> a -> b -> (a, b)
unAccF :: a -> b -> (a, b) }

collectLHNames :: Data a => a -> (a, [LHResolvedName])
collectLHNames :: forall a. Data a => a -> (a, [LHResolvedName])
collectLHNames a
t =
    let ((Word
_, HashMap LHResolvedName Word
_, [LHResolvedName]
xs), a
t') = (Word, HashMap LHResolvedName Word, [LHResolvedName])
-> a -> ((Word, HashMap LHResolvedName Word, [LHResolvedName]), a)
forall a.
Data a =>
(Word, HashMap LHResolvedName Word, [LHResolvedName])
-> a -> ((Word, HashMap LHResolvedName Word, [LHResolvedName]), a)
go (Word
0, HashMap LHResolvedName Word
forall k v. HashMap k v
M.empty, []) a
t
     in (a
t', [LHResolvedName] -> [LHResolvedName]
forall a. [a] -> [a]
reverse [LHResolvedName]
xs)
  where
    go
      :: Data a
      => (Word, M.HashMap LHResolvedName Word, [LHResolvedName])
      -> a
      -> ((Word, M.HashMap LHResolvedName Word, [LHResolvedName]), a)
    go :: forall a.
Data a =>
(Word, HashMap LHResolvedName Word, [LHResolvedName])
-> a -> ((Word, HashMap LHResolvedName Word, [LHResolvedName]), a)
go = (forall a.
 Data a =>
 (Word, HashMap LHResolvedName Word, [LHResolvedName])
 -> a -> ((Word, HashMap LHResolvedName Word, [LHResolvedName]), a))
-> (Word, HashMap LHResolvedName Word, [LHResolvedName])
-> a
-> ((Word, HashMap LHResolvedName Word, [LHResolvedName]), a)
forall d a.
Data d =>
(forall e. Data e => a -> e -> (a, e)) -> a -> d -> (a, d)
gmapAccumT ((forall a.
  Data a =>
  (Word, HashMap LHResolvedName Word, [LHResolvedName])
  -> a -> ((Word, HashMap LHResolvedName Word, [LHResolvedName]), a))
 -> (Word, HashMap LHResolvedName Word, [LHResolvedName])
 -> a
 -> ((Word, HashMap LHResolvedName Word, [LHResolvedName]), a))
-> (forall a.
    Data a =>
    (Word, HashMap LHResolvedName Word, [LHResolvedName])
    -> a -> ((Word, HashMap LHResolvedName Word, [LHResolvedName]), a))
-> (Word, HashMap LHResolvedName Word, [LHResolvedName])
-> a
-> ((Word, HashMap LHResolvedName Word, [LHResolvedName]), a)
forall a b. (a -> b) -> a -> b
$ AccF (Word, HashMap LHResolvedName Word, [LHResolvedName]) e
-> (Word, HashMap LHResolvedName Word, [LHResolvedName])
-> e
-> ((Word, HashMap LHResolvedName Word, [LHResolvedName]), e)
forall a b. AccF a b -> a -> b -> (a, b)
unAccF (AccF (Word, HashMap LHResolvedName Word, [LHResolvedName]) e
 -> (Word, HashMap LHResolvedName Word, [LHResolvedName])
 -> e
 -> ((Word, HashMap LHResolvedName Word, [LHResolvedName]), e))
-> AccF (Word, HashMap LHResolvedName Word, [LHResolvedName]) e
-> (Word, HashMap LHResolvedName Word, [LHResolvedName])
-> e
-> ((Word, HashMap LHResolvedName Word, [LHResolvedName]), e)
forall a b. (a -> b) -> a -> b
$ ((Word, HashMap LHResolvedName Word, [LHResolvedName])
 -> e -> ((Word, HashMap LHResolvedName Word, [LHResolvedName]), e))
-> AccF (Word, HashMap LHResolvedName Word, [LHResolvedName]) e
forall a b. (a -> b -> (a, b)) -> AccF a b
AccF (Word, HashMap LHResolvedName Word, [LHResolvedName])
-> e -> ((Word, HashMap LHResolvedName Word, [LHResolvedName]), e)
forall a.
Data a =>
(Word, HashMap LHResolvedName Word, [LHResolvedName])
-> a -> ((Word, HashMap LHResolvedName Word, [LHResolvedName]), a)
go AccF (Word, HashMap LHResolvedName Word, [LHResolvedName]) e
-> AccF
     (Word, HashMap LHResolvedName Word, [LHResolvedName])
     LHResolvedName
-> AccF (Word, HashMap LHResolvedName Word, [LHResolvedName]) e
forall a b (c :: * -> *).
(Typeable a, Typeable b) =>
c a -> c b -> c a
`ext0` ((Word, HashMap LHResolvedName Word, [LHResolvedName])
 -> LHResolvedName
 -> ((Word, HashMap LHResolvedName Word, [LHResolvedName]),
     LHResolvedName))
-> AccF
     (Word, HashMap LHResolvedName Word, [LHResolvedName])
     LHResolvedName
forall a b. (a -> b -> (a, b)) -> AccF a b
AccF (Word, HashMap LHResolvedName Word, [LHResolvedName])
-> LHResolvedName
-> ((Word, HashMap LHResolvedName Word, [LHResolvedName]),
    LHResolvedName)
forall {a}.
Hashable a =>
(Word, HashMap a Word, [a])
-> a -> ((Word, HashMap a Word, [a]), LHResolvedName)
collectName

    collectName :: (Word, HashMap a Word, [a])
-> a -> ((Word, HashMap a Word, [a]), LHResolvedName)
collectName acc :: (Word, HashMap a Word, [a])
acc@(Word
sz, HashMap a Word
m, [a]
xs) a
n = case a -> HashMap a Word -> Maybe Word
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
M.lookup a
n HashMap a Word
m of
      Just Word
i -> ((Word, HashMap a Word, [a])
acc, Word -> LHResolvedName
LHRIndex Word
i)
      Maybe Word
Nothing -> ((Word
sz Word -> Word -> Word
forall a. Num a => a -> a -> a
+ Word
1, a -> Word -> HashMap a Word -> HashMap a Word
forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
M.insert a
n Word
sz HashMap a Word
m, a
n a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
xs), Word -> LHResolvedName
LHRIndex Word
sz)