Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add merge #45

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion dependent-map.cabal
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
name: dependent-map
version: 0.4.0.0
version: 0.4.1.0
stability: provisional

cabal-version: >= 1.6
Expand Down Expand Up @@ -35,6 +35,7 @@ Library
ghc-options: -fwarn-unused-imports -fwarn-unused-binds
exposed-modules: Data.Dependent.Map,
Data.Dependent.Map.Lens,
Data.Dependent.Map.Merge,
Data.Dependent.Map.Internal
other-modules: Data.Dependent.Map.PtrEquality
build-depends: base >= 4.9 && < 5,
Expand Down
72 changes: 12 additions & 60 deletions src/Data/Dependent/Map.hs
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,9 @@ module Data.Dependent.Map
, intersection
, intersectionWithKey

-- ** General combining functions
-- | See "Data.Dependent.Map.Merge"

-- * Traversal
-- ** Map
, map
Expand All @@ -72,6 +75,7 @@ module Data.Dependent.Map
, traverseWithKey_
, forWithKey_
, traverseWithKey
, traverseMaybeWithKey
, forWithKey
, mapAccumLWithKey
, mapAccumRWithKey
Expand Down Expand Up @@ -659,7 +663,7 @@ difference t1 Tip = t1
difference t1 (Bin _ k2 _x2 l2 r2) = case split k2 t1 of
(l1, r1)
| size t1 == size l1l2 + size r1r2 -> t1
| otherwise -> merge l1l2 r1r2
| otherwise -> link2 l1l2 r1r2
where
!l1l2 = l1 `difference` l2
!r1r2 = r1 `difference` r2
Expand All @@ -675,7 +679,7 @@ differenceWithKey f (Bin _ k1 x1 l1 r1) t2 = case splitLookup k1 t2 of
(l2, mx2, r2) -> case mx2 of
Nothing -> combine k1 x1 l1l2 r1r2
Just x2 -> case f k1 x1 x2 of
Nothing -> merge l1l2 r1r2
Nothing -> link2 l1l2 r1r2
Just x1x2 -> combine k1 x1x2 l1l2 r1r2
where !l1l2 = differenceWithKey f l1 l2
!r1r2 = differenceWithKey f r1 r2
Expand All @@ -698,7 +702,7 @@ intersection t1@(Bin s1 k1 x1 l1 r1) t2 =
then if l1l2 `ptrEq` l1 && r1r2 `ptrEq` r1
then t1
else combine k1 x1 l1l2 r1r2
else merge l1l2 r1r2
else link2 l1l2 r1r2

-- | /O(m * log (n\/m + 1), m <= n/. Intersection with a combining function.
intersectionWithKey :: GCompare k => (forall v. k v -> f v -> g v -> h v) -> DMap k f -> DMap k g -> DMap k h
Expand All @@ -709,7 +713,7 @@ intersectionWithKey f (Bin s1 k1 x1 l1 r1) t2 =
!l1l2 = intersectionWithKey f l1 l2
!r1r2 = intersectionWithKey f r1 r2
in case found of
Nothing -> merge l1l2 r1r2
Nothing -> link2 l1l2 r1r2
Just x2 -> combine k1 (f k1 x1 x2) l1l2 r1r2

{--------------------------------------------------------------------
Expand Down Expand Up @@ -766,19 +770,6 @@ isProperSubmapOfBy f t1 t2
Filter and partition
--------------------------------------------------------------------}

-- | /O(n)/. Filter all keys\/values that satisfy the predicate.
filterWithKey :: GCompare k => (forall v. k v -> f v -> Bool) -> DMap k f -> DMap k f
filterWithKey p = go
where
go Tip = Tip
go t@(Bin _ kx x l r)
| p kx x = if l' `ptrEq` l && r' `ptrEq` r
then t
else combine kx x l' r'
| otherwise = merge l' r'
where !l' = go l
!r' = go r

-- | /O(n)/. Partition the map according to a predicate. The first
-- map contains all elements that satisfy the predicate, the second all
-- elements that fail the predicate. See also 'split'.
Expand All @@ -788,8 +779,8 @@ partitionWithKey p0 m0 = toPair (go p0 m0)
go :: GCompare k => (forall v. k v -> f v -> Bool) -> DMap k f -> (DMap k f :*: DMap k f)
go _ Tip = (Tip :*: Tip)
go p (Bin _ kx x l r)
| p kx x = (combine kx x l1 r1 :*: merge l2 r2)
| otherwise = (merge l1 r1 :*: combine kx x l2 r2)
| p kx x = (combine kx x l1 r1 :*: link2 l2 r2)
| otherwise = (link2 l1 r1 :*: combine kx x l2 r2)
where
(l1 :*: l2) = go p l
(r1 :*: r2) = go p r
Expand All @@ -798,15 +789,6 @@ partitionWithKey p0 m0 = toPair (go p0 m0)
mapMaybe :: GCompare k => (forall v. f v -> Maybe (g v)) -> DMap k f -> DMap k g
mapMaybe f = mapMaybeWithKey (const f)

-- | /O(n)/. Map keys\/values and collect the 'Just' results.
mapMaybeWithKey :: GCompare k => (forall v. k v -> f v -> Maybe (g v)) -> DMap k f -> DMap k g
mapMaybeWithKey f = go
where
go Tip = Tip
go (Bin _ kx x l r) = case f kx x of
Just y -> combine kx y (go l) (go r)
Nothing -> merge (go l) (go r)

-- | /O(n)/. Map keys\/values and separate the 'Left' and 'Right' results.
mapEitherWithKey :: GCompare k =>
(forall v. k v -> f v -> Either (g v) (h v)) -> DMap k f -> (DMap k g, DMap k h)
Expand All @@ -817,8 +799,8 @@ mapEitherWithKey f0 = toPair . go f0
-> DMap k f -> (DMap k g :*: DMap k h)
go _ Tip = (Tip :*: Tip)
go f (Bin _ kx x l r) = case f kx x of
Left y -> (combine kx y l1 r1 :*: merge l2 r2)
Right z -> (merge l1 r1 :*: combine kx z l2 r2)
Left y -> (combine kx y l1 r1 :*: link2 l2 r2)
Right z -> (link2 l1 r1 :*: combine kx z l2 r2)
where
(l1,l2) = mapEitherWithKey f l
(r1,r2) = mapEitherWithKey f r
Expand All @@ -840,13 +822,6 @@ map f = go
ffor :: DMap k f -> (forall v. f v -> g v) -> DMap k g
ffor m f = map f m

-- | /O(n)/. Map a function over all values in the map.
mapWithKey :: (forall v. k v -> f v -> g v) -> DMap k f -> DMap k g
mapWithKey f = go
where
go Tip = Tip
go (Bin sx kx x l r) = Bin sx kx (f kx x) (go l) (go r)

-- | /O(n)/.
-- @'fforWithKey' == 'flip' 'mapWithKey'@ except we cannot actually use
-- 'flip' because of the lack of impredicative types.
Expand All @@ -870,17 +845,6 @@ traverseWithKey_ f = go
forWithKey_ :: Applicative t => DMap k f -> (forall v. k v -> f v -> t ()) -> t ()
forWithKey_ m f = traverseWithKey_ f m

-- | /O(n)/.
-- @'traverseWithKey' f m == 'fromList' <$> 'traverse' (\(k, v) -> (,) k <$> f k v) ('toList' m)@
-- That is, behaves exactly like a regular 'traverse' except that the traversing
-- function also has access to the key associated with a value.
traverseWithKey :: Applicative t => (forall v. k v -> f v -> t (g v)) -> DMap k f -> t (DMap k g)
traverseWithKey f = go
where
go Tip = pure Tip
go (Bin 1 k v _ _) = (\v' -> Bin 1 k v' Tip Tip) <$> f k v
go (Bin s k v l r) = flip (Bin s k) <$> go l <*> f k v <*> go r

-- | /O(n)/.
-- @'forWithKey' == 'flip' 'traverseWithKey'@ except we cannot actually use
-- 'flip' because of the lack of impredicative types.
Expand Down Expand Up @@ -1112,18 +1076,6 @@ split k = toPair . go
GEQ -> (l :*: r)
{-# INLINABLE split #-}

-- | /O(log n)/. The expression (@'splitLookup' k map@) splits a map just
-- like 'split' but also returns @'lookup' k map@.
splitLookup :: forall k f v. GCompare k => k v -> DMap k f -> (DMap k f, Maybe (f v), DMap k f)
splitLookup k = toTriple . go
where
go :: DMap k f -> Triple' (DMap k f) (Maybe (f v)) (DMap k f)
go Tip = Triple' Tip Nothing Tip
go (Bin _ kx x l r) = case gcompare k kx of
GLT -> let !(Triple' lt z gt) = go l in Triple' lt z (combine kx x gt r)
GGT -> let !(Triple' lt z gt) = go r in Triple' (combine kx x l lt) z gt
GEQ -> Triple' l (Just x) r

-- | /O(log n)/. The expression (@'splitMember' k map@) splits a map just
-- like 'split' but also returns @'member' k map@.
splitMember :: forall k f v. GCompare k => k v -> DMap k f -> (DMap k f, Bool, DMap k f)
Expand Down
106 changes: 95 additions & 11 deletions src/Data/Dependent/Map/Internal.hs
Original file line number Diff line number Diff line change
@@ -1,17 +1,21 @@
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE Safe #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeOperators #-}
module Data.Dependent.Map.Internal where

import Control.Applicative (liftA3)
import Data.Dependent.Sum (DSum((:=>)))
import Data.GADT.Compare (GCompare, GOrdering(..), gcompare)
import Data.Some (Some, mkSome, withSome)
import Data.Typeable (Typeable)

import Data.Dependent.Map.PtrEquality (ptrEq)

-- |Dependent maps: 'k' is a GADT-like thing with a facility for
-- rediscovering its type parameter, elements of which function as identifiers
-- tagged with the type of the thing they identify. Real GADTs are one
Expand All @@ -34,6 +38,87 @@ data DMap k f where
-> DMap k f
deriving Typeable

-- | /O(n)/. Filter all keys\/values that satisfy the predicate.
filterWithKey :: GCompare k => (forall v. k v -> f v -> Bool) -> DMap k f -> DMap k f
filterWithKey p = go
where
go Tip = Tip
go t@(Bin _ kx x l r)
| p kx x = if l' `ptrEq` l && r' `ptrEq` r
then t
else combine kx x l' r'
| otherwise = link2 l' r'
where !l' = go l
!r' = go r

-- | /O(n)/. Filter keys and values using an 'Applicative'
-- predicate.
filterWithKeyA :: (GCompare k, Applicative t) => (forall v. k v -> f v -> t Bool) -> DMap k f -> t (DMap k f)
filterWithKeyA _ Tip = pure Tip
filterWithKeyA p t@(Bin _ kx x l r) =
liftA3 combine' (p kx x) (filterWithKeyA p l) (filterWithKeyA p r)
where
combine' True pl pr
| pl `ptrEq` l && pr `ptrEq` r = t
| otherwise = combine kx x pl pr
combine' False pl pr = link2 pl pr

-- | /O(n)/. Map keys\/values and collect the 'Just' results.
mapMaybeWithKey :: GCompare k => (forall v. k v -> f v -> Maybe (g v)) -> DMap k f -> DMap k g
mapMaybeWithKey f = go
where
go Tip = Tip
go (Bin _ kx x l r) = case f kx x of
Just y -> combine kx y (go l) (go r)
Nothing -> link2 (go l) (go r)

-- | /O(n)/. Map a function over all values in the map.
mapWithKey :: (forall v. k v -> f v -> g v) -> DMap k f -> DMap k g
mapWithKey f = go
where
go Tip = Tip
go (Bin sx kx x l r) = Bin sx kx (f kx x) (go l) (go r)

-- | /O(n)/.
-- @'traverseWithKey' f m == 'fromList' <$> 'traverse' (\(k, v) -> (,) k <$> f k v) ('toList' m)@
-- That is, behaves exactly like a regular 'traverse' except that the traversing
-- function also has access to the key associated with a value.
traverseWithKey :: Applicative t => (forall v. k v -> f v -> t (g v)) -> DMap k f -> t (DMap k g)
traverseWithKey f = go
where
go Tip = pure Tip
go (Bin 1 k v _ _) = (\v' -> Bin 1 k v' Tip Tip) <$> f k v
go (Bin s k v l r) = flip (Bin s k) <$> go l <*> f k v <*> go r

-- | /O(n)/. Traverse keys\/values and collect the 'Just' results.
--
-- @since UNRELEASED
traverseMaybeWithKey
:: (GCompare k, Applicative t)
=> (forall v. k v -> f v -> t (Maybe (g v))) -> DMap k f -> t (DMap k g)
traverseMaybeWithKey f = go
where
go Tip = pure Tip
go (Bin _ kx x Tip Tip) = maybe Tip (\x' -> Bin 1 kx x' Tip Tip) <$> f kx x
go (Bin _ kx x l r) = liftA3 combine' (go l) (f kx x) (go r)
where
combine' !l' mx !r' = case mx of
Nothing -> link2 l' r'
Just x' -> combine kx x' l' r'


-- | /O(log n)/. The expression (@'splitLookup' k map@) splits a map just
-- like 'split' but also returns @'lookup' k map@.
splitLookup :: forall k f v. GCompare k => k v -> DMap k f -> (DMap k f, Maybe (f v), DMap k f)
splitLookup k = toTriple . go
where
go :: DMap k f -> Triple' (DMap k f) (Maybe (f v)) (DMap k f)
go Tip = Triple' Tip Nothing Tip
go (Bin _ kx x l r) = case gcompare k kx of
GLT -> let !(Triple' lt z gt) = go l in Triple' lt z (combine kx x gt r)
GGT -> let !(Triple' lt z gt) = go r in Triple' (combine kx x l lt) z gt
GEQ -> Triple' l (Just x) r

{--------------------------------------------------------------------
Construction
--------------------------------------------------------------------}
Expand Down Expand Up @@ -112,10 +197,10 @@ lookupAssoc sk = withSome sk $ \k ->
are valid:
[glue l r] Glues [l] and [r] together. Assumes that [l] and
[r] are already balanced with respect to each other.
[merge l r] Merges two trees and restores balance.
[link2 l r] Merges two trees and restores balance.

Note: in contrast to Adam's paper, we use (<=) comparisons instead
of (<) comparisons in [combine], [merge] and [balance].
of (<) comparisons in [combine], [link2] and [balance].
Quickcheck (on [difference]) showed that this was necessary in order
to maintain the invariants. It is quite unsatisfactory that I haven't
been able to find out why this is actually the case! Fortunately, it
Expand All @@ -133,7 +218,6 @@ combine kx x l@(Bin sizeL ky y ly ry) r@(Bin sizeR kz z lz rz)
| delta*sizeR <= sizeL = balance ky y ly (combine kx x ry r)
| otherwise = bin kx x l r


-- insertMin and insertMax don't perform potentially expensive comparisons.
insertMax,insertMin :: k v -> f v -> DMap k f -> DMap k f
insertMax kx x t
Expand All @@ -149,14 +233,14 @@ insertMin kx x t
-> balance ky y (insertMin kx x l) r

{--------------------------------------------------------------------
[merge l r]: merges two trees.
[link2 l r]: merges two trees.
--------------------------------------------------------------------}
merge :: DMap k f -> DMap k f -> DMap k f
merge Tip r = r
merge l Tip = l
merge l@(Bin sizeL kx x lx rx) r@(Bin sizeR ky y ly ry)
| delta*sizeL <= sizeR = balance ky y (merge l ly) ry
| delta*sizeR <= sizeL = balance kx x lx (merge rx r)
link2 :: DMap k f -> DMap k f -> DMap k f
link2 Tip r = r
link2 l Tip = l
link2 l@(Bin sizeL kx x lx rx) r@(Bin sizeR ky y ly ry)
| delta*sizeL <= sizeR = balance ky y (link2 l ly) ry
| delta*sizeR <= sizeL = balance kx x lx (link2 rx r)
| otherwise = glue l r

{--------------------------------------------------------------------
Expand Down
Loading