-- | A module that implements a dictionary/hash table

module Text.Regex.Deriv.Dictionary where


import qualified Data.IntMap as IM
import Data.Char
import Prelude hiding (all, words)

class Key a where
    hash :: a -> [Int]


instance Key Int where
    hash i = [i]

instance Key Char where
    hash c = [(ord c)]

instance (Key a, Key b) => Key (a,b) where
    hash (a,b) = hash a ++ hash b

instance (Key a, Key b, Key c) => Key (a,b,c) where
    hash (a,b,c) = hash a ++ hash b ++ hash c


instance Key a => Key [a] where
    hash as = concatMap hash as


-- an immutable dictionary
newtype Dictionary a = Dictionary (Trie a)

primeL :: Int
primeL = 757
primeR :: Int
primeR = 577

empty :: Dictionary a
empty = Dictionary emptyTrie

-- insert and overwrite
insert :: Key k => k -> a -> Dictionary a -> Dictionary a
insert key val (Dictionary trie) =
    let key_hash = hash key
    in key_hash `seq` Dictionary (insertTrie True key_hash val trie)

-- insert not overwrite
insertNotOverwrite :: Key k => k -> a -> Dictionary a -> Dictionary a
insertNotOverwrite key val (Dictionary trie) =
    let key_hash = hash key
    in key_hash `seq` Dictionary (insertTrie False key_hash val trie)



lookup :: Key k => k -> Dictionary a -> Maybe a
lookup key (Dictionary trie) =
    let key_hash = hash key
    in key_hash `seq`
       case lookupTrie key_hash trie of
        Just (Trie (x:_) _) -> Just x
        _ -> Nothing

lookupAll :: Key k => k -> Dictionary a -> [a]
lookupAll key (Dictionary trie) =
    let key_hash = hash key
    in key_hash `seq`
       case lookupTrie key_hash trie of
        Just (Trie xs _) -> xs
        _ -> []

member :: Key k => k -> Dictionary a -> Bool
member key d =
   case Text.Regex.Deriv.Dictionary.lookup key d of
     { Just _ -> True
     ; Nothing -> False
     }

fromList :: Key k => [(k,a)] -> Dictionary a
fromList l = foldl (\d (key,val) -> insert key val d) empty l

fromListNotOverwrite :: Key k => [(k,a)] -> Dictionary a
fromListNotOverwrite l = foldl (\d (key,val) -> insertNotOverwrite key val d) empty l

update :: Key k => k -> a -> Dictionary a -> Dictionary a
update key val (Dictionary trie) =
    let key_hash = hash key
        trie'     = key_hash `seq` updateTrie key_hash val trie
    in Dictionary trie'


-- The following are some special functions we implemented for
-- an special instance of the dictionary 'Dictionary (k,a)'
-- in which we store both the key k together with the actual value a,
-- i.e. we map (hash k) to list of (k,a) value pairs

-- ^ the dictionary (k,a) version of elem
isIn :: (Key k, Eq k) => k -> Dictionary (k,a) -> Bool
isIn k dict =
    let all = lookupAll (hash k) dict
    in k `elem` (map fst all)

nub :: (Key k, Eq k) => [k] -> [k]
nub ks = nubSub ks empty

nubSub :: (Key k, Eq k) => [k] -> Dictionary (k,()) -> [k]
nubSub [] _ = []
nubSub (x:xs) d
    | x `isIn` d = nubSub xs d
    | otherwise = let d' = insertNotOverwrite x (x,()) d
                  in x:(nubSub xs d')




-- An internal trie which we use to implement the dictoinar

data Trie a = Trie ![a] !(IM.IntMap (Trie a))

emptyTrie :: Trie a
emptyTrie = Trie [] (IM.empty)


insertTrie :: Bool -> [Int] -> a -> Trie a -> Trie a
insertTrie overwrite [] i (Trie is maps)
    | overwrite  =  Trie [i] maps
    | otherwise  =  Trie (i:is) maps
insertTrie overwrite (word:words) i (Trie is maps) =
    let key = word
    in key `seq` case IM.lookup key maps of
    { Just trie ->  let trie' = insertTrie overwrite words i trie
                        maps' = trie' `seq` IM.update (\_ -> Just trie') key maps
                    in maps' `seq` Trie is maps'
    ; Nothing ->    let trie = emptyTrie
                        trie' = insertTrie overwrite words i trie 
                        maps' = trie' `seq` IM.insert key trie' maps
                    in maps' `seq` Trie is maps'
    }

lookupTrie :: [Int] -> Trie a -> Maybe (Trie a)
lookupTrie [] trie = Just trie
lookupTrie (word:words) (Trie _ maps) =
    let key = word
    in case IM.lookup key maps of
        Just trie -> lookupTrie words trie
        Nothing   -> Nothing

-- we only update the first one, not the collided ones
updateTrie :: [Int] -> a -> Trie a -> Trie a
updateTrie [] y (Trie (_:xs) maps) = Trie (y:xs) maps
updateTrie (word:words) v  (Trie is maps) =
    let key = word
    in case IM.lookup key maps of
        Just trie ->    let trie' = updateTrie words v trie
                            maps'  = IM.update (\_ -> Just trie') key maps
                        in Trie is maps'
        Nothing   -> Trie is maps
