A Haskell implementation of Okasaki and Gill's Fast Mergeable Integer Maps

IntMap.hs

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