{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE DeriveDataTypeable #-}
module Language.Haskell.Liquid.GHC.Plugin.Util (
        partitionMaybe
      , extractSpecComments

      -- * Serialising and deserialising things from/to specs.
      , serialiseLiquidLib
      , deserialiseLiquidLib

      -- * Aborting the plugin execution
      , pluginAbort
      ) where

import           GhcPlugins                              as GHC
import           UniqDFM
import           IfaceSyn
import           Panic                                    ( throwGhcExceptionIO, GhcException(..) )
import           Data.Foldable                            ( asum )

import           Control.Monad.IO.Class
import           Control.Monad

import qualified Data.Binary                             as B
import           Data.Binary                              ( Binary )
import qualified Data.ByteString.Lazy                    as B
import           Data.Typeable
import           Data.Maybe                               ( listToMaybe )
import           Data.Data
import           Data.Either                              ( partitionEithers )

import           Language.Haskell.Liquid.GHC.Plugin.Types ( SpecComment
                                                          , LiquidLib
                                                          )


pluginAbort :: MonadIO m => String -> m a
pluginAbort :: String -> m a
pluginAbort = IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO a -> m a) -> (String -> IO a) -> String -> m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. GhcException -> IO a
forall a. GhcException -> IO a
throwGhcExceptionIO (GhcException -> IO a)
-> (String -> GhcException) -> String -> IO a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> GhcException
ProgramError

-- | Courtesy of [inspection testing](https://github.com/nomeata/inspection-testing/blob/master/src/Test/Inspection/Plugin.hs)
partitionMaybe :: (a -> Maybe b) -> [a] -> ([a], [b])
partitionMaybe :: (a -> Maybe b) -> [a] -> ([a], [b])
partitionMaybe a -> Maybe b
f = [Either a b] -> ([a], [b])
forall a b. [Either a b] -> ([a], [b])
partitionEithers ([Either a b] -> ([a], [b]))
-> ([a] -> [Either a b]) -> [a] -> ([a], [b])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Either a b) -> [a] -> [Either a b]
forall a b. (a -> b) -> [a] -> [b]
map (\a
x -> Either a b -> (b -> Either a b) -> Maybe b -> Either a b
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (a -> Either a b
forall a b. a -> Either a b
Left a
x) b -> Either a b
forall a b. b -> Either a b
Right (a -> Maybe b
f a
x))

-- | Extracts the spec comments from the Core 'Annotation's. It returns a
-- "cleaned" 'ModGuts' which doesn't contain the deserialised 'Annotation's.
-- This also means that these 'Annotation's /won't/ land into an interface file,
-- and we won't be able to retrieve them back later on.
extractSpecComments :: ModGuts -> (ModGuts, [SpecComment])
extractSpecComments :: ModGuts -> (ModGuts, [SpecComment])
extractSpecComments = ModGuts -> (ModGuts, [SpecComment])
forall a. (Typeable a, Data a) => ModGuts -> (ModGuts, [a])
extractModuleAnnotations

-- | Tries to deserialise the 'Annotation's in the input 'ModGuts', pruning the latter
-- upon successful deserialisation.
extractModuleAnnotations :: forall a. (Typeable a, Data a) => ModGuts -> (ModGuts, [a])
extractModuleAnnotations :: ModGuts -> (ModGuts, [a])
extractModuleAnnotations ModGuts
guts = (ModGuts
guts', [a]
extracted)
  where
    thisModule :: Module
thisModule = ModGuts -> Module
mg_module ModGuts
guts
    ([Annotation]
anns_clean, [a]
extracted) = (Annotation -> Maybe a) -> [Annotation] -> ([Annotation], [a])
forall a b. (a -> Maybe b) -> [a] -> ([a], [b])
partitionMaybe Annotation -> Maybe a
tryDeserialise (ModGuts -> [Annotation]
mg_anns ModGuts
guts)
    guts' :: ModGuts
guts' = ModGuts
guts { mg_anns :: [Annotation]
mg_anns = [Annotation]
anns_clean }

    tryDeserialise :: Annotation -> Maybe a
    tryDeserialise :: Annotation -> Maybe a
tryDeserialise (Annotation (ModuleTarget Module
m) AnnPayload
payload)
        | Module
thisModule Module -> Module -> Bool
forall a. Eq a => a -> a -> Bool
== Module
m = ([Word8] -> a) -> AnnPayload -> Maybe a
forall a. Typeable a => ([Word8] -> a) -> AnnPayload -> Maybe a
fromSerialized [Word8] -> a
forall a. Data a => [Word8] -> a
deserializeWithData AnnPayload
payload
        | Bool
otherwise       = Maybe a
forall a. Maybe a
Nothing
    tryDeserialise (Annotation (NamedTarget Name
_) AnnPayload
payload) --NOTE(adn) What is the correct behaviour here?
        | Just a
a <- ([Word8] -> a) -> AnnPayload -> Maybe a
forall a. Typeable a => ([Word8] -> a) -> AnnPayload -> Maybe a
fromSerialized [Word8] -> a
forall a. Data a => [Word8] -> a
deserializeWithData AnnPayload
payload
        = a -> Maybe a
forall a. a -> Maybe a
Just a
a
    tryDeserialise Annotation
_
        = Maybe a
forall a. Maybe a
Nothing

--
-- Serialising and deserialising Specs
--

deserialiseBinaryObject :: forall a. (Typeable a, Binary a)
                        => Module
                        -> ExternalPackageState
                        -> HomePackageTable
                        -> Maybe a
deserialiseBinaryObject :: Module -> ExternalPackageState -> HomePackageTable -> Maybe a
deserialiseBinaryObject Module
thisModule ExternalPackageState
eps HomePackageTable
hpt = [Maybe a] -> Maybe a
forall (t :: * -> *) (f :: * -> *) a.
(Foldable t, Alternative f) =>
t (f a) -> f a
asum [Maybe a
extractFromHpt, Maybe a
extractFromEps]
  where
    extractFromEps :: Maybe a
    extractFromEps :: Maybe a
extractFromEps = [a] -> Maybe a
forall a. [a] -> Maybe a
listToMaybe ([a] -> Maybe a) -> [a] -> Maybe a
forall a b. (a -> b) -> a -> b
$ ([Word8] -> a) -> AnnEnv -> AnnTarget Name -> [a]
forall a.
Typeable a =>
([Word8] -> a) -> AnnEnv -> AnnTarget Name -> [a]
findAnns [Word8] -> a
deserialise (ExternalPackageState -> AnnEnv
eps_ann_env ExternalPackageState
eps) (Module -> AnnTarget Name
forall name. Module -> AnnTarget name
ModuleTarget Module
thisModule)

    extractFromHpt :: Maybe a
    extractFromHpt :: Maybe a
extractFromHpt = do
      HomeModInfo
modInfo <- HomePackageTable -> ModuleName -> Maybe HomeModInfo
forall key elt. Uniquable key => UniqDFM elt -> key -> Maybe elt
lookupUDFM HomePackageTable
hpt (Module -> ModuleName
moduleName Module
thisModule)
      Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Module
thisModule Module -> Module -> Bool
forall a. Eq a => a -> a -> Bool
== (ModIface_ 'ModIfaceFinal -> Module
forall (phase :: ModIfacePhase). ModIface_ phase -> Module
mi_module (ModIface_ 'ModIfaceFinal -> Module)
-> (HomeModInfo -> ModIface_ 'ModIfaceFinal)
-> HomeModInfo
-> Module
forall b c a. (b -> c) -> (a -> b) -> a -> c
. HomeModInfo -> ModIface_ 'ModIfaceFinal
hm_iface (HomeModInfo -> Module) -> HomeModInfo -> Module
forall a b. (a -> b) -> a -> b
$ HomeModInfo
modInfo))
      [a]
xs <- (IfaceAnnotation -> Maybe a) -> [IfaceAnnotation] -> Maybe [a]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (([Word8] -> a) -> AnnPayload -> Maybe a
forall a. Typeable a => ([Word8] -> a) -> AnnPayload -> Maybe a
fromSerialized [Word8] -> a
deserialise (AnnPayload -> Maybe a)
-> (IfaceAnnotation -> AnnPayload) -> IfaceAnnotation -> Maybe a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IfaceAnnotation -> AnnPayload
ifAnnotatedValue) (ModIface_ 'ModIfaceFinal -> [IfaceAnnotation]
forall (phase :: ModIfacePhase).
ModIface_ phase -> [IfaceAnnotation]
mi_anns (ModIface_ 'ModIfaceFinal -> [IfaceAnnotation])
-> (HomeModInfo -> ModIface_ 'ModIfaceFinal)
-> HomeModInfo
-> [IfaceAnnotation]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. HomeModInfo -> ModIface_ 'ModIfaceFinal
hm_iface (HomeModInfo -> [IfaceAnnotation])
-> HomeModInfo -> [IfaceAnnotation]
forall a b. (a -> b) -> a -> b
$ HomeModInfo
modInfo)
      [a] -> Maybe a
forall a. [a] -> Maybe a
listToMaybe [a]
xs

    deserialise :: [B.Word8] -> a
    deserialise :: [Word8] -> a
deserialise [Word8]
payload = ByteString -> a
forall a. Binary a => ByteString -> a
B.decode ([Word8] -> ByteString
B.pack [Word8]
payload)

serialiseBinaryObject :: forall a. (Binary a, Typeable a) => a -> Module -> Annotation
serialiseBinaryObject :: a -> Module -> Annotation
serialiseBinaryObject a
obj Module
thisModule = Annotation
serialised
  where
    serialised :: Annotation
    serialised :: Annotation
serialised = AnnTarget Name -> AnnPayload -> Annotation
Annotation (Module -> AnnTarget Name
forall name. Module -> AnnTarget name
ModuleTarget Module
thisModule) ((a -> [Word8]) -> a -> AnnPayload
forall a. Typeable a => (a -> [Word8]) -> a -> AnnPayload
toSerialized (ByteString -> [Word8]
B.unpack (ByteString -> [Word8]) -> (a -> ByteString) -> a -> [Word8]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> ByteString
forall a. Binary a => a -> ByteString
B.encode) a
obj)

-- | Serialise a 'LiquidLib', removing the termination checks from the target.
serialiseLiquidLib :: LiquidLib -> Module -> Annotation
serialiseLiquidLib :: LiquidLib -> Module -> Annotation
serialiseLiquidLib LiquidLib
lib = LiquidLib -> Module -> Annotation
forall a. (Binary a, Typeable a) => a -> Module -> Annotation
serialiseBinaryObject @LiquidLib LiquidLib
lib

deserialiseLiquidLib :: Module -> ExternalPackageState -> HomePackageTable -> Maybe LiquidLib
deserialiseLiquidLib :: Module
-> ExternalPackageState -> HomePackageTable -> Maybe LiquidLib
deserialiseLiquidLib Module
thisModule = Module
-> ExternalPackageState -> HomePackageTable -> Maybe LiquidLib
forall a.
(Typeable a, Binary a) =>
Module -> ExternalPackageState -> HomePackageTable -> Maybe a
deserialiseBinaryObject @LiquidLib Module
thisModule