{-# OPTIONS_GHC -Wunused-imports #-}

------------------------------------------------------------------------
-- | Hash tables.
------------------------------------------------------------------------

module Agda.Utils.HashTable
  ( HashTable
  , empty
  , insert
  , lookup
  , toList
  , keySet
  ) where

import Prelude hiding (lookup)

import Data.Hashable
import qualified Data.Vector.Hashtables as H
import qualified Data.Vector.Mutable as VM
import qualified Data.Vector as V
import Data.Set (Set)
import qualified Data.Set as Set

-- | Hash tables.

-- A very limited amount of (possibly outdated) testing indicates
-- that, for the use in Agda's serialiser/deserialiser,
-- Data.HashTable.IO.CuckooHashTable is somewhat slower than
-- Data.HashTable.IO.BasicHashTable, and that
-- Data.HashTable.IO.LinearHashTable and the hashtables from
-- Data.Hashtable are much slower. However, other (also possibly
-- outdated) testing suggests that Data.HashTable.IO.CuckooHashTable
-- is quite a bit faster than Data.HashTable.IO.BasicHashTable for
-- 64-bit Windows. Some more recent, also limited, testing suggests
-- that the following hash table implementation from
-- Data.Vector.Hashtables is quite a bit faster than
-- Data.HashTable.IO.BasicHashTable (see issue #5966).

newtype HashTable k v =
  HashTable (H.Dictionary (H.PrimState IO) VM.MVector k VM.MVector v)

-- | An empty hash table.

empty :: IO (HashTable k v)
empty :: forall k v. IO (HashTable k v)
empty = Dictionary RealWorld MVector k MVector v -> HashTable k v
Dictionary (PrimState IO) MVector k MVector v -> HashTable k v
forall k v.
Dictionary (PrimState IO) MVector k MVector v -> HashTable k v
HashTable (Dictionary RealWorld MVector k MVector v -> HashTable k v)
-> IO (Dictionary RealWorld MVector k MVector v)
-> IO (HashTable k v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> IO (Dictionary (PrimState IO) MVector k MVector v)
forall (ks :: * -> * -> *) k (vs :: * -> * -> *) v (m :: * -> *).
(MVector ks k, MVector vs v, PrimMonad m) =>
Int -> m (Dictionary (PrimState m) ks k vs v)
H.initialize Int
0

-- | Inserts the key and the corresponding value into the hash table.

insert :: (Eq k, Hashable k) => HashTable k v -> k -> v -> IO ()
insert :: forall k v. (Eq k, Hashable k) => HashTable k v -> k -> v -> IO ()
insert (HashTable Dictionary (PrimState IO) MVector k MVector v
h) = Dictionary (PrimState IO) MVector k MVector v -> k -> v -> IO ()
forall (ks :: * -> * -> *) k (vs :: * -> * -> *) v (m :: * -> *).
(MVector ks k, MVector vs v, PrimMonad m, Hashable k, Eq k) =>
Dictionary (PrimState m) ks k vs v -> k -> v -> m ()
H.insert Dictionary (PrimState IO) MVector k MVector v
h
{-# INLINABLE insert #-}

-- | Tries to find a value corresponding to the key in the hash table.

lookup :: (Eq k, Hashable k) => HashTable k v -> k -> IO (Maybe v)
lookup :: forall k v.
(Eq k, Hashable k) =>
HashTable k v -> k -> IO (Maybe v)
lookup (HashTable Dictionary (PrimState IO) MVector k MVector v
h) = Dictionary (PrimState IO) MVector k MVector v -> k -> IO (Maybe v)
forall (ks :: * -> * -> *) k (vs :: * -> * -> *) v (m :: * -> *).
(MVector ks k, MVector vs v, PrimMonad m, Hashable k, Eq k) =>
Dictionary (PrimState m) ks k vs v -> k -> m (Maybe v)
H.lookup Dictionary (PrimState IO) MVector k MVector v
h
{-# INLINABLE lookup #-}

-- | Converts the hash table to a list.
--
-- The order of the elements in the list is unspecified.

toList :: (Eq k, Hashable k) => HashTable k v -> IO [(k, v)]
toList :: forall k v. (Eq k, Hashable k) => HashTable k v -> IO [(k, v)]
toList (HashTable Dictionary (PrimState IO) MVector k MVector v
h) = Dictionary (PrimState IO) MVector k MVector v -> IO [(k, v)]
forall (ks :: * -> * -> *) k (vs :: * -> * -> *) v (m :: * -> *).
(MVector ks k, MVector vs v, PrimMonad m, Hashable k, Eq k) =>
Dictionary (PrimState m) ks k vs v -> m [(k, v)]
H.toList Dictionary (PrimState IO) MVector k MVector v
h
{-# INLINABLE toList #-}

keySet :: forall k v. Ord k => HashTable k v -> IO (Set k)
keySet :: forall k v. Ord k => HashTable k v -> IO (Set k)
keySet (HashTable Dictionary (PrimState IO) MVector k MVector v
h) = do
  (ks :: V.Vector k) <- Dictionary (PrimState IO) (Mutable Vector) k MVector v
-> IO (Vector k)
forall (ks :: * -> *) k (m :: * -> *) (vs :: * -> * -> *) v.
(Vector ks k, PrimMonad m) =>
Dictionary (PrimState m) (Mutable ks) k vs v -> m (ks k)
H.keys Dictionary (PrimState IO) MVector k MVector v
Dictionary (PrimState IO) (Mutable Vector) k MVector v
h
  pure $! V.foldl' (flip Set.insert) mempty ks
{-# INLINABLE keySet #-}