{-# 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
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 = forall (t :: * -> *). Foldable t => t Bool -> Bool
and [
      forall a. Hashable a => a -> Int
hash k
a forall a. Eq a => a -> a -> Bool
== forall a. Hashable a => a -> Int
hash k
b
    ,      k
a 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) = forall a. Monoid a => [a] -> a
mconcat [
        forall a. Ord a => a -> a -> Ordering
compare (forall a. Hashable a => a -> Int
hash k
a) (forall a. Hashable a => a -> Int
hash k
b)
      , forall a. Ord a => a -> a -> Ordering
compare       k
a        k
b
      ]

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

newtype SmallHashMap k a = Wrap { forall k a. SmallHashMap k a -> Map (Hashed k) a
unwrap :: Map (Hashed k) a }
  deriving (Int -> SmallHashMap k a -> ShowS
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 :: forall a b. (a -> b) -> SmallHashMap k a -> SmallHashMap k b
fmap a -> b
f = forall k a. Map (Hashed k) a -> SmallHashMap k a
Wrap forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b k. (a -> b) -> Map k a -> Map k b
Map.map a -> b
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. SmallHashMap k a -> Map (Hashed k) a
unwrap

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

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

empty :: forall k a. SmallHashMap k a
empty :: forall k a. SmallHashMap k a
empty = coerce :: forall a b. Coercible a b => a -> b
coerce forall a b. (a -> b) -> a -> b
$ 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 :: forall k a. (Hashable k, Ord k) => k -> SmallHashMap k a -> Maybe a
lookup = coerce :: forall a b. Coercible a b => a -> b
coerce forall a b. (a -> b) -> a -> b
$ 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 :: forall k a. (Hashable k, Ord k) => k -> SmallHashMap k a -> Bool
member = coerce :: forall a b. Coercible a b => a -> b
coerce forall a b. (a -> b) -> a -> b
$ 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 :: forall k a.
(Hashable k, Ord k) =>
k -> a -> SmallHashMap k a -> SmallHashMap k a
insert = coerce :: forall a b. Coercible a b => a -> b
coerce forall a b. (a -> b) -> a -> b
$ 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 :: forall k a. SmallHashMap k a -> [(k, a)]
toList = coerce :: forall a b. Coercible a b => a -> b
coerce forall a b. (a -> b) -> a -> b
$ 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 :: forall k a.
(Hashable k, Ord k) =>
(Maybe a -> Maybe a) -> k -> SmallHashMap k a -> SmallHashMap k a
alter = coerce :: forall a b. Coercible a b => a -> b
coerce forall a b. (a -> b) -> a -> b
$ 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 :: forall k a b.
(Hashable k, Ord k) =>
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
  | forall k a. SmallHashMap k a -> Bool
null SmallHashMap k a
m    = forall a. Maybe a
Nothing
  | Bool
otherwise =
      forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second forall k a. Map (Hashed k) a -> SmallHashMap k a
Wrap forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> (b, a)
swap)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall x y. (x, Maybe y) -> Maybe (x, y)
distrib
    forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b c. (a -> b -> c) -> b -> a -> c
flip forall s a. State s a -> s -> (a, s)
runState forall a. Maybe a
Nothing
    forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 -> State (Maybe b) (Maybe a)
f' (forall k. k -> Hashed k
Hashed k
k)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. SmallHashMap k a -> Map (Hashed k) a
unwrap
    forall a b. (a -> b) -> a -> b
$ SmallHashMap k a
m
  where
    f' :: Maybe a -> State (Maybe b) (Maybe a)
    f' :: Maybe a -> State (Maybe b) (Maybe a)
f' Maybe a
Nothing  = forall s (m :: * -> *) a. MonadState s m => (s -> (a, s)) -> m a
state forall a b. (a -> b) -> a -> b
$ \Maybe b
_ -> (forall a. Maybe a
Nothing, forall a. Maybe a
Nothing)
    f' (Just a
a) = forall s (m :: * -> *) a. MonadState s m => (s -> (a, s)) -> m a
state forall a b. (a -> b) -> a -> b
$ \Maybe b
_ -> forall a b. (a, b) -> (b, a)
swap forall a b. (a -> b) -> a -> b
$ forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first forall a. a -> Maybe a
Just (a -> (b, Maybe a)
f a
a)

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