Skip to content

Instantly share code, notes, and snippets.

@kseo
Created February 14, 2014 04:35
Show Gist options
  • Save kseo/8995875 to your computer and use it in GitHub Desktop.
Save kseo/8995875 to your computer and use it in GitHub Desktop.
A Haskell implementation of Okasaki and Gill's Fast Mergeable Integer Maps
module IntMap where
import Data.Bits
import Prelude hiding (lookup)
-- Fast Mergeable lnteger Maps*
-- http://ittc.ku.edu/~andygill/papers/IntMap98.pdf
-- Little-endian implementation
type Key = Int
data IntMap v =
Empty
| Leaf Int v
| Branch Int Int (IntMap v) (IntMap v) deriving (Show)
-- the first integer is the prefix and the second integer is the branching bit
empty :: IntMap v
empty = Empty
lookup :: Key -> IntMap v -> Maybe v
lookup _ Empty = Nothing
lookup k (Leaf j x) = if j == k then Just x else Nothing
lookup k (Branch _ m t0 t1) = if zeroBit k m then lookup k t0 else lookup k t1
insert :: (v -> v -> v) -> Key -> v -> IntMap v -> IntMap v
insert c k x = ins
where ins Empty = Leaf k x
ins t@(Leaf j y) =
if j == k
then Leaf k $ c x y
else join k (Leaf k x) j t
ins t@(Branch p m t0 t1) =
if matchPrefix k p m
then if zeroBit k m
then Branch p m (ins t0) t1
else Branch p m t0 (ins t1)
else join k (Leaf k x) p t
merge :: (v -> v -> v) -> IntMap v -> IntMap v -> IntMap v
merge c = mrg
where mrg Empty t = t
mrg t Empty = t
mrg (Leaf k x) t = insert c k x t
mrg t (Leaf k x) = insert (flip c) k x t
-- The trees have the same prefix. Merge the subtrees.
mrg (Branch p m s0 s1) (Branch q n t0 t1) | m == n && p == q = Branch p m (mrg s0 t0) (mrg s1 t1)
-- q contains p. Merge t with a subtree of s.
mrg (Branch p m s0 s1) t@(Branch q n _ _) | m < n && matchPrefix q p m =
if zeroBit q m then Branch p m (mrg s0 t) s1
else Branch p m s0 (mrg s1 t)
-- p contains q. Merge s with a subtree of t.
mrg s@(Branch p m _ _) (Branch q n t0 t1) | m > n && matchPrefix p q n =
if zeroBit p n then Branch q n (mrg s t0) t1
else Branch q n t0 (mrg s t1)
-- The prefixes disagree.
mrg s@(Branch p _ _ _) t@(Branch q _ _ _) = join p s q t
-- TODO: implement delete, intersection, difference
join :: Int -> IntMap v -> Int -> IntMap v -> IntMap v
join p0 t0 p1 t1 =
let m = branchingBit p0 p1
in if zeroBit p0 m
then Branch (mask p0 m) m t0 t1
else Branch (mask p0 m) m t1 t0
zeroBit :: Int -> Int -> Bool
zeroBit k m = (k .&. m) == 0
mask :: Int -> Int -> Int
mask k m = k .&. (m - 1)
matchPrefix :: Int -> Int -> Int -> Bool
matchPrefix k p m = mask k m == p
-- Finds the first bit at wich p0 and p1 disagree.
branchingBit :: Int -> Int -> Int
branchingBit p0 p1 = lowestBit $ p0 `xor` p1
lowestBit :: Int -> Int
lowestBit x = x .&. (complement x + 1)
-- A smart constructor to collapse each empty subtree into a single Empty node
br :: Int -> Int -> IntMap v -> IntMap v -> IntMap v
br _ _ Empty Empty = Empty
br _ _ Empty t@(Leaf _ _) = t
br _ _ t@(Leaf _ _) Empty = t
br p m t0 t1 = Branch p m t0 t1
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment