------------------------------------------------------------------------
-- |
-- Module           : Data.Parameterized.HashTable
-- Description      : a hash table for parameterized keys and values
-- Copyright        : (c) Galois, Inc 2014-2019
-- Maintainer       : Joe Hendrix <jhendrix@galois.com>
--
-- This module provides a 'ST'-based hashtable for parameterized keys and values.
--
-- NOTE: This API makes use of 'unsafeCoerce' to implement the parameterized
-- hashtable abstraction.  This should be type-safe provided that the
-- 'TestEquality' instance on the key type is implemented soundly.
------------------------------------------------------------------------
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE Trustworthy #-}
module Data.Parameterized.HashTable
  ( HashTable
  , new
  , newSized
  , clone
  , lookup
  , insert
  , member
  , delete
  , clear
  , Data.Parameterized.Classes.HashableF(..)
  , Control.Monad.ST.RealWorld
  ) where

import Control.Applicative
import Control.Monad.ST
import qualified Data.HashTable.ST.Basic as H
import Data.Kind
import GHC.Exts (Any)
import Unsafe.Coerce

import Prelude hiding (lookup)

import Data.Parameterized.Classes
import Data.Parameterized.Some

-- | A hash table mapping nonces to values.
newtype HashTable s (key :: k -> Type) (val :: k -> Type)
      = HashTable (H.HashTable s (Some key) Any)

-- | Create a new empty table.
new :: ST s (HashTable s key val)
new :: ST s (HashTable s key val)
new = HashTable s (Some key) Any -> HashTable s key val
forall k s (key :: k -> *) (val :: k -> *).
HashTable s (Some key) Any -> HashTable s key val
HashTable (HashTable s (Some key) Any -> HashTable s key val)
-> ST s (HashTable s (Some key) Any) -> ST s (HashTable s key val)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ST s (HashTable s (Some key) Any)
forall s k v. ST s (HashTable s k v)
H.new

-- | Create a new empty table to hold 'n' elements.
newSized :: Int -> ST s (HashTable s k v)
newSized :: Int -> ST s (HashTable s k v)
newSized Int
n = HashTable s (Some k) Any -> HashTable s k v
forall k s (key :: k -> *) (val :: k -> *).
HashTable s (Some key) Any -> HashTable s key val
HashTable (HashTable s (Some k) Any -> HashTable s k v)
-> ST s (HashTable s (Some k) Any) -> ST s (HashTable s k v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> ST s (HashTable s (Some k) Any)
forall s k v. Int -> ST s (HashTable s k v)
H.newSized Int
n

-- | Create a hash table that is a copy of the current one.
clone :: (HashableF key, TestEquality key)
      => HashTable s key val
      -> ST s (HashTable s key val)
clone :: HashTable s key val -> ST s (HashTable s key val)
clone (HashTable HashTable s (Some key) Any
tbl) = do
  -- Create a new table
  HashTable s (Some key) Any
r <- ST s (HashTable s (Some key) Any)
forall s k v. ST s (HashTable s k v)
H.new
  -- Insert existing elements in
  ((Some key, Any) -> ST s ())
-> HashTable s (Some key) Any -> ST s ()
forall k v s b. ((k, v) -> ST s b) -> HashTable s k v -> ST s ()
H.mapM_ ((Some key -> Any -> ST s ()) -> (Some key, Any) -> ST s ()
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry (HashTable s (Some key) Any -> Some key -> Any -> ST s ()
forall k s v.
(Eq k, Hashable k) =>
HashTable s k v -> k -> v -> ST s ()
H.insert HashTable s (Some key) Any
r)) HashTable s (Some key) Any
tbl
  -- Return table
  HashTable s key val -> ST s (HashTable s key val)
forall (m :: * -> *) a. Monad m => a -> m a
return (HashTable s key val -> ST s (HashTable s key val))
-> HashTable s key val -> ST s (HashTable s key val)
forall a b. (a -> b) -> a -> b
$! HashTable s (Some key) Any -> HashTable s key val
forall k s (key :: k -> *) (val :: k -> *).
HashTable s (Some key) Any -> HashTable s key val
HashTable HashTable s (Some key) Any
r

-- | Lookup value of key in table.
lookup :: (HashableF key, TestEquality key)
       => HashTable s key val
       -> key tp
       -> ST s (Maybe (val tp))
lookup :: HashTable s key val -> key tp -> ST s (Maybe (val tp))
lookup (HashTable HashTable s (Some key) Any
h) key tp
k = (Any -> val tp) -> Maybe Any -> Maybe (val tp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Any -> val tp
forall a b. a -> b
unsafeCoerce (Maybe Any -> Maybe (val tp))
-> ST s (Maybe Any) -> ST s (Maybe (val tp))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> HashTable s (Some key) Any -> Some key -> ST s (Maybe Any)
forall k s v.
(Eq k, Hashable k) =>
HashTable s k v -> k -> ST s (Maybe v)
H.lookup HashTable s (Some key) Any
h (key tp -> Some key
forall k (f :: k -> *) (x :: k). f x -> Some f
Some key tp
k)
{-# INLINE lookup #-}

-- | Insert new key and value mapping into table.
insert :: (HashableF key, TestEquality key)
       => HashTable s (key :: k -> Type) (val :: k -> Type)
       -> key tp
       -> val tp
       -> ST s ()
insert :: HashTable s key val -> key tp -> val tp -> ST s ()
insert (HashTable HashTable s (Some key) Any
h) key tp
k val tp
v = HashTable s (Some key) Any -> Some key -> Any -> ST s ()
forall k s v.
(Eq k, Hashable k) =>
HashTable s k v -> k -> v -> ST s ()
H.insert HashTable s (Some key) Any
h (key tp -> Some key
forall k (f :: k -> *) (x :: k). f x -> Some f
Some key tp
k) (val tp -> Any
forall a b. a -> b
unsafeCoerce val tp
v)

-- | Return true if the key is in the hash table.
member :: (HashableF key, TestEquality key)
       => HashTable s (key :: k -> Type) (val :: k -> Type)
       -> key (tp :: k)
       -> ST s Bool
member :: HashTable s key val -> key tp -> ST s Bool
member (HashTable HashTable s (Some key) Any
h) key tp
k = Maybe Any -> Bool
forall a. Maybe a -> Bool
isJust (Maybe Any -> Bool) -> ST s (Maybe Any) -> ST s Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> HashTable s (Some key) Any -> Some key -> ST s (Maybe Any)
forall k s v.
(Eq k, Hashable k) =>
HashTable s k v -> k -> ST s (Maybe v)
H.lookup HashTable s (Some key) Any
h (key tp -> Some key
forall k (f :: k -> *) (x :: k). f x -> Some f
Some key tp
k)

-- | Delete an element from the hash table.
delete :: (HashableF key, TestEquality key)
       => HashTable s (key :: k -> Type) (val :: k -> Type)
       -> key (tp :: k)
       -> ST s ()
delete :: HashTable s key val -> key tp -> ST s ()
delete (HashTable HashTable s (Some key) Any
h) key tp
k = HashTable s (Some key) Any -> Some key -> ST s ()
forall k s v. (Hashable k, Eq k) => HashTable s k v -> k -> ST s ()
H.delete HashTable s (Some key) Any
h (key tp -> Some key
forall k (f :: k -> *) (x :: k). f x -> Some f
Some key tp
k)

clear :: (HashableF key, TestEquality key)
      => HashTable s (key :: k -> Type) (val :: k -> Type) -> ST s ()
clear :: HashTable s key val -> ST s ()
clear (HashTable HashTable s (Some key) Any
h) = ((Some key, Any) -> ST s ())
-> HashTable s (Some key) Any -> ST s ()
forall k v s b. ((k, v) -> ST s b) -> HashTable s k v -> ST s ()
H.mapM_ (\(Some key
k,Any
_) -> HashTable s (Some key) Any -> Some key -> ST s ()
forall k s v. (Hashable k, Eq k) => HashTable s k v -> k -> ST s ()
H.delete HashTable s (Some key) Any
h Some key
k) HashTable s (Some key) Any
h