Data.Map.updateAt problem

[Christian Maeder]

Andriy Palamarchuk schrieb:
> I'm currently improving Data.Map documentation.
> Behavior of updateAt in this module differs from what I would expect.
> 
> The documentation says:
> 
>      updateAt :: (k -> a -> Maybe a) -> Int -> Map k a -> Map k a
>      O(log n). Update the element at index. Calls error when an invalid index is used.
> 
> So I expected this code:
> 
>     let m = fromList [(3,"b"),(5,"a")]
>     let f k a = Just "x"
>     updateAt f 1 m
> 
> to return
> 
>     fromList [(3,"b"),(5,"x")]
> 
> but it returns
> 
>     fromList [(5,"x")]
> 
> deleting the key 3.
> I tested this using ghci 6.6.
> 
> Is this a bug or I missed something here?

This is a bug, the tree is not properly reconstructed (in the LT and GT
cases). "updateAt f 0 m" works as expected, because the left tree
happens to be empty and updateAt goes into the EQ case. "deleteAt 1 m"
is also wrong.

Cheers Christian

The code looks as follows:

 -- | /O(log n)/. Update the element at /index/. Calls 'error' when an
-- invalid index is used.
updateAt :: (k -> a -> Maybe a) -> Int -> Map k a -> Map k a
updateAt f i Tip  = error "Map.updateAt: index out of range"
updateAt f i (Bin sx kx x l r)
  = case compare i sizeL of
      LT -> updateAt f i l
      GT -> updateAt f (i-sizeL-1) r
      EQ -> case f kx x of
              Just x' -> Bin sx kx x' l r
              Nothing -> glue l r
  where
    sizeL = size l

-- | /O(log n)/. Delete the element at /index/.
-- Defined as (@'deleteAt' i map = 'updateAt' (\k x -> 'Nothing') i map@).
deleteAt :: Int -> Map k a -> Map k a
deleteAt i map
  = updateAt (\k x -> Nothing) i map