{-# LANGUAGE CPP #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TupleSections       #-}
{-# LANGUAGE TypeApplications    #-}

-- | Strict hash table designed for small hash tables
--
-- Currently this is is just a wrapper around a 'Map'. We do not use 'HashMap',
-- since for small hash tables the overhead from having to copy all the small
-- arrays defeats the purpose of having a 'Diff' in the first place.
--
-- Having this as a separate abstraction also allows us to easily change the
-- representation of the 'HashMap' without affecting the rest of the code.
--
-- Intended for qualified import.
--
-- > import Data.Record.Anon.Internal.Util.SmallHashMap (SmallHashMap)
-- > import qualified Data.Record.Anon.Internal.Util.SmallHashMap as HashMap
module Data.Record.Anon.Internal.Util.SmallHashMap (
    SmallHashMap
    -- * Standard operations
  , null
  , empty
  , lookup
  , member
  , insert
  , toList
  , alter
    -- * Non-standard operations
  , alterExisting
  ) where

import Prelude hiding (lookup, null)

import Control.Monad.State
import Data.Bifunctor
import Data.Coerce (coerce)
import Data.Hashable (Hashable(hash))
import Data.Map.Strict (Map)
import Data.Tuple (swap)

import qualified Data.Map.Strict as Map

{-------------------------------------------------------------------------------
  Wrapper to compare keys based on their hash first
-------------------------------------------------------------------------------}

newtype Hashed k = Hashed k
  deriving (Int -> Hashed k -> ShowS
[Hashed k] -> ShowS
Hashed k -> String
(Int -> Hashed k -> ShowS)
-> (Hashed k -> String) -> ([Hashed k] -> ShowS) -> Show (Hashed k)
forall k. Show k => Int -> Hashed k -> ShowS
forall k. Show k => [Hashed k] -> ShowS
forall k. Show k => Hashed k -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Hashed k] -> ShowS
$cshowList :: forall k. Show k => [Hashed k] -> ShowS
show :: Hashed k -> String
$cshow :: forall k. Show k => Hashed k -> String
showsPrec :: Int -> Hashed k -> ShowS
$cshowsPrec :: forall k. Show k => Int -> Hashed k -> ShowS
Show)

#if MIN_VERSION_hashable(1,4,0)
instance Hashable k => Eq (Hashed k) where
#else
instance (Hashable k, Eq k) => Eq (Hashed k) where
#endif
  Hashed k
a == :: Hashed k -> Hashed k -> Bool
== Hashed k
b = [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
and [
      k -> Int
forall a. Hashable a => a -> Int
hash k
a Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== k -> Int
forall a. Hashable a => a -> Int
hash k
b
    ,      k
a k -> k -> Bool
forall a. Eq a => a -> a -> Bool
==      k
b
    ]

instance (Hashable k, Ord k) => Ord (Hashed k) where
  compare :: Hashed k -> Hashed k -> Ordering
compare (Hashed k
a) (Hashed k
b) = [Ordering] -> Ordering
forall a. Monoid a => [a] -> a
mconcat [
        Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (k -> Int
forall a. Hashable a => a -> Int
hash k
a) (k -> Int
forall a. Hashable a => a -> Int
hash k
b)
      , k -> k -> Ordering
forall a. Ord a => a -> a -> Ordering
compare       k
a        k
b
      ]

{-------------------------------------------------------------------------------
  Definition of the HashMap proper
-------------------------------------------------------------------------------}

newtype SmallHashMap k a = Wrap { SmallHashMap k a -> Map (Hashed k) a
unwrap :: Map (Hashed k) a }
  deriving (Int -> SmallHashMap k a -> ShowS
[SmallHashMap k a] -> ShowS
SmallHashMap k a -> String
(Int -> SmallHashMap k a -> ShowS)
-> (SmallHashMap k a -> String)
-> ([SmallHashMap k a] -> ShowS)
-> Show (SmallHashMap k a)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall k a. (Show k, Show a) => Int -> SmallHashMap k a -> ShowS
forall k a. (Show k, Show a) => [SmallHashMap k a] -> ShowS
forall k a. (Show k, Show a) => SmallHashMap k a -> String
showList :: [SmallHashMap k a] -> ShowS
$cshowList :: forall k a. (Show k, Show a) => [SmallHashMap k a] -> ShowS
show :: SmallHashMap k a -> String
$cshow :: forall k a. (Show k, Show a) => SmallHashMap k a -> String
showsPrec :: Int -> SmallHashMap k a -> ShowS
$cshowsPrec :: forall k a. (Show k, Show a) => Int -> SmallHashMap k a -> ShowS
Show)

-- | Cannot derive 'Functor' because the 'Functor' instance for 'Map' is wrong
-- (not strict)
instance Functor (SmallHashMap k) where
  fmap :: (a -> b) -> SmallHashMap k a -> SmallHashMap k b
fmap a -> b
f = Map (Hashed k) b -> SmallHashMap k b
forall k a. Map (Hashed k) a -> SmallHashMap k a
Wrap (Map (Hashed k) b -> SmallHashMap k b)
-> (SmallHashMap k a -> Map (Hashed k) b)
-> SmallHashMap k a
-> SmallHashMap k b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> b) -> Map (Hashed k) a -> Map (Hashed k) b
forall a b k. (a -> b) -> Map k a -> Map k b
Map.map a -> b
f (Map (Hashed k) a -> Map (Hashed k) b)
-> (SmallHashMap k a -> Map (Hashed k) a)
-> SmallHashMap k a
-> Map (Hashed k) b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SmallHashMap k a -> Map (Hashed k) a
forall k a. SmallHashMap k a -> Map (Hashed k) a
unwrap

{-------------------------------------------------------------------------------
  Standard operations
-------------------------------------------------------------------------------}

null :: forall k a. SmallHashMap k a -> Bool
null :: SmallHashMap k a -> Bool
null = (Map (Hashed k) a -> Bool) -> SmallHashMap k a -> Bool
coerce ((Map (Hashed k) a -> Bool) -> SmallHashMap k a -> Bool)
-> (Map (Hashed k) a -> Bool) -> SmallHashMap k a -> Bool
forall a b. (a -> b) -> a -> b
$ Map (Hashed k) a -> Bool
forall k a. Map k a -> Bool
Map.null @(Hashed k) @a

empty :: forall k a. SmallHashMap k a
empty :: SmallHashMap k a
empty = Map (Hashed k) a -> SmallHashMap k a
coerce (Map (Hashed k) a -> SmallHashMap k a)
-> Map (Hashed k) a -> SmallHashMap k a
forall a b. (a -> b) -> a -> b
$ Map (Hashed k) a
forall k a. Map k a
Map.empty @(Hashed k) @a

lookup :: forall k a. (Hashable k, Ord k) => k -> SmallHashMap k a -> Maybe a
lookup :: k -> SmallHashMap k a -> Maybe a
lookup = (Hashed k -> Map (Hashed k) a -> Maybe a)
-> k -> SmallHashMap k a -> Maybe a
coerce ((Hashed k -> Map (Hashed k) a -> Maybe a)
 -> k -> SmallHashMap k a -> Maybe a)
-> (Hashed k -> Map (Hashed k) a -> Maybe a)
-> k
-> SmallHashMap k a
-> Maybe a
forall a b. (a -> b) -> a -> b
$ Ord (Hashed k) => Hashed k -> Map (Hashed k) a -> Maybe a
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup @(Hashed k) @a

member :: forall k a. (Hashable k, Ord k) => k -> SmallHashMap k a -> Bool
member :: k -> SmallHashMap k a -> Bool
member = (Hashed k -> Map (Hashed k) a -> Bool)
-> k -> SmallHashMap k a -> Bool
coerce ((Hashed k -> Map (Hashed k) a -> Bool)
 -> k -> SmallHashMap k a -> Bool)
-> (Hashed k -> Map (Hashed k) a -> Bool)
-> k
-> SmallHashMap k a
-> Bool
forall a b. (a -> b) -> a -> b
$ Ord (Hashed k) => Hashed k -> Map (Hashed k) a -> Bool
forall k a. Ord k => k -> Map k a -> Bool
Map.member @(Hashed k) @a

insert :: forall k a.
     (Hashable k, Ord k)
  => k -> a -> SmallHashMap k a -> SmallHashMap k a
insert :: k -> a -> SmallHashMap k a -> SmallHashMap k a
insert = (Hashed k -> a -> Map (Hashed k) a -> Map (Hashed k) a)
-> k -> a -> SmallHashMap k a -> SmallHashMap k a
coerce ((Hashed k -> a -> Map (Hashed k) a -> Map (Hashed k) a)
 -> k -> a -> SmallHashMap k a -> SmallHashMap k a)
-> (Hashed k -> a -> Map (Hashed k) a -> Map (Hashed k) a)
-> k
-> a
-> SmallHashMap k a
-> SmallHashMap k a
forall a b. (a -> b) -> a -> b
$ Ord (Hashed k) =>
Hashed k -> a -> Map (Hashed k) a -> Map (Hashed k) a
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert @(Hashed k) @a

toList :: forall k a. SmallHashMap k a -> [(k, a)]
toList :: SmallHashMap k a -> [(k, a)]
toList = (Map (Hashed k) a -> [(Hashed k, a)])
-> SmallHashMap k a -> [(k, a)]
coerce ((Map (Hashed k) a -> [(Hashed k, a)])
 -> SmallHashMap k a -> [(k, a)])
-> (Map (Hashed k) a -> [(Hashed k, a)])
-> SmallHashMap k a
-> [(k, a)]
forall a b. (a -> b) -> a -> b
$ Map (Hashed k) a -> [(Hashed k, a)]
forall k a. Map k a -> [(k, a)]
Map.toList @(Hashed k) @a

alter :: forall k a.
     (Hashable k, Ord k)
  => (Maybe a -> Maybe a) -> k -> SmallHashMap k a -> SmallHashMap k a
alter :: (Maybe a -> Maybe a) -> k -> SmallHashMap k a -> SmallHashMap k a
alter = ((Maybe a -> Maybe a)
 -> Hashed k -> Map (Hashed k) a -> Map (Hashed k) a)
-> (Maybe a -> Maybe a)
-> k
-> SmallHashMap k a
-> SmallHashMap k a
coerce (((Maybe a -> Maybe a)
  -> Hashed k -> Map (Hashed k) a -> Map (Hashed k) a)
 -> (Maybe a -> Maybe a)
 -> k
 -> SmallHashMap k a
 -> SmallHashMap k a)
-> ((Maybe a -> Maybe a)
    -> Hashed k -> Map (Hashed k) a -> Map (Hashed k) a)
-> (Maybe a -> Maybe a)
-> k
-> SmallHashMap k a
-> SmallHashMap k a
forall a b. (a -> b) -> a -> b
$ Ord (Hashed k) =>
(Maybe a -> Maybe a)
-> Hashed k -> Map (Hashed k) a -> Map (Hashed k) a
forall k a.
Ord k =>
(Maybe a -> Maybe a) -> k -> Map k a -> Map k a
Map.alter @(Hashed k) @a

{-------------------------------------------------------------------------------
  Non-standard operations
-------------------------------------------------------------------------------}

-- | Alter an existing key
--
-- Returns 'Nothing' if the key does not exist.
--
-- @O(1)@.
alterExisting :: forall k a b.
     (Hashable k, Ord k)
  => k -> (a -> (b, Maybe a)) -> SmallHashMap k a -> Maybe (b, SmallHashMap k a)
alterExisting :: k
-> (a -> (b, Maybe a))
-> SmallHashMap k a
-> Maybe (b, SmallHashMap k a)
alterExisting k
k a -> (b, Maybe a)
f SmallHashMap k a
m
  | SmallHashMap k a -> Bool
forall k a. SmallHashMap k a -> Bool
null SmallHashMap k a
m    = Maybe (b, SmallHashMap k a)
forall a. Maybe a
Nothing
  | Bool
otherwise =
      ((Map (Hashed k) a, b) -> (b, SmallHashMap k a))
-> Maybe (Map (Hashed k) a, b) -> Maybe (b, SmallHashMap k a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Map (Hashed k) a -> SmallHashMap k a)
-> (b, Map (Hashed k) a) -> (b, SmallHashMap k a)
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second Map (Hashed k) a -> SmallHashMap k a
forall k a. Map (Hashed k) a -> SmallHashMap k a
Wrap ((b, Map (Hashed k) a) -> (b, SmallHashMap k a))
-> ((Map (Hashed k) a, b) -> (b, Map (Hashed k) a))
-> (Map (Hashed k) a, b)
-> (b, SmallHashMap k a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Map (Hashed k) a, b) -> (b, Map (Hashed k) a)
forall a b. (a, b) -> (b, a)
swap)
    (Maybe (Map (Hashed k) a, b) -> Maybe (b, SmallHashMap k a))
-> (SmallHashMap k a -> Maybe (Map (Hashed k) a, b))
-> SmallHashMap k a
-> Maybe (b, SmallHashMap k a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Map (Hashed k) a, Maybe b) -> Maybe (Map (Hashed k) a, b)
forall x y. (x, Maybe y) -> Maybe (x, y)
distrib
    ((Map (Hashed k) a, Maybe b) -> Maybe (Map (Hashed k) a, b))
-> (SmallHashMap k a -> (Map (Hashed k) a, Maybe b))
-> SmallHashMap k a
-> Maybe (Map (Hashed k) a, b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (State (Maybe b) (Map (Hashed k) a)
 -> Maybe b -> (Map (Hashed k) a, Maybe b))
-> Maybe b
-> State (Maybe b) (Map (Hashed k) a)
-> (Map (Hashed k) a, Maybe b)
forall a b c. (a -> b -> c) -> b -> a -> c
flip State (Maybe b) (Map (Hashed k) a)
-> Maybe b -> (Map (Hashed k) a, Maybe b)
forall s a. State s a -> s -> (a, s)
runState Maybe b
forall a. Maybe a
Nothing
    (State (Maybe b) (Map (Hashed k) a) -> (Map (Hashed k) a, Maybe b))
-> (SmallHashMap k a -> State (Maybe b) (Map (Hashed k) a))
-> SmallHashMap k a
-> (Map (Hashed k) a, Maybe b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Maybe a -> StateT (Maybe b) Identity (Maybe a))
-> Hashed k
-> Map (Hashed k) a
-> State (Maybe b) (Map (Hashed k) a)
forall (f :: * -> *) k a.
(Functor f, Ord k) =>
(Maybe a -> f (Maybe a)) -> k -> Map k a -> f (Map k a)
Map.alterF Maybe a -> StateT (Maybe b) Identity (Maybe a)
f' (k -> Hashed k
forall k. k -> Hashed k
Hashed k
k)
    (Map (Hashed k) a -> State (Maybe b) (Map (Hashed k) a))
-> (SmallHashMap k a -> Map (Hashed k) a)
-> SmallHashMap k a
-> State (Maybe b) (Map (Hashed k) a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SmallHashMap k a -> Map (Hashed k) a
forall k a. SmallHashMap k a -> Map (Hashed k) a
unwrap
    (SmallHashMap k a -> Maybe (b, SmallHashMap k a))
-> SmallHashMap k a -> Maybe (b, SmallHashMap k a)
forall a b. (a -> b) -> a -> b
$ SmallHashMap k a
m
  where
    f' :: Maybe a -> State (Maybe b) (Maybe a)
    f' :: Maybe a -> StateT (Maybe b) Identity (Maybe a)
f' Maybe a
Nothing  = (Maybe b -> (Maybe a, Maybe b))
-> StateT (Maybe b) Identity (Maybe a)
forall s (m :: * -> *) a. MonadState s m => (s -> (a, s)) -> m a
state ((Maybe b -> (Maybe a, Maybe b))
 -> StateT (Maybe b) Identity (Maybe a))
-> (Maybe b -> (Maybe a, Maybe b))
-> StateT (Maybe b) Identity (Maybe a)
forall a b. (a -> b) -> a -> b
$ \Maybe b
_ -> (Maybe a
forall a. Maybe a
Nothing, Maybe b
forall a. Maybe a
Nothing)
    f' (Just a
a) = (Maybe b -> (Maybe a, Maybe b))
-> StateT (Maybe b) Identity (Maybe a)
forall s (m :: * -> *) a. MonadState s m => (s -> (a, s)) -> m a
state ((Maybe b -> (Maybe a, Maybe b))
 -> StateT (Maybe b) Identity (Maybe a))
-> (Maybe b -> (Maybe a, Maybe b))
-> StateT (Maybe b) Identity (Maybe a)
forall a b. (a -> b) -> a -> b
$ \Maybe b
_ -> (Maybe b, Maybe a) -> (Maybe a, Maybe b)
forall a b. (a, b) -> (b, a)
swap ((Maybe b, Maybe a) -> (Maybe a, Maybe b))
-> (Maybe b, Maybe a) -> (Maybe a, Maybe b)
forall a b. (a -> b) -> a -> b
$ (b -> Maybe b) -> (b, Maybe a) -> (Maybe b, Maybe a)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first b -> Maybe b
forall a. a -> Maybe a
Just (a -> (b, Maybe a)
f a
a)

    distrib :: (x, Maybe y) -> Maybe (x, y)
    distrib :: (x, Maybe y) -> Maybe (x, y)
distrib (x
x, Maybe y
my) = (x
x,) (y -> (x, y)) -> Maybe y -> Maybe (x, y)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe y
my