{-# LANGUAGE RecordWildCards #-}
{- |

'Data.CMap' is a wrapper around 'Data.Map' which transparently
compresses all the values (but not the keys) stored in the map.

To do this, the values must be an instance of the class 'Serialize'
from the the 'cereal' package.

This implementation is just proof-of-concept. In practice, it will
probably require more RAM than Data.Map since each value is compressed
individually.

A better implementation would probably store a common dictionary in
the 'CMap' which would be used when compressing the values. This would
result in a higher overall compression ratio.

-}
module Data.CMap where

import Codec.Compression.GZip ( CompressParams, compressWith
                              , decompress, defaultCompressParams)
import Control.Arrow          (second)
import Data.ByteString.Lazy   (ByteString)
import qualified Data.Map     as Map
import Data.Serialize         (Serialize, encodeLazy, decodeLazy)

------------------------------------------------------------------------------
-- CMap
------------------------------------------------------------------------------

data CMap key value = CMap 
    { params :: CompressParams
    , c_map  :: Map.Map key ByteString 
    }

------------------------------------------------------------------------------
-- fromList
------------------------------------------------------------------------------

-- | create a CMap from a list of key/value pairs
--
-- uses 'defaultCompressParams'
fromList :: (Ord key, Serialize value) =>
            [(key, value)] -- ^ list of key/value pairs
         -> CMap key value
fromList = fromListWith defaultCompressParams


-- | create a CMap from a list of key/value pairs
--
-- similar to 'fromList' but allows you to specify 'CompressParams'
fromListWith :: (Ord key, Serialize value) =>
            CompressParams -- ^ compressor settings
         -> [(key, value)] -- ^ list of key/value pairs
         -> CMap key value
fromListWith p kvs = 
    CMap { params = p
         , c_map  = Map.fromList (map (second $ compressor p) kvs)
         }

------------------------------------------------------------------------------
-- toList
------------------------------------------------------------------------------

-- | convert a 'CMap' back to a list of key/value pairs
toList :: (Serialize value) =>
          CMap key value  -- ^ CMap
       -> [(key, value)]
toList cmap = 
    map (second decoder) $ Map.toList (c_map cmap)

------------------------------------------------------------------------------
-- insert
------------------------------------------------------------------------------

-- | insert a new key/value pair into a 'CMap'
insert :: (Ord key, Serialize value) => 
          key             -- ^ key
       -> value           -- ^ value
       -> CMap key value  -- ^ CMap
       -> CMap key value
insert k v cmap@(CMap{..}) = 
    cmap { c_map = Map.insert k (compressor params v) c_map }

------------------------------------------------------------------------------
-- lookup
------------------------------------------------------------------------------


-- | lookup a key and return the value if a match is found
lookup :: (Ord key, Serialize value) => 
          key             -- ^ key to look for
       -> CMap key value  -- ^ 'CMap' to look in
       -> Maybe value
lookup k CMap{..} =
    fmap decoder $ Map.lookup k c_map

------------------------------------------------------------------------------
-- internal helpers for compressing/decompressing
------------------------------------------------------------------------------

decoder :: (Serialize a) => ByteString -> a
decoder = right . decodeLazy . decompress
    where
      right (Left e)  = error "toList: decode failed. :("
      right (Right v) = v

compressor :: (Serialize a) => CompressParams -> a -> ByteString
compressor params a =
    compressWith params (encodeLazy a)