quickset-0.1.0: Data/QuickSet.hs

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

module Data.QuickSet
    ( -- * QuickSet
      QuickSet,

      -- * Construction
      fromList,
      fromListN,
      fromVector,

      -- * Query
      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)


-- | QuickSets are sets that use a compact unboxed vector as the
-- internal representation.  As such QuickSets are always strict in the
-- values.

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

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

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


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

fromList :: (Ord a, Unbox a) => [a] -> QuickSet a
fromList = fromVector . Vu.fromList


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

fromListN :: (Ord a, Unbox a) => Int -> [a] -> QuickSet a
fromListN n = fromVector . Vu.fromListN n


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

fromVector :: (Ord a, Unbox a) => Vu.Vector a -> QuickSet a
fromVector = QuickSet . Vu.modify Va.sort


-- | Check whether the given value is in the set.

member :: (Ord a, Unbox a) => a -> QuickSet a -> Bool
member k (QuickSet vec') =
    runST $ do
        vec <- Vu.unsafeThaw vec'
        i <- Va.binarySearchL vec k
        let k' = vec' Vu.! i
        return (if i < Vu.length vec'
                  then compare k k' == EQ
                  else False)