Skip to content

Commit 8f7ebc5

Browse files
sjakobisjakobi
authored
Introduce nextShift (#432)
Context: #426
1 parent 19674b5 commit 8f7ebc5

File tree

2 files changed

+66
-60
lines changed

2 files changed

+66
-60
lines changed

‎Data/HashMap/Internal.hs‎

Lines changed: 49 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -118,6 +118,7 @@ module Data.HashMap.Internal
118118
, index
119119
, bitsPerSubkey
120120
, fullBitmap
121+
, nextShift
121122
, sparseIndex
122123
, two
123124
, unionArrayBy
@@ -689,10 +690,10 @@ lookupCont absent present !h0 !k0 !s0 !m0 = go h0 k0 s0 m0
689690
go h k s (BitmapIndexed b v)
690691
| b .&. m == 0 = absent (# #)
691692
| otherwise =
692-
go h k (s+bitsPerSubkey) (A.index v (sparseIndex b m))
693+
go h k (nextShift s) (A.index v (sparseIndex b m))
693694
where m = mask h s
694695
go h k s (Full v) =
695-
go h k (s+bitsPerSubkey) (A.index v (index h s))
696+
go h k (nextShift s) (A.index v (index h s))
696697
go h k _ (Collision hx v)
697698
| h == hx = lookupInArrayCont absent present k v
698699
| otherwise = absent (# #)
@@ -786,15 +787,15 @@ insert' h0 k0 v0 m0 = go h0 k0 v0 0 m0
786787
in bitmapIndexedOrFull (b .|. m) ary'
787788
| otherwise =
788789
let !st = A.index ary i
789-
!st' = go h k x (s+bitsPerSubkey) st
790+
!st' = go h k x (nextShift s) st
790791
in if st' `ptrEq` st
791792
then t
792793
else BitmapIndexed b (A.update ary i st')
793794
where m = mask h s
794795
i = sparseIndex b m
795796
go h k x s t@(Full ary) =
796797
let !st = A.index ary i
797-
!st' = go h k x (s+bitsPerSubkey) st
798+
!st' = go h k x (nextShift s) st
798799
in if st' `ptrEq` st
799800
then t
800801
else Full (update32 ary i st')
@@ -824,13 +825,13 @@ insertNewKey !h0 !k0 x0 !m0 = go h0 k0 x0 0 m0
824825
in bitmapIndexedOrFull (b .|. m) ary'
825826
| otherwise =
826827
let !st = A.index ary i
827-
!st' = go h k x (s+bitsPerSubkey) st
828+
!st' = go h k x (nextShift s) st
828829
in BitmapIndexed b (A.update ary i st')
829830
where m = mask h s
830831
i = sparseIndex b m
831832
go h k x s (Full ary) =
832833
let !st = A.index ary i
833-
!st' = go h k x (s+bitsPerSubkey) st
834+
!st' = go h k x (nextShift s) st
834835
in Full (update32 ary i st')
835836
where i = index h s
836837
go h k x s t@(Collision hy v)
@@ -860,13 +861,13 @@ insertKeyExists !collPos0 !h0 !k0 x0 !m0 = go collPos0 h0 k0 x0 0 m0
860861
in bitmapIndexedOrFull (b .|. m) ary'
861862
| otherwise =
862863
let !st = A.index ary i
863-
!st' = go collPos h k x (s+bitsPerSubkey) st
864+
!st' = go collPos h k x (nextShift s) st
864865
in BitmapIndexed b (A.update ary i st')
865866
where m = mask h s
866867
i = sparseIndex b m
867868
go collPos h k x s (Full ary) =
868869
let !st = A.index ary i
869-
!st' = go collPos h k x (s+bitsPerSubkey) st
870+
!st' = go collPos h k x (nextShift s) st
870871
in Full (update32 ary i st')
871872
where i = index h s
872873
go collPos h k x _s (Collision _hy v)
@@ -903,14 +904,14 @@ unsafeInsert k0 v0 m0 = runST (go h0 k0 v0 0 m0)
903904
return $! bitmapIndexedOrFull (b .|. m) ary'
904905
| otherwise = do
905906
st <- A.indexM ary i
906-
st' <- go h k x (s+bitsPerSubkey) st
907+
st' <- go h k x (nextShift s) st
907908
A.unsafeUpdateM ary i st'
908909
return t
909910
where m = mask h s
910911
i = sparseIndex b m
911912
go h k x s t@(Full ary) = do
912913
st <- A.indexM ary i
913-
st' <- go h k x (s+bitsPerSubkey) st
914+
st' <- go h k x (nextShift s) st
914915
A.unsafeUpdateM ary i st'
915916
return t
916917
where i = index h s
@@ -932,7 +933,7 @@ two = go
932933
where
933934
go s h1 k1 v1 h2 t2
934935
| bp1 == bp2 = do
935-
st <- go (s+bitsPerSubkey) h1 k1 v1 h2 t2
936+
st <- go (nextShift s) h1 k1 v1 h2 t2
936937
ary <- A.singletonM st
937938
return $ BitmapIndexed bp1 ary
938939
| otherwise = do
@@ -985,7 +986,7 @@ insertModifying x f k0 m0 = go h0 k0 0 m0
985986
in bitmapIndexedOrFull (b .|. m) ary'
986987
| otherwise =
987988
let !st = A.index ary i
988-
!st' = go h k (s+bitsPerSubkey) st
989+
!st' = go h k (nextShift s) st
989990
ary' = A.update ary i $! st'
990991
in if ptrEq st st'
991992
then t
@@ -994,7 +995,7 @@ insertModifying x f k0 m0 = go h0 k0 0 m0
994995
i = sparseIndex b m
995996
go h k s t@(Full ary) =
996997
let !st = A.index ary i
997-
!st' = go h k (s+bitsPerSubkey) st
998+
!st' = go h k (nextShift s) st
998999
ary' = update32 ary i $! st'
9991000
in if ptrEq st st'
10001001
then t
@@ -1052,14 +1053,14 @@ unsafeInsertWithKey f k0 v0 m0 = runST (go h0 k0 v0 0 m0)
10521053
return $! bitmapIndexedOrFull (b .|. m) ary'
10531054
| otherwise = do
10541055
st <- A.indexM ary i
1055-
st' <- go h k x (s+bitsPerSubkey) st
1056+
st' <- go h k x (nextShift s) st
10561057
A.unsafeUpdateM ary i st'
10571058
return t
10581059
where m = mask h s
10591060
i = sparseIndex b m
10601061
go h k x s t@(Full ary) = do
10611062
st <- A.indexM ary i
1062-
st' <- go h k x (s+bitsPerSubkey) st
1063+
st' <- go h k x (nextShift s) st
10631064
A.unsafeUpdateM ary i st'
10641065
return t
10651066
where i = index h s
@@ -1085,7 +1086,7 @@ delete' h0 k0 m0 = go h0 k0 0 m0
10851086
| b .&. m == 0 = t
10861087
| otherwise =
10871088
let !st = A.index ary i
1088-
!st' = go h k (s+bitsPerSubkey) st
1089+
!st' = go h k (nextShift s) st
10891090
in if st' `ptrEq` st
10901091
then t
10911092
else case st' of
@@ -1104,7 +1105,7 @@ delete' h0 k0 m0 = go h0 k0 0 m0
11041105
i = sparseIndex b m
11051106
go h k s t@(Full ary) =
11061107
let !st = A.index ary i
1107-
!st' = go h k (s+bitsPerSubkey) st
1108+
!st' = go h k (nextShift s) st
11081109
in if st' `ptrEq` st
11091110
then t
11101111
else case st' of
@@ -1141,7 +1142,7 @@ deleteKeyExists !collPos0 !h0 !k0 !m0 = go collPos0 h0 k0 0 m0
11411142
go !_collPos !_h !_k !_s (Leaf _ _) = Empty
11421143
go collPos h k s (BitmapIndexed b ary) =
11431144
let !st = A.index ary i
1144-
!st' = go collPos h k (s+bitsPerSubkey) st
1145+
!st' = go collPos h k (nextShift s) st
11451146
in case st' of
11461147
Empty | A.length ary == 1 -> Empty
11471148
| A.length ary == 2 ->
@@ -1158,7 +1159,7 @@ deleteKeyExists !collPos0 !h0 !k0 !m0 = go collPos0 h0 k0 0 m0
11581159
i = sparseIndex b m
11591160
go collPos h k s (Full ary) =
11601161
let !st = A.index ary i
1161-
!st' = go collPos h k (s+bitsPerSubkey) st
1162+
!st' = go collPos h k (nextShift s) st
11621163
in case st' of
11631164
Empty ->
11641165
let ary' = A.delete ary i
@@ -1202,7 +1203,7 @@ adjust# f k0 m0 = go h0 k0 0 m0
12021203
go h k s t@(BitmapIndexed b ary)
12031204
| b .&. m == 0 = t
12041205
| otherwise = let !st = A.index ary i
1205-
!st' = go h k (s+bitsPerSubkey) st
1206+
!st' = go h k (nextShift s) st
12061207
ary' = A.update ary i $! st'
12071208
in if ptrEq st st'
12081209
then t
@@ -1212,7 +1213,7 @@ adjust# f k0 m0 = go h0 k0 0 m0
12121213
go h k s t@(Full ary) =
12131214
let i = index h s
12141215
!st = A.index ary i
1215-
!st' = go h k (s+bitsPerSubkey) st
1216+
!st' = go h k (nextShift s) st
12161217
ary' = update32 ary i $! st'
12171218
in if ptrEq st st'
12181219
then t
@@ -1459,21 +1460,21 @@ isSubmapOfBy comp !m1 !m2 = go 0 m1 m2
14591460
go s t1@(Collision h1 _) (BitmapIndexed b ls2)
14601461
| b .&. m == 0 = False
14611462
| otherwise =
1462-
go (s+bitsPerSubkey) t1 (A.index ls2 (sparseIndex b m))
1463+
go (nextShift s) t1 (A.index ls2 (sparseIndex b m))
14631464
where m = mask h1 s
14641465

14651466
-- Similar to the previous case we need to traverse l2 at the index for the hash h1.
14661467
go s t1@(Collision h1 _) (Full ls2) =
1467-
go (s+bitsPerSubkey) t1 (A.index ls2 (index h1 s))
1468+
go (nextShift s) t1 (A.index ls2 (index h1 s))
14681469

14691470
-- In cases where the first and second map are BitmapIndexed or Full,
14701471
-- traverse down the tree at the appropriate indices.
14711472
go s (BitmapIndexed b1 ls1) (BitmapIndexed b2 ls2) =
1472-
submapBitmapIndexed (go (s+bitsPerSubkey)) b1 ls1 b2 ls2
1473+
submapBitmapIndexed (go (nextShift s)) b1 ls1 b2 ls2
14731474
go s (BitmapIndexed b1 ls1) (Full ls2) =
1474-
submapBitmapIndexed (go (s+bitsPerSubkey)) b1 ls1 fullBitmap ls2
1475+
submapBitmapIndexed (go (nextShift s)) b1 ls1 fullBitmap ls2
14751476
go s (Full ls1) (Full ls2) =
1476-
submapBitmapIndexed (go (s+bitsPerSubkey)) fullBitmap ls1 fullBitmap ls2
1477+
submapBitmapIndexed (go (nextShift s)) fullBitmap ls1 fullBitmap ls2
14771478

14781479
-- Collision and Full nodes always contain at least two entries. Hence it
14791480
-- cannot be a map of a leaf.
@@ -1559,16 +1560,16 @@ unionWithKey f = go 0
15591560
-- branch vs. branch
15601561
go s (BitmapIndexed b1 ary1) (BitmapIndexed b2 ary2) =
15611562
let b' = b1 .|. b2
1562-
ary' = unionArrayBy (go (s+bitsPerSubkey)) b1 b2 ary1 ary2
1563+
ary' = unionArrayBy (go (nextShift s)) b1 b2 ary1 ary2
15631564
in bitmapIndexedOrFull b' ary'
15641565
go s (BitmapIndexed b1 ary1) (Full ary2) =
1565-
let ary' = unionArrayBy (go (s+bitsPerSubkey)) b1 fullBitmap ary1 ary2
1566+
let ary' = unionArrayBy (go (nextShift s)) b1 fullBitmap ary1 ary2
15661567
in Full ary'
15671568
go s (Full ary1) (BitmapIndexed b2 ary2) =
1568-
let ary' = unionArrayBy (go (s+bitsPerSubkey)) fullBitmap b2 ary1 ary2
1569+
let ary' = unionArrayBy (go (nextShift s)) fullBitmap b2 ary1 ary2
15691570
in Full ary'
15701571
go s (Full ary1) (Full ary2) =
1571-
let ary' = unionArrayBy (go (s+bitsPerSubkey)) fullBitmap fullBitmap
1572+
let ary' = unionArrayBy (go (nextShift s)) fullBitmap fullBitmap
15721573
ary1 ary2
15731574
in Full ary'
15741575
-- leaf vs. branch
@@ -1577,7 +1578,7 @@ unionWithKey f = go 0
15771578
b' = b1 .|. m2
15781579
in bitmapIndexedOrFull b' ary'
15791580
| otherwise = let ary' = A.updateWith' ary1 i $ \st1 ->
1580-
go (s+bitsPerSubkey) st1 t2
1581+
go (nextShift s) st1 t2
15811582
in BitmapIndexed b1 ary'
15821583
where
15831584
h2 = leafHashCode t2
@@ -1588,7 +1589,7 @@ unionWithKey f = go 0
15881589
b' = b2 .|. m1
15891590
in bitmapIndexedOrFull b' ary'
15901591
| otherwise = let ary' = A.updateWith' ary2 i $ \st2 ->
1591-
go (s+bitsPerSubkey) t1 st2
1592+
go (nextShift s) t1 st2
15921593
in BitmapIndexed b2 ary'
15931594
where
15941595
h1 = leafHashCode t1
@@ -1597,20 +1598,20 @@ unionWithKey f = go 0
15971598
go s (Full ary1) t2 =
15981599
let h2 = leafHashCode t2
15991600
i = index h2 s
1600-
ary' = update32With' ary1 i $ \st1 -> go (s+bitsPerSubkey) st1 t2
1601+
ary' = update32With' ary1 i $ \st1 -> go (nextShift s) st1 t2
16011602
in Full ary'
16021603
go s t1 (Full ary2) =
16031604
let h1 = leafHashCode t1
16041605
i = index h1 s
1605-
ary' = update32With' ary2 i $ \st2 -> go (s+bitsPerSubkey) t1 st2
1606+
ary' = update32With' ary2 i $ \st2 -> go (nextShift s) t1 st2
16061607
in Full ary'
16071608

16081609
leafHashCode (Leaf h _) = h
16091610
leafHashCode (Collision h _) = h
16101611
leafHashCode _ = error "leafHashCode"
16111612

16121613
goDifferentHash s h1 h2 t1 t2
1613-
| m1 == m2 = BitmapIndexed m1 (A.singleton $! goDifferentHash (s+bitsPerSubkey) h1 h2 t1 t2)
1614+
| m1 == m2 = BitmapIndexed m1 (A.singleton $! goDifferentHash (nextShift s) h1 h2 t1 t2)
16141615
| m1 < m2 = BitmapIndexed (m1 .|. m2) (A.pair t1 t2)
16151616
| otherwise = BitmapIndexed (m1 .|. m2) (A.pair t2 t1)
16161617
where
@@ -1812,30 +1813,30 @@ intersectionWithKey# f = go 0
18121813
go _ (Collision h1 ls1) (Collision h2 ls2) = intersectionCollisions f h1 h2 ls1 ls2
18131814
-- branch vs. branch
18141815
go s (BitmapIndexed b1 ary1) (BitmapIndexed b2 ary2) =
1815-
intersectionArrayBy (go (s + bitsPerSubkey)) b1 b2 ary1 ary2
1816+
intersectionArrayBy (go (nextShift s)) b1 b2 ary1 ary2
18161817
go s (BitmapIndexed b1 ary1) (Full ary2) =
1817-
intersectionArrayBy (go (s + bitsPerSubkey)) b1 fullBitmap ary1 ary2
1818+
intersectionArrayBy (go (nextShift s)) b1 fullBitmap ary1 ary2
18181819
go s (Full ary1) (BitmapIndexed b2 ary2) =
1819-
intersectionArrayBy (go (s + bitsPerSubkey)) fullBitmap b2 ary1 ary2
1820+
intersectionArrayBy (go (nextShift s)) fullBitmap b2 ary1 ary2
18201821
go s (Full ary1) (Full ary2) =
1821-
intersectionArrayBy (go (s + bitsPerSubkey)) fullBitmap fullBitmap ary1 ary2
1822+
intersectionArrayBy (go (nextShift s)) fullBitmap fullBitmap ary1 ary2
18221823
-- collision vs. branch
18231824
go s (BitmapIndexed b1 ary1) t2@(Collision h2 _ls2)
18241825
| b1 .&. m2 == 0 = Empty
1825-
| otherwise = go (s + bitsPerSubkey) (A.index ary1 i) t2
1826+
| otherwise = go (nextShift s) (A.index ary1 i) t2
18261827
where
18271828
m2 = mask h2 s
18281829
i = sparseIndex b1 m2
18291830
go s t1@(Collision h1 _ls1) (BitmapIndexed b2 ary2)
18301831
| b2 .&. m1 == 0 = Empty
1831-
| otherwise = go (s + bitsPerSubkey) t1 (A.index ary2 i)
1832+
| otherwise = go (nextShift s) t1 (A.index ary2 i)
18321833
where
18331834
m1 = mask h1 s
18341835
i = sparseIndex b2 m1
1835-
go s (Full ary1) t2@(Collision h2 _ls2) = go (s + bitsPerSubkey) (A.index ary1 i) t2
1836+
go s (Full ary1) t2@(Collision h2 _ls2) = go (nextShift s) (A.index ary1 i) t2
18361837
where
18371838
i = index h2 s
1838-
go s t1@(Collision h1 _ls1) (Full ary2) = go (s + bitsPerSubkey) t1 (A.index ary2 i)
1839+
go s t1@(Collision h1 _ls1) (Full ary2) = go (nextShift s) t1 (A.index ary2 i)
18391840
where
18401841
i = index h1 s
18411842
{-# INLINE intersectionWithKey# #-}
@@ -2435,6 +2436,11 @@ fullBitmap :: Bitmap
24352436
fullBitmap = complement (complement 0 `shiftL` maxChildren)
24362437
{-# INLINE fullBitmap #-}
24372438

2439+
-- | Increment a 'Shift' for use at the next deeper level.
2440+
nextShift :: Shift -> Shift
2441+
nextShift s = s + bitsPerSubkey
2442+
{-# INLINE nextShift #-}
2443+
24382444
------------------------------------------------------------------------
24392445
-- Pointer equality
24402446

0 commit comments

Comments
 (0)