Library proposal: add a Location interface for element-wise operations on Data.Map (#4887)

[Bas van Dijk]

On Sat, Jan 8, 2011 at 1:49 AM, Ross Paterson <ross at soi.city.ac.uk> wrote:
> On Sat, Jan 08, 2011 at 01:07:48AM +0100, Bas van Dijk wrote:
>> (The following is just some brainstorming)
>>
>> Why not store the value in the Location as in:
>>
>> data Location k a
>>     = Empty !k !(Path k a)
>>     | Full {-# UNPACK #-} !Size !k a !(Path k a) !(Map k a) !(Map k a)
>>
>> [...]
>> We do need a function to retrieve the value:
>>
>> value :: Location k a -> Maybe a
>> value (Empty _ _) = Nothing
>> value (Full _ _ v _ _ _) = Just v
>
> That would work well for search, but then index, minLocation and maxLocation
> would return Locations that value was always mapped to Just something.
> Extra invariants like that feel wrong to me.
>
> _______________________________________________
> Libraries mailing list
> Libraries at haskell.org
> http://www.haskell.org/mailman/listinfo/libraries
>

You could go with something like the following. However I don't think
it's worth the trouble:

data Location k a = E !(Empty k a)
                  | F !(Full  k a)

location :: (Empty k a -> b) -> (Full k a -> b) -> Location k a -> b
location f _ (E empty) = f empty
location _ g (F full)  = g full

data Empty k a = Empty !k !(Path k a)
data Full k a = Full {-# UNPACK #-} !Size !k a !(Path k a) !(Map k a) !(Map k a)

data Path k a
    = Root
    | LeftBin  {-# UNPACK #-} !Size !k a !(Path k a) !(Map k a)
    | RightBin {-# UNPACK #-} !Size !k a !(Map k a)  !(Path k a)

search :: Ord k => k -> Map k a -> Location k a
search k = k `seq` go Root
  where
    go path Tip = E $ Empty k path
    go path (Bin sx kx x l r) = case compare k kx of
       LT -> go (LeftBin sx kx x path r) l
       GT -> go (RightBin sx kx x l path) r
       EQ -> F $ Full sx kx x path l r

index :: Int -> Map k a -> Full k a
index = go Root
  where
    STRICT_2_OF_3(go)
    go _path _i Tip = error "Map.index: out of range"
    go path i (Bin sx kx x l r) = case compare i size_l of
        LT -> go (LeftBin sx kx x path r) i l
        GT -> go (RightBin sx kx x l path) (i-size_l-1) r
        EQ -> Full sx kx x path l r
      where size_l = size l

minLocation :: Map k a -> Full k a
minLocation = go Root
  where
    go _path Tip = error "Map.least: empty map"
    go path (Bin sx kx x Tip r) = Full sx kx x path Tip r
    go path (Bin sx kx x l r) = go (LeftBin sx kx x path r) l

maxLocation :: Map k a -> Full k a
maxLocation = go Root
  where
    go _path Tip = error "Map.greatest: empty map"
    go path (Bin sx kx x l Tip) = Full sx kx x path l Tip
    go path (Bin sx kx x l r) = go (RightBin sx kx x l path) r

class Key m where
    key :: m k a -> k

instance Key Empty where
    key (Empty kx _path) = kx

instance Key Full where
    key (Full _sx kx _x _path _l _r) = kx

instance Key Location where
    key = location key key

value :: Full k a -> a
value (Full _sx _kx x _path _l _r) = x

class Before m where
    before :: Ord k => m k a -> Map k a

instance Before Empty where
    before (Empty _k path) = buildBefore Tip path

instance Before Full where
    before (Full _sx _kx _x path l _r) = buildBefore l path

instance Before Location where
    before = location before before

buildBefore :: Ord k => Map k a -> Path k a -> Map k a
buildBefore t Root = t
buildBefore t (LeftBin _sx _kx _x path _r) = buildBefore t path
buildBefore t (RightBin _sx kx x l path) = buildBefore (join kx x l t) path

class After m where
    after :: Ord k => m k a -> Map k a

instance After Empty where
    after (Empty _k path) = buildAfter Tip path

instance After Full where
    after (Full _sx _kx _x path l _r) = buildAfter l path

instance After Location where
    after = location after after

buildAfter :: Ord k => Map k a -> Path k a -> Map k a
buildAfter t Root = t
buildAfter t (LeftBin _sx kx x path r) = buildAfter (join kx x t r) path
buildAfter t (RightBin _sx _kx _x _l path) = buildAfter t path

class Assign m where
    assign :: a -> m k a -> Map k a

instance Assign Empty where
    assign x (Empty k path) = rebuildGT (singleton k x) path

instance Assign Full where
    assign x (Full sx kx _x path l r) = rebuildEQ (Bin sx kx x l r) path

instance Assign Location where
    assign x = location (assign x) (assign x)

class Clear m where clear :: m k a -> Map k a

instance Clear Empty
    where clear (Empty _k path) = rebuildEQ Tip path

instance Clear Full
    where clear (Full _sx _kx _x path l r) = rebuildLT (glue l r) path

instance Clear Location
    where clear = location clear clear

-- Rebuild the tree the same size as it was, so no rebalancing is needed.
rebuildEQ :: Map k a -> Path k a -> Map k a
rebuildEQ t Root = t
rebuildEQ l (LeftBin sx kx x path r) = rebuildEQ (Bin sx kx x l r) path
rebuildEQ r (RightBin sx kx x l path) = rebuildEQ (Bin sx kx x l r) path

-- Rebuild the tree one entry smaller than it was, rebalancing as we go.
rebuildLT :: Map k a -> Path k a -> Map k a
rebuildLT t Root = t
rebuildLT l (LeftBin _sx kx x path r) = rebuildLT (balanceR kx x l r) path
rebuildLT r (RightBin _sx kx x l path) = rebuildLT (balanceL kx x l r) path

-- Rebuild the tree one entry larger than it was, rebalancing as we go.
rebuildGT :: Map k a -> Path k a -> Map k a
rebuildGT t Root = t
rebuildGT l (LeftBin _sx kx x path r) = rebuildGT (balanceL kx x l r) path
rebuildGT r (RightBin _sx kx x l path) = rebuildGT (balanceR kx x l r) path

Regards,

Bas