content: -- Copyright (c) 1982-1999 Lennart Augustsson, Thomas Johnsson
  -- See LICENSE for the full license.
  --
  module IntMap (
  	IntMap, 
  	empty, singleton, union, unionMany, add, (//), addKeep,
  	-- union_C, unionMany_C, addMany_C,
  	add_C,
  	-- intersect, 
  	delete, deleteMany, 
  	--minus, 
  	foldr, imap, filter,
  	-- partition, foldl,
  	toList, fromList,
  	length,
  	null, isSingleton,
  	-- intersecting, subset
  	elems, indices,
  	(!),
  	lookup, lookupWithDefault --, lookupWithContinuation
  	) where
  
  -- @@ Mapping from Int to any type.  Similar to an array with Int index, but
  -- @@ without any bounds on the index.
  
  import Prelude hiding (foldr,length,lookup,filter,null)
  import qualified Prelude
  
  data IntMap a = Nil | Leaf !Int a | Fork !(IntMap a) !(IntMap a)
  instance (Show a) => Show (IntMap a) where
  	--showsType _ = showString "IntMap a"
  	showsPrec _ Nil = showString "{}"
  	showsPrec _ s = showString "{" . f (toList s) . showString "}"
  		where f [x] = g x
  		      f (x:xs) = g x . showString ", " . f xs
  		      g (i, r) = shows i . showString "->" . shows r
  
  instance (Eq a) => Eq (IntMap a) where
  	Nil == Nil = True
  	Leaf x y == Leaf x' y' = x == x' && y == y'
  	Fork l r == Fork l' r' = l == l' && r == r'
  	_ == _ = False
  
  empty :: IntMap a
  empty = Nil
  
  singleton :: (Int, a) -> IntMap a
  singleton (x, y) = Leaf x y
  
  null :: IntMap a -> Bool
  null Nil = True
  null (Leaf _ _) = False
  null (Fork _ _) = False
  
  add :: (Int, a) -> IntMap a -> IntMap a
  add (x, y) t = add' x y t
  add' :: Int -> a -> IntMap a -> IntMap a
  add' x y Nil = Leaf x y
  add' x y (Leaf x' y') =
  	if x == x' then
  	    Leaf x y
  	else
  	    add' x y (add' x' y' (Fork Nil Nil))
  add' x y (Fork l r) =
  	if odd x then
  	    Fork l (add' (x `div` 2) y r)
  	else
  	    Fork (add' (x `div` 2) y l) r
  
  -- similar to add, but does not overwrite the old contents
  addKeep :: (Int, a) -> IntMap a -> IntMap a
  addKeep (x, y) t = addKeep' x y t
  addKeep' :: Int -> a -> IntMap a -> IntMap a
  addKeep' x y Nil = Leaf x y
  addKeep' x y t@(Leaf x' y') =
  	if x == x' then
  	    t
  	else
  	    addKeep' x y (addKeep' x' y' (Fork Nil Nil))
  addKeep' x y (Fork l r) =
  	if odd x then
  	    Fork l (addKeep' (x `div` 2) y r)
  	else
  	    Fork (addKeep' (x `div` 2) y l) r
  
  lookupWithDefault :: IntMap a -> a -> Int -> a
  lookupWithDefault Nil d x = if x==x then d else d			-- force it to be strict in x
  lookupWithDefault (Leaf x' y) d x = if x == x' then y else d
  lookupWithDefault (Fork l r) d x =
  	if odd x then
  	    lookupWithDefault r d (x `div` 2)
  	else
  	    lookupWithDefault l d (x `div` 2)
  
  lookup :: Int -> IntMap a -> Maybe a
  lookup x Nil = Nothing
  lookup x (Leaf x' y) = if x == x' then Just y else Nothing
  lookup x (Fork l r) =
  	if odd x then
  	    lookup (x `div` 2) r
  	else
  	    lookup (x `div` 2) l
  
  (!) :: IntMap a -> Int -> a
  t ! x = case lookup x t of Nothing -> error "IntMap.!: index not found"; Just y -> y
  
  
  union :: IntMap a -> IntMap a -> IntMap a
  union Nil t = t
  union (Leaf x y) t = add' x y t
  union t Nil = t
  union t (Leaf x y) = addKeep' x y t
  union (Fork l r) (Fork l' r') = Fork (union l l') (union r r')
  
  unionMany :: [IntMap a] -> IntMap a
  unionMany = Prelude.foldr union empty
  
  fromList :: [(Int, a)] -> IntMap a
  fromList xs = Prelude.foldr (\ (x,y) -> \ m -> add' x y m) empty xs
  
  toList :: IntMap a -> [(Int, a)]
  toList t = foldr (:) [] t
  {-
  toList :: IntMap a -> [(Int, a)]
  toList Nil = []
  toList (Leaf x y) = [(x, y)]
  toList (Fork l r) = [(2*x, y) | (x, y) <- toList l] ++ [ (2*x+1, y) | (x, y) <- toList r]
  -}
  
  elems :: IntMap a -> [a]
  elems = Prelude.map snd . toList
  
  indices :: IntMap a -> [Int]
  indices = Prelude.map fst . toList
  
  length :: IntMap a -> Int
  length Nil = 0
  length (Leaf _ _) = 1
  length (Fork l r) = length l + length r
  
  isSingleton :: IntMap a -> Bool
  isSingleton t = length t == 1
  
  add_C :: (a->a->a) -> (Int, a) -> IntMap a -> IntMap a
  add_C comb (x, y) t = add_C' comb x y t
  add_C' :: (a->a->a) -> Int -> a -> IntMap a -> IntMap a
  add_C' comb x y Nil = Leaf x y
  add_C' comb x y (Leaf x' y') =
  	if x == x' then
  	    Leaf x (comb y y')
  	else
  	    add_C' comb x y (add_C' comb x' y' (Fork Nil Nil))
  add_C' comb x y (Fork l r) =
  	if odd x then
  	    Fork l (add_C' comb (x `div` 2) y r)
  	else
  	    Fork (add_C' comb (x `div` 2) y l) r
  
  (//) :: IntMap a -> [(Int, a)] -> IntMap a
  t // [] = t
  t // ((x,y):xys) = add' x y t // xys
  
  instance Functor IntMap where
      fmap f Nil = Nil
      fmap f (Leaf x y) = Leaf x (f y)
      fmap f (Fork l r) = Fork (fmap f l) (fmap f r)
  
  delete :: Int -> IntMap a -> IntMap a
  delete x Nil           = Nil
  delete x t@(Leaf x' y) = if x == x' then Nil else t
  delete x (Fork l r)    = if odd x then
                               fork l (delete (x `div` 2) r)
  			 else
  			     fork (delete (x `div` 2) l) r
  
  deleteMany :: [Int] -> IntMap a -> IntMap a
  deleteMany is s = Prelude.foldr delete s is
  
  fork Nil        Nil        = Nil
  fork Nil        (Leaf x y) = Leaf (x*2+1) y
  fork (Leaf x y) Nil        = Leaf (x*2) y
  fork l          r          = Fork l r
  
  foldr :: ((Int, a) -> b -> b) -> b -> IntMap a -> b
  foldr f z Nil = z
  foldr f z (Leaf x y) = f (x,y) z
  foldr f z (Fork l r) = foldr g (foldr h z r) l
      where g (x,y) z = f (2*x,y) z
            h (x,y) z = f (2*x+1,y) z
  
  imap :: ((Int, a) -> (Int, b)) -> IntMap a -> IntMap b
  imap f t = foldr (add . f) empty t
  
  filter :: ((Int, a) -> Bool) -> IntMap a -> IntMap a
  filter p t = foldr (\x l -> if p x then add x l else l) empty t
  