Ticket #4887: MapOps.hs

-- Element-wise operations from Data.Map re-defined using Locations.

module MapOps where

import Data.Map (Map, empty, null,
        Location, key, before, after, assign, clear,
        search, index, minLocation, maxLocation)

import Prelude hiding (null)

-- ** Insertion

insert :: Ord k => k -> a -> Map k a -> Map k a
insert k v m = assign v (snd (search k m))

insertWith :: Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
insertWith f k v m = case search k m of
        (Nothing, loc) -> assign v loc
        (Just oldv, loc) -> assign (f v oldv) loc

insertWith' :: Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
insertWith' f k v m = case search k m of
        (Nothing, loc) -> assign v loc
        (Just oldv, loc) -> flip assign loc $! f v oldv

insertWithKey :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a -> Map k a
insertWithKey f k v m = case search k m of
        (Nothing, loc) -> assign v loc
        (Just oldv, loc) -> assign (f (key loc) v oldv) loc

insertWithKey' :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a -> Map k a
insertWithKey' f k v m = case search k m of
        (Nothing, loc) -> assign v loc
        (Just oldv, loc) -> flip assign loc $! f (key loc) v oldv

insertLookupWithKey :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a -> (Maybe a, Map k a)
insertLookupWithKey f k v m = case search k m of
        (Nothing, loc) -> (Nothing, assign v loc)
        (Just oldv, loc) -> (Just oldv, assign (f (key loc) v oldv) loc)

insertLookupWithKey' :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a -> (Maybe a, Map k a)
insertLookupWithKey' f k v m = case search k m of
        (Nothing, loc) -> (Nothing, assign v loc)
        (Just oldv, loc) -> v' `seq` (Just oldv, assign v' loc)
          where v' = f (key loc) v oldv

-- ** Delete/Update

delete :: Ord k => k -> Map k a -> Map k a
delete k m = clear (snd (search k m))

adjust :: Ord k => (a -> a) -> k -> Map k a -> Map k a
adjust f k m = case search k m of
        (Nothing, _) -> m
        (Just v, loc) -> assign (f v) loc

adjustWithKey :: Ord k => (k -> a -> a) -> k -> Map k a -> Map k a
adjustWithKey f k m = case search k m of
        (Nothing, _) -> m
        (Just v, loc) -> assign (f (key loc) v) loc

update :: Ord k => (a -> Maybe a) -> k -> Map k a -> Map k a
update f k m = case search k m of
        (Nothing, _) -> m
        (Just v, loc) -> maybeAssign (f v) loc

updateWithKey :: Ord k => (k -> a -> Maybe a) -> k -> Map k a -> Map k a
updateWithKey f k m = case search k m of
        (Nothing, _) -> m
        (Just v, loc) -> maybeAssign (f (key loc) v) loc

updateLookupWithKey :: Ord k => (k -> a -> Maybe a) -> k -> Map k a -> (Maybe a, Map k a)
updateLookupWithKey f k m = case search k m of
        (Nothing, _) -> (Nothing, m)
        (Just v, loc) -> (Just v, maybeAssign (f (key loc) v) loc)

alter :: Ord k => (Maybe a -> Maybe a) -> k -> Map k a -> Map k a
alter f k m = case search k m of
        (v, loc) -> maybeAssign (f v) loc

-- * Filter

split :: Ord k => k -> Map k a -> (Map k a, Map k a)
split k m = case search k m of
        (_, loc) -> (before loc, after loc)

splitLookup :: Ord k => k -> Map k a -> (Map k a, Maybe a, Map k a)
splitLookup k m = case search k m of
        (res, loc) -> (before loc, res, after loc)

-- * Indexed

updateAt :: (k -> a -> Maybe a) -> Int -> Map k a -> Map k a
updateAt f i m = case index i m of
        (v, loc) -> maybeAssign (f (key loc) v) loc

deleteAt :: Int -> Map k a -> Map k a
deleteAt i m = clear (snd (index i m))

-- * Min/Max

deleteMin :: Map k a -> Map k a
deleteMin m
  | null m = empty
  | otherwise = clear (snd (minLocation m))

deleteMax :: Map k a -> Map k a
deleteMax m
  | null m = empty
  | otherwise = clear (snd (maxLocation m))

deleteFindMin :: Map k a -> ((k, a), Map k a)
deleteFindMin m
  | null m = (error "Map.deleteFindMin: empty map", m)
  | otherwise = case minLocation m of
        (x, loc) -> ((key loc, x), clear loc)

deleteFindMax :: Map k a -> ((k, a), Map k a)
deleteFindMax m
  | null m = (error "Map.deleteFindMax: empty map", m)
  | otherwise = case maxLocation m of
        (x, loc) -> ((key loc, x), clear loc)

updateMin :: (a -> Maybe a) -> Map k a -> Map k a
updateMin f m
  | null m = m
  | otherwise = case minLocation m of
        (x, loc) -> maybeAssign (f x) loc

updateMax :: (a -> Maybe a) -> Map k a -> Map k a
updateMax f m
  | null m = m
  | otherwise = case maxLocation m of
        (x, loc) -> maybeAssign (f x) loc

updateMinWithKey :: (k -> a -> Maybe a) -> Map k a -> Map k a
updateMinWithKey f m
  | null m = m
  | otherwise = case minLocation m of
        (x, loc) -> maybeAssign (f (key loc) x) loc

updateMaxWithKey :: (k -> a -> Maybe a) -> Map k a -> Map k a
updateMaxWithKey f m
  | null m = m
  | otherwise = case maxLocation m of
        (x, loc) -> maybeAssign (f (key loc) x) loc

minView :: Map k a -> Maybe (a, Map k a)
minView m
  | null m = Nothing
  | otherwise = case minLocation m of
        (x, loc) -> Just (x, clear loc)

maxView :: Map k a -> Maybe (a, Map k a)
maxView m
  | null m = Nothing
  | otherwise = case maxLocation m of
        (x, loc) -> Just (x, clear loc)

minViewWithKey :: Map k a -> Maybe ((k, a), Map k a)
minViewWithKey m
  | null m = Nothing
  | otherwise = case minLocation m of
        (x, loc) -> Just ((key loc, x), clear loc)

maxViewWithKey :: Map k a -> Maybe ((k, a), Map k a)
maxViewWithKey m
  | null m = Nothing
  | otherwise = case maxLocation m of
        (x, loc) -> Just ((key loc, x), clear loc)

-- utility

maybeAssign :: Maybe a -> Location k a -> Map k a
maybeAssign = maybe clear assign