@@ -0,0 +1,155 @@
1+module Data.HashMap.Monadic (MEq (..), insert, lookup) where
2+
3+import Control.Monad.ST (runST)
4+import Data.Bits
5+import Data.Functor ((<&>))
6+import Data.HashMap.Internal
7+ ( Hash,
8+ HashMap (..),
9+ Leaf (..),
10+ bitmapIndexedOrFull,
11+ collision,
12+ hash,
13+ index,
14+ mask,
15+ nextShift,
16+ ptrEq,
17+ sparseIndex,
18+ two,
19+ update32,
20+ )
21+import Data.HashMap.Internal.Array qualified as A
22+import Data.Hashable
23+import Prelude hiding (lookup)
24+
25+class MEq x m where
26+ (===) :: x -> x -> m Bool
27+
28+insert :: (MEq k m, Hashable k, Monad m) => k -> v -> HashMap k v -> m (HashMap k v)
29+insert k v m = insert' (Data.HashMap.Internal.hash k) k v m
30+
31+insert' :: forall m k v. (MEq k m, Monad m) => Hash -> k -> v -> HashMap k v -> m (HashMap k v)
32+insert' h0 k0 v0 m0 = go h0 k0 v0 0 m0
33+ where
34+ -- go :: _ -> _ -> _ -> _ -> _ -> m _
35+ go !h !k x !_ Empty = pure @m $ Leaf h (L k x)
36+ go h k x s t@(Leaf hy l@(L ky y))
37+ | hy == h = do
38+ ky === k >>= \eq ->
39+ if eq
40+ then
41+ if x `ptrEq` y
42+ then pure @m t
43+ else pure @m $ Leaf h (L k x)
44+ else pure $ collision h l (L k x)
45+ | otherwise = pure $ runST (two s h k x hy t)
46+ go h k x s t@(BitmapIndexed b ary)
47+ | b .&. m == 0 =
48+ let !ary' = A.insert ary i $! Leaf h (L k x)
49+ in pure $ bitmapIndexedOrFull (b .|. m) ary'
50+ | otherwise = do
51+ let !st = A.index ary i
52+ !st' <- go h k x (nextShift s) st
53+ if st' `ptrEq` st
54+ then pure t
55+ else pure $ BitmapIndexed b (A.update ary i st')
56+ where
57+ m = mask h s
58+ i = sparseIndex b m
59+ go h k x s t@(Full ary) = do
60+ let !st = A.index ary i
61+ !st' <- go h k x (nextShift s) st
62+ if st' `ptrEq` st
63+ then pure t
64+ else pure $ Full (update32 ary i st')
65+ where
66+ i = index h s
67+ go h k x s t@(Collision hy v)
68+ | h == hy = Collision h <$> (updateOrSnocWith (\a _ -> (# a #)) k x v)
69+ | otherwise = go h k x s $ BitmapIndexed (mask hy s) (A.singleton t)
70+
71+updateOrSnocWith ::
72+ (MEq k m, Monad m) =>
73+ (v -> v -> (# v #)) ->
74+ k ->
75+ v ->
76+ A.Array (Leaf k v) ->
77+ m (A.Array (Leaf k v))
78+updateOrSnocWith f = updateOrSnocWithKey (const f)
79+{-# INLINEABLE updateOrSnocWith #-}
80+
81+updateOrSnocWithKey ::
82+ (MEq k m, Monad m) =>
83+ (k -> v -> v -> (# v #)) ->
84+ k ->
85+ v ->
86+ A.Array (Leaf k v) ->
87+ m (A.Array (Leaf k v))
88+updateOrSnocWithKey f k0 v0 ary0 = go k0 v0 ary0 0 (A.length ary0)
89+ where
90+ go !k v !ary !i !n
91+ -- Not found, append to the end.
92+ | i >= n = pure $ A.snoc ary $ L k v
93+ | otherwise =
94+ let L kx y = A.index ary i
95+ (# v2 #) = f k v y
96+ in k === kx >>= \eq ->
97+ if eq
98+ then pure $ A.update ary i (L k v2)
99+ else go k v ary (i + 1) n
100+
101+lookup :: (MEq k m, Hashable k, Monad m) => k -> HashMap k v -> m (Maybe v)
102+lookup k m = lookupCont Nothing (\v _i -> Just v) (Data.HashMap.Internal.hash k) k 0 m
103+
104+lookupCont ::
105+ forall m r k v.
106+ (MEq k m, Monad m) =>
107+ r -> -- Absent continuation
108+ (v -> Int -> r) -> -- Present continuation
109+ Hash -> -- The hash of the key
110+ k ->
111+ Int -> -- The offset of the subkey in the hash.
112+ HashMap k v ->
113+ m r
114+lookupCont absent present !h0 !k0 !s0 !m0 = go h0 k0 s0 m0
115+ where
116+ go :: Hash -> k -> Int -> HashMap k v -> m r
117+ go !_ !_ !_ Empty = pure absent
118+ go h k _ (Leaf hx (L kx x))
119+ | h == hx =
120+ k === kx <&> \eq ->
121+ if eq
122+ then present x (-1)
123+ else absent
124+ | otherwise = pure absent
125+ go h k s (BitmapIndexed b v)
126+ | b .&. m == 0 = pure absent
127+ | otherwise =
128+ go h k (nextShift s) (A.index v (sparseIndex b m))
129+ where
130+ m = mask h s
131+ go h k s (Full v) =
132+ go h k (nextShift s) (A.index v (index h s))
133+ go h k _ (Collision hx v)
134+ | h == hx = lookupInArrayCont absent present k v
135+ | otherwise = pure absent
136+
137+lookupInArrayCont ::
138+ forall m r k v.
139+ (MEq k m, Monad m) =>
140+ r ->
141+ (v -> Int -> r) ->
142+ k ->
143+ A.Array (Leaf k v) ->
144+ m r
145+lookupInArrayCont absent present k0 ary0 = go k0 ary0 0 (A.length ary0)
146+ where
147+ go :: k -> A.Array (Leaf k v) -> Int -> Int -> m r
148+ go !k !ary !i !n
149+ | i >= n = pure absent
150+ | otherwise = case A.index ary i of
151+ (L kx v) ->
152+ k === kx >>= \eq ->
153+ if eq
154+ then pure $ present v i
155+ else go k ary (i + 1) n