98 lines
3.1 KiB
Haskell
98 lines
3.1 KiB
Haskell
module HAMT where
|
|
|
|
import Data.Bits
|
|
import Data.Word
|
|
import Prelude hiding (lookup, (++))
|
|
|
|
import Data.Vector (Vector, singleton, unsafeDrop, unsafeIndex,
|
|
unsafeTake, unsafeUpdate, (++))
|
|
|
|
type Key = Word
|
|
type Bitmap = Word
|
|
type Shift = Int
|
|
type Subkey = Int -- we need to use this to do shifts, so an Int it is
|
|
|
|
data HAMT a = Empty
|
|
| BitmapIndexed {-# UNPACK #-} !Bitmap !(Vector (HAMT a))
|
|
| Leaf {-# UNPACK #-} !Key a
|
|
| Full {-# UNPACK #-} !(Vector (HAMT a))
|
|
deriving (Show)
|
|
|
|
|
|
-- These architecture dependent constants
|
|
|
|
bitsPerSubkey :: Int
|
|
bitsPerSubkey = floor . logBase 2 . fromIntegral . bitSize $ (undefined :: Word)
|
|
|
|
subkeyMask :: Bitmap
|
|
subkeyMask = 1 `shiftL` bitsPerSubkey - 1
|
|
|
|
maskIndex :: Bitmap -> Bitmap -> Int
|
|
maskIndex b m = popCount (b .&. (m - 1))
|
|
|
|
mask :: Key -> Shift -> Bitmap
|
|
mask k s = shiftL 1 (subkey k s)
|
|
|
|
{-# INLINE subkey #-}
|
|
subkey :: Key -> Shift -> Int
|
|
subkey k s = fromIntegral $ shiftR k s .&. subkeyMask
|
|
|
|
empty :: HAMT a
|
|
empty = Empty
|
|
|
|
lookup :: Key -> HAMT a -> Maybe a
|
|
lookup k t = lookup' k 0 t
|
|
|
|
lookup' :: Key -> Shift -> HAMT a -> Maybe a
|
|
lookup' k s t
|
|
= case t of
|
|
Empty -> Nothing
|
|
Leaf kx x
|
|
| k == kx -> Just x
|
|
| otherwise -> Nothing
|
|
BitmapIndexed b v ->
|
|
let m = mask k s in
|
|
if b .&. m == 0
|
|
then Nothing
|
|
else lookup' k (s+bitsPerSubkey) (unsafeIndex v (maskIndex b m))
|
|
Full v -> lookup' k (s+bitsPerSubkey) (unsafeIndex v (subkey k s))
|
|
|
|
insert :: Key -> a -> HAMT a -> HAMT a
|
|
insert k v t = insert' k 0 v t
|
|
|
|
insert' :: Key -> Shift -> a -> HAMT a -> HAMT a
|
|
insert' kx s x t
|
|
= case t of
|
|
Empty -> Leaf kx x
|
|
Leaf ky y
|
|
| ky == kx -> Leaf kx x
|
|
| otherwise ->
|
|
insert' kx s x $ BitmapIndexed (mask ky s) (singleton t)
|
|
BitmapIndexed b v -> {-# SCC "i-Bitmap" #-}
|
|
let m = mask kx s
|
|
i = maskIndex b m in
|
|
if b .&. m == 0
|
|
then let l = Leaf kx x
|
|
v' = unsafeTake i v ++ singleton l ++ unsafeDrop i v
|
|
b' = b .|. m in
|
|
if b' == 0xFFFFFFFF
|
|
then Full v'
|
|
else BitmapIndexed b' v'
|
|
else {-# SCC "i-Bitmap-conflict" #-}
|
|
let st = unsafeIndex v i
|
|
st' = insert' kx (s+bitsPerSubkey) x st
|
|
v' = {-# SCC "i-Bitmap-update" #-}
|
|
unsafeUpdate v (singleton (i, st'))
|
|
in BitmapIndexed b v'
|
|
Full v -> {-# SCC "i-Full" #-}
|
|
let i = subkey kx s
|
|
st = unsafeIndex v i
|
|
st' = insert' kx (s+bitsPerSubkey) x st
|
|
v' = {-# SCC "i-Full-update" #-}
|
|
unsafeUpdate v (singleton (i, st'))
|
|
in Full v'
|
|
|
|
-- lazy, but that's not necessarily a bad thing
|
|
fromList :: [(Key, a)] -> HAMT a
|
|
fromList = foldl (flip $ uncurry insert) empty
|