Table.hs

-- | Simple hash tables, implemented as @Map hash (Map key value)@.
--
-- To be Haskell98 compatible (no multi-param type classes), when constructing 
-- a new hash table, we have to support the function computing (or just fetching, if
-- it is cached) the hash value. This function is then stored in the data type.
--
module Data.Generics.Fixplate.Hash.Table
  ( HashTable 
  , getHashValue , unHashTable
    -- * Construction and deconstruction
  , empty , singleton
  , fromList , toList 
  , bag
    -- * Membership
  , lookup , member
    -- * Insert
  , insert , insertWith
  )
  where

--------------------------------------------------------------------------------

import Prelude hiding ( lookup )

import Data.List ( foldl' )

import qualified Data.Map as Map ; import Data.Map (Map) 
-- import qualified Data.Set as Set ; import Data.Set (Set) 

--------------------------------------------------------------------------------
-- helper functions

mapInsertWith :: Ord k => (a -> v) -> (a -> v -> v) -> k -> a -> Map k v ->  Map k v
mapInsertWith f g k x = x `seq` Map.alter worker k where
  worker Nothing   =          Just $! (f x)
  worker (Just y)  = y `seq` (Just $! (g x y))

mapIsSingleton :: Map k v -> Maybe (k,v)
mapIsSingleton table = if Map.size table == 1 
  then let [(k,v)] = Map.toList table in Just (k,v)
  else Nothing

mapIsSingleton_ :: Map k v -> Maybe v
mapIsSingleton_ table = if Map.size table == 1 
  then let [(_,v)] = Map.toList table in Just v
  else Nothing

--------------------------------------------------------------------------------

--newtype HashTable hash k v = HashTable { unHashTable :: Map hash (Map k v) }

data HashTable hash k v = HashTable 
  { getHashValue :: k -> hash
  , unHashTable  :: Map hash (Map k v) 
  }

empty :: (Ord hash, Ord k) => (k -> hash) -> HashTable hash k v
empty gethash = HashTable gethash (Map.empty)

singleton :: (Ord hash, Ord k) => (k -> hash) -> k -> v -> HashTable hash k v
singleton gethash k v = HashTable gethash $ Map.singleton h (Map.singleton k v) where
  h = gethash k

fromList :: (Ord hash, Ord k) => (k -> hash) -> [(k,v)] -> HashTable hash k v
fromList gethash = foldl' (\old (k,v) -> insert k v old) (empty gethash)

toList :: Ord k => HashTable hash k v -> [(k,v)]
toList (HashTable _ table) = concat [ Map.toList sub | sub <- Map.elems table ]

--------------------------------------------------------------------------------

lookup :: (Ord hash, Ord k) => k -> HashTable hash k v -> Maybe v
lookup key (HashTable gethash table) = 
  case Map.lookup h table of
    Just sub -> case mapIsSingleton_ sub of
      Just v  -> Just v
      Nothing -> Map.lookup key sub
    Nothing  -> Nothing      
  where
    h = gethash key

member :: (Ord hash, Ord k) => k -> HashTable hash k v -> Bool
member key table = case lookup key table of
  Just _  -> True
  Nothing -> False

--------------------------------------------------------------------------------

insert :: (Ord hash, Ord k) => k -> v -> HashTable hash k v -> HashTable hash k v
insert k v (HashTable gethash table) = HashTable gethash $ mapInsertWith f g h v table where
  h = gethash k
  f v     = Map.singleton k v
  g v sub = Map.insert    k v sub

insertWith :: (Ord hash, Ord k) => (a -> v) -> (a -> v -> v) -> k -> a -> HashTable hash k v -> HashTable hash k v
insertWith ff gg k x (HashTable gethash table) = HashTable gethash $ mapInsertWith f g h x table where
  h = gethash k
  f x     = Map.singleton k (ff x)
  g x sub = mapInsertWith ff gg k x sub

--------------------------------------------------------------------------------

-- | Creates a multi-set from a list.
bag :: (Ord hash, Ord k) => (k -> hash) -> [k] -> HashTable hash k Int
bag gethash = foldl' (\old k -> insertWith id (+) k 1 old) (empty gethash)

--------------------------------------------------------------------------------