{-# LANGUAGE CPP, ExistentialQuantification, DeriveDataTypeable, ScopedTypeVariables #-}
-----------------------------------------------------------------------------
--
-- Module      :  Data.Binary.Shared
-- Copyright   :  2007-2010 Juergen "jutaro" Nicklisch-Franken
-- License     :  GPL
--
-- Maintainer  :  Jutaro <jutaro@leksah.org>
-- Stability   :  provisional
-- Portability :
--
-- | Binary serializing with sharing
--
-----------------------------------------------------------------------------

module Data.Binary.Shared (
    BinaryShared(..)
,   encodeFileSer
,   encodeSer
,   decodeSer
) where

import Data.Typeable (cast,Typeable(..))
#if MIN_VERSION_base(4,6,0)
import Data.Typeable (typeOf)
#else
import Data.Typeable (typeRepKey)
import System.IO.Unsafe (unsafePerformIO)
#endif
import qualified Control.Monad.State as St  (StateT(..),get,put)
import Data.Map (Map(..))
import qualified Data.Map as Map  (empty,fromDistinctAscList,toAscList,Map(..),insert,lookup)
import Data.IntMap (IntMap(..))
import qualified Data.IntMap as IMap  (empty,IntMap(..),insert,lookup)
import qualified Data.Binary as Bin (getWord8,putWord8,Get(..),Binary(..))
import Data.Binary.Put (runPut,PutM(..),putWord64be)
import Control.Monad.Trans (lift)
import Control.Monad (liftM2,replicateM,liftM)
import qualified Data.Set as Set  (fromDistinctAscList,toAscList,Set(..))
import qualified Data.ByteString.Lazy as L
import Data.ByteString (ByteString(..))
import Control.Monad.State.Lazy (evalStateT)
import Data.Binary.Get (runGet,getWord64be)

-- | A class for storing Binary instances with shared nodes.
-- Cycles are not supported, cause put and get is a one path process.

class (Typeable alpha, Ord alpha, Eq alpha, Show alpha) => BinaryShared alpha  where
    -- | Encode a value in the Put monad.
    put :: alpha  -> PutShared

    putShared :: (alpha -> PutShared) -> alpha -> PutShared
    putShared alpha -> PutShared
fput alpha
v = do
        (Map Object Int
dict, Int
unique) <- forall s (m :: * -> *). MonadState s m => m s
St.get
        case (forall alpha.
(Typeable alpha, Ord alpha, Eq alpha, Show alpha) =>
alpha -> Object
ObjC alpha
v) forall k a. Ord k => k -> Map k a -> Maybe a
`Map.lookup` Map Object Int
dict of
            Just Int
i  -> forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (Word8 -> PutM ()
Bin.putWord8 Word8
0 forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Word64 -> PutM ()
putWord64be (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
i))
            Maybe Int
Nothing -> do
                forall s (m :: * -> *). MonadState s m => s -> m ()
St.put (Map Object Int
dict,Int
unique forall a. Num a => a -> a -> a
+ Int
1)
                forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (Word8 -> PutM ()
Bin.putWord8 Word8
1)
                forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (Word64 -> PutM ()
putWord64be (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
unique))
                alpha -> PutShared
fput alpha
v
                (Map Object Int
dict2, Int
unique2) <- forall s (m :: * -> *). MonadState s m => m s
St.get
                let newDict :: Map Object Int
newDict = forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert (forall alpha.
(Typeable alpha, Ord alpha, Eq alpha, Show alpha) =>
alpha -> Object
ObjC alpha
v) Int
unique Map Object Int
dict2
                forall s (m :: * -> *). MonadState s m => s -> m ()
St.put (Map Object Int
newDict,Int
unique2)

    -- | Decode a value in the Get monad
    get :: GetShared alpha

    getShared :: GetShared alpha -> GetShared alpha
    getShared GetShared alpha
f = do
        IntMap Object
dict <- forall s (m :: * -> *). MonadState s m => m s
St.get
        Word8
w <- forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift Get Word8
Bin.getWord8
        case Word8
w of
            Word8
0 -> do
                Int
i   <- forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM forall a b. (Integral a, Num b) => a -> b
fromIntegral (Get Word64
getWord64be))
                case  forall a. Int -> IntMap a -> Maybe a
IMap.lookup Int
i IntMap Object
dict of
                    Just (ObjC alpha
obj) -> forall (m :: * -> *) a. Monad m => a -> m a
return (forall alpha. Maybe alpha -> String -> alpha
forceJust (forall a b. (Typeable a, Typeable b) => a -> Maybe b
cast alpha
obj)
                                            String
"Shared>>getShared: Cast failed")
                    Maybe Object
Nothing -> forall a. HasCallStack => String -> a
error forall a b. (a -> b) -> a -> b
$ String
"Shared>>getShared : Dont find in Map " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Int
i
            Word8
1 -> do
                Int
i   <- forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM forall a b. (Integral a, Num b) => a -> b
fromIntegral (Get Word64
getWord64be))
                alpha
obj <- GetShared alpha
f
                IntMap Object
dict2 <- forall s (m :: * -> *). MonadState s m => m s
St.get
                forall s (m :: * -> *). MonadState s m => s -> m ()
St.put (forall a. Int -> a -> IntMap a -> IntMap a
IMap.insert Int
i (forall alpha.
(Typeable alpha, Ord alpha, Eq alpha, Show alpha) =>
alpha -> Object
ObjC alpha
obj) IntMap Object
dict2)
                forall (m :: * -> *) a. Monad m => a -> m a
return alpha
obj
            Word8
_ -> forall a. HasCallStack => String -> a
error forall a b. (a -> b) -> a -> b
$ String
"Shared>>getShared : Encoding error"


-- * How to call this

encodeSer :: BinaryShared a => a -> L.ByteString
encodeSer :: forall a. BinaryShared a => a -> ByteString
encodeSer a
v = PutM () -> ByteString
runPut (forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m a
evalStateT (forall alpha. BinaryShared alpha => alpha -> PutShared
put a
v) (forall k a. Map k a
Map.empty,Int
0))

encodeFileSer :: BinaryShared a => FilePath -> a -> IO ()
encodeFileSer :: forall a. BinaryShared a => String -> a -> IO ()
encodeFileSer String
f a
v = String -> ByteString -> IO ()
L.writeFile String
f (forall a. BinaryShared a => a -> ByteString
encodeSer a
v)

decodeSer :: BinaryShared alpha  => L.ByteString -> alpha
decodeSer :: forall alpha. BinaryShared alpha => ByteString -> alpha
decodeSer =  forall a. Get a -> ByteString -> a
runGet (forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m a
evalStateT forall alpha. BinaryShared alpha => GetShared alpha
get forall a. IntMap a
IMap.empty)

-- * The types needed internally

data Object = forall alpha. (Typeable alpha, Ord alpha, Eq alpha, Show alpha) => ObjC {()
unObj :: alpha}

instance Eq Object where
    (ObjC alpha
a) == :: Object -> Object -> Bool
== (ObjC alpha
b) = if forall a. Typeable a => a -> TypeRep
typeOf alpha
a forall a. Eq a => a -> a -> Bool
/= forall a. Typeable a => a -> TypeRep
typeOf alpha
b
                                then Bool
False
                                else (forall a. a -> Maybe a
Just alpha
a) forall a. Eq a => a -> a -> Bool
== forall a b. (Typeable a, Typeable b) => a -> Maybe b
cast alpha
b

instance Ord Object where
    compare :: Object -> Object -> Ordering
compare (ObjC alpha
a) (ObjC alpha
b) = if forall a. Typeable a => a -> TypeRep
typeOf alpha
a forall a. Eq a => a -> a -> Bool
/= forall a. Typeable a => a -> TypeRep
typeOf alpha
b
#if MIN_VERSION_base(4,6,0)
                                then forall a. Ord a => a -> a -> Ordering
compare (forall a. Typeable a => a -> TypeRep
typeOf alpha
a) (forall a. Typeable a => a -> TypeRep
typeOf alpha
b)
#else
                                then compare ((unsafePerformIO . typeRepKey . typeOf) a)
                                                ((unsafePerformIO . typeRepKey . typeOf) b)
#endif
                                else forall a. Ord a => a -> a -> Ordering
compare (forall a. a -> Maybe a
Just alpha
a) (forall a b. (Typeable a, Typeable b) => a -> Maybe b
cast alpha
b)

type PutShared = St.StateT (Map Object Int, Int) PutM ()
type GetShared = St.StateT (IntMap Object) Bin.Get

-----------
-- * Some standard instances, but very incomplete

instance BinaryShared a => BinaryShared [a] where
    put :: [a] -> PutShared
put    = forall alpha.
BinaryShared alpha =>
(alpha -> PutShared) -> alpha -> PutShared
putShared (\[a]
l -> forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (forall t. Binary t => t -> PutM ()
Bin.put (forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
l)) forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ forall alpha. BinaryShared alpha => alpha -> PutShared
put [a]
l)
    get :: GetShared [a]
get    = forall alpha.
BinaryShared alpha =>
GetShared alpha -> GetShared alpha
getShared (do
                Int
n <- forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (forall t. Binary t => Get t
Bin.get :: Bin.Get Int)
                forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
n forall alpha. BinaryShared alpha => GetShared alpha
get)

instance (BinaryShared a) => BinaryShared (Maybe a) where
    put :: Maybe a -> PutShared
put Maybe a
Nothing  = forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (Word8 -> PutM ()
Bin.putWord8 Word8
0)
    put (Just a
x) = forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (Word8 -> PutM ()
Bin.putWord8 Word8
1) forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall alpha. BinaryShared alpha => alpha -> PutShared
put a
x
    get :: GetShared (Maybe a)
get = do
        Word8
w <- forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (Get Word8
Bin.getWord8)
        case Word8
w of
            Word8
0 -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a. Maybe a
Nothing
            Word8
_ -> forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM forall a. a -> Maybe a
Just forall alpha. BinaryShared alpha => GetShared alpha
get

instance (BinaryShared a, BinaryShared b) => BinaryShared (a,b) where
    put :: (a, b) -> PutShared
put (a
a,b
b)           = forall alpha. BinaryShared alpha => alpha -> PutShared
put a
a forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall alpha. BinaryShared alpha => alpha -> PutShared
put b
b
    get :: GetShared (a, b)
get                 = forall (m :: * -> *) a1 a2 r.
Monad m =>
(a1 -> a2 -> r) -> m a1 -> m a2 -> m r
liftM2 (,) forall alpha. BinaryShared alpha => GetShared alpha
get forall alpha. BinaryShared alpha => GetShared alpha
get

instance BinaryShared a => BinaryShared (Set.Set a) where
    put :: Set a -> PutShared
put Set a
s = forall alpha. BinaryShared alpha => alpha -> PutShared
put (forall a. Set a -> [a]
Set.toAscList Set a
s)
    get :: GetShared (Set a)
get   = forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM forall a. [a] -> Set a
Set.fromDistinctAscList forall alpha. BinaryShared alpha => GetShared alpha
get

instance (BinaryShared k, BinaryShared e) => BinaryShared (Map.Map k e) where
    put :: Map k e -> PutShared
put Map k e
m = forall alpha. BinaryShared alpha => alpha -> PutShared
put (forall k a. Map k a -> [(k, a)]
Map.toAscList Map k e
m)
    get :: GetShared (Map k e)
get   = forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM forall k a. [(k, a)] -> Map k a
Map.fromDistinctAscList forall alpha. BinaryShared alpha => GetShared alpha
get

instance BinaryShared Bool where
    put :: Bool -> PutShared
put = forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall t. Binary t => t -> PutM ()
Bin.put
    get :: GetShared Bool
get = forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall t. Binary t => Get t
Bin.get

instance BinaryShared Char where
    put :: Char -> PutShared
put = forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall t. Binary t => t -> PutM ()
Bin.put
    get :: GetShared Char
get = forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall t. Binary t => Get t
Bin.get

instance BinaryShared Int where
    put :: Int -> PutShared
put = forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall t. Binary t => t -> PutM ()
Bin.put
    get :: GetShared Int
get = forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall t. Binary t => Get t
Bin.get

instance BinaryShared Integer where
    put :: Integer -> PutShared
put = forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall t. Binary t => t -> PutM ()
Bin.put
    get :: GetShared Integer
get = forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall t. Binary t => Get t
Bin.get

instance BinaryShared ByteString where
    put :: ByteString -> PutShared
put = forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall t. Binary t => t -> PutM ()
Bin.put
    get :: GetShared ByteString
get = forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall t. Binary t => Get t
Bin.get


forceJust :: Maybe alpha -> String -> alpha
forceJust :: forall alpha. Maybe alpha -> String -> alpha
forceJust Maybe alpha
mb String
str = case Maybe alpha
mb of
			Maybe alpha
Nothing -> forall a. HasCallStack => String -> a
error String
str
			Just alpha
it -> alpha
it