-- |
-- Module:     Data.QuickMap
-- Copyright:  (c) 2012 Ertugrul Soeylemez
-- License:    BSD3
-- Maintainer: Ertugrul Soeylemez <es@ertes.de>
--
-- This module implements very fast and compact query-only maps.

module Data.QuickMap
    ( -- * QuickMap
      QuickMap,

      -- * Construction
      fromList,
      fromListN,
      fromVector,

      -- * Query
      lookup,
      member
    )
    where

import qualified Data.Vector.Unboxed as Vu
import qualified Data.Vector.Algorithms.Heap as Va
import qualified Data.Vector.Algorithms.Search as Va
import Control.Monad
import Control.Monad.ST
import Data.Data
import Data.Ord
import Data.Vector.Unboxed (Unbox)
import Prelude hiding (lookup)


-- | QuickMaps are maps from keys to values that use a compact unboxed
-- vector as the internal representation.  As such QuickMaps are always
-- strict in both the keys and values.

newtype QuickMap k a =
    QuickMap (Vu.Vector (k, a))
    deriving (Data, Eq, Ord, Typeable)

instance (Ord k, Read a, Read k, Unbox a, Unbox k) => Read (QuickMap k a) where
    readsPrec pr =
        map (\(vec, r) -> (fromVector vec, r)) . readsPrec pr

instance (Show a, Show k, Unbox a, Unbox k) => Show (QuickMap k a) where
    show (QuickMap vec) = show vec


-- | Convert a list to a 'QuickMap'.

fromList :: (Ord k, Unbox a, Unbox k) => [(k, a)] -> QuickMap k a
fromList xs =
    (QuickMap . Vu.create)
    (do vec <- Vu.unsafeThaw (Vu.fromList xs)
        Va.sortBy (comparing fst) vec
        return vec)


-- | Convert a prefix of the given length of the given list to a
-- 'QuickMap'.

fromListN :: (Ord k, Unbox a, Unbox k) => Int -> [(k, a)] -> QuickMap k a
fromListN n xs =
    (QuickMap . Vu.create)
    (do vec <- Vu.unsafeThaw (Vu.fromListN n xs)
        Va.sortBy (comparing fst) vec
        return vec)


-- | Convert an unboxed vector to a 'QuickMap'.

fromVector :: (Ord k, Unbox a, Unbox k) => Vu.Vector (k, a) -> QuickMap k a
fromVector =
    QuickMap .
    Vu.modify (Va.sortBy (comparing fst))


-- | Try to look up a key.

lookup :: (Ord k, Unbox a, Unbox k) => k -> QuickMap k a -> Maybe a
lookup k (QuickMap vec') =
    runST $ do
        vec <- Vu.unsafeThaw vec'
        i <- Va.binarySearchLBy (comparing fst) vec (k, undefined)
        let (k', x) = vec' Vu.! i
        return $ do
            guard (i < Vu.length vec')
            case compare k k' of
              EQ -> return x
              _  -> Nothing


-- | Check whether the given key is in the map.

member :: (Ord k, Unbox a, Unbox k) => k -> QuickMap k a -> Bool
member k = maybe False (const True) . lookup k