{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE Safe #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeOperators #-}
module Data.Dependent.Map.Internal where

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

-- |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
-- useful instantiation of @k@, as are 'Tag's from "Data.Unique.Tag" in the
-- 'prim-uniq' package.
--
-- Semantically, @'DMap' k f@ is equivalent to a set of @'DSum' k f@ where no two
-- elements have the same tag.
--
-- More informally, 'DMap' is to dependent products as 'M.Map' is to @(->)@.
-- Thus it could also be thought of as a partial (in the sense of \"partial
-- function\") dependent product.
data DMap k f where
    Tip :: DMap k f
    Bin :: {- sz    -} !Int
        -> {- key   -} !(k v)
        -> {- value -} f v
        -> {- left  -} !(DMap k f)
        -> {- right -} !(DMap k f)
        -> DMap k f
    deriving Typeable

{--------------------------------------------------------------------
  Construction
--------------------------------------------------------------------}

-- | /O(1)/. The empty map.
--
-- > empty      == fromList []
-- > size empty == 0
empty :: DMap k f
empty :: DMap k f
empty = DMap k f
forall k (k :: k -> *) (f :: k -> *). DMap k f
Tip

-- | /O(1)/. A map with a single element.
--
-- > singleton 1 'a'        == fromList [(1, 'a')]
-- > size (singleton 1 'a') == 1
singleton :: k v -> f v -> DMap k f
singleton :: k v -> f v -> DMap k f
singleton k v
k f v
x = Int -> k v -> f v -> DMap k f -> DMap k f -> DMap k f
forall k (k :: k -> *) (v :: k) (f :: k -> *).
Int -> k v -> f v -> DMap k f -> DMap k f -> DMap k f
Bin Int
1 k v
k f v
x DMap k f
forall k (k :: k -> *) (f :: k -> *). DMap k f
Tip DMap k f
forall k (k :: k -> *) (f :: k -> *). DMap k f
Tip

{--------------------------------------------------------------------
  Query
--------------------------------------------------------------------}

-- | /O(1)/. Is the map empty?
null :: DMap k f -> Bool
null :: DMap k f -> Bool
null DMap k f
Tip    = Bool
True
null Bin{}  = Bool
False

-- | /O(1)/. The number of elements in the map.
size :: DMap k f -> Int
size :: DMap k f -> Int
size DMap k f
Tip                = Int
0
size (Bin Int
n k v
_ f v
_ DMap k f
_ DMap k f
_)    = Int
n

-- | /O(log n)/. Lookup the value at a key in the map.
--
-- The function will return the corresponding value as @('Just' value)@,
-- or 'Nothing' if the key isn't in the map.
lookup :: forall k f v. GCompare k => k v -> DMap k f -> Maybe (f v)
lookup :: k v -> DMap k f -> Maybe (f v)
lookup k v
k = k v
k k v -> (DMap k f -> Maybe (f v)) -> DMap k f -> Maybe (f v)
`seq` DMap k f -> Maybe (f v)
go
    where
        go :: DMap k f -> Maybe (f v)
        go :: DMap k f -> Maybe (f v)
go DMap k f
Tip = Maybe (f v)
forall a. Maybe a
Nothing
        go (Bin Int
_ k v
kx f v
x DMap k f
l DMap k f
r) =
            case k v -> k v -> GOrdering v v
forall k (f :: k -> *) (a :: k) (b :: k).
GCompare f =>
f a -> f b -> GOrdering a b
gcompare k v
k k v
kx of
                GOrdering v v
GLT -> DMap k f -> Maybe (f v)
go DMap k f
l
                GOrdering v v
GGT -> DMap k f -> Maybe (f v)
go DMap k f
r
                GOrdering v v
GEQ -> f v -> Maybe (f v)
forall a. a -> Maybe a
Just f v
x

lookupAssoc :: forall k f v. GCompare k => Some k -> DMap k f -> Maybe (DSum k f)
lookupAssoc :: Some k -> DMap k f -> Maybe (DSum k f)
lookupAssoc Some k
sk = Some k
-> (forall (a :: k). k a -> DMap k f -> Maybe (DSum k f))
-> DMap k f
-> Maybe (DSum k f)
forall k (tag :: k -> *) b.
Some tag -> (forall (a :: k). tag a -> b) -> b
withSome Some k
sk ((forall (a :: k). k a -> DMap k f -> Maybe (DSum k f))
 -> DMap k f -> Maybe (DSum k f))
-> (forall (a :: k). k a -> DMap k f -> Maybe (DSum k f))
-> DMap k f
-> Maybe (DSum k f)
forall a b. (a -> b) -> a -> b
$ \k a
k ->
  let
    go :: DMap k f -> Maybe (DSum k f)
    go :: DMap k f -> Maybe (DSum k f)
go DMap k f
Tip = Maybe (DSum k f)
forall a. Maybe a
Nothing
    go (Bin Int
_ k v
kx f v
x DMap k f
l DMap k f
r) =
        case k a -> k v -> GOrdering a v
forall k (f :: k -> *) (a :: k) (b :: k).
GCompare f =>
f a -> f b -> GOrdering a b
gcompare k a
k k v
kx of
            GOrdering a v
GLT -> DMap k f -> Maybe (DSum k f)
go DMap k f
l
            GOrdering a v
GGT -> DMap k f -> Maybe (DSum k f)
go DMap k f
r
            GOrdering a v
GEQ -> DSum k f -> Maybe (DSum k f)
forall a. a -> Maybe a
Just (k v
kx k v -> f v -> DSum k f
forall k (tag :: k -> *) (f :: k -> *) (a :: k).
tag a -> f a -> DSum tag f
:=> f v
x)
  in k a
k k a
-> (DMap k f -> Maybe (DSum k f)) -> DMap k f -> Maybe (DSum k f)
`seq` DMap k f -> Maybe (DSum k f)
go

{--------------------------------------------------------------------
  Utility functions that maintain the balance properties of the tree.
  All constructors assume that all values in [l] < [k] and all values
  in [r] > [k], and that [l] and [r] are valid trees.

  In order of sophistication:
    [Bin sz k x l r]  The type constructor.
    [bin k x l r]     Maintains the correct size, assumes that both [l]
                      and [r] are balanced with respect to each other.
    [balance k x l r] Restores the balance and size.
                      Assumes that the original tree was balanced and
                      that [l] or [r] has changed by at most one element.
    [combine k x l r] Restores balance and size.

  Furthermore, we can construct a new tree from two trees. Both operations
  assume that all values in [l] < all values in [r] and that [l] and [r]
  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.

  Note: in contrast to Adam's paper, we use (<=) comparisons instead
  of (<) comparisons in [combine], [merge] 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
  doesn't hurt to be a bit more conservative.
--------------------------------------------------------------------}

{--------------------------------------------------------------------
  Combine
--------------------------------------------------------------------}
combine :: GCompare k => k v -> f v -> DMap k f -> DMap k f -> DMap k f
combine :: k v -> f v -> DMap k f -> DMap k f -> DMap k f
combine k v
kx f v
x DMap k f
Tip DMap k f
r  = k v -> f v -> DMap k f -> DMap k f
forall k (k :: k -> *) (v :: k) (f :: k -> *).
k v -> f v -> DMap k f -> DMap k f
insertMin k v
kx f v
x DMap k f
r
combine k v
kx f v
x DMap k f
l DMap k f
Tip  = k v -> f v -> DMap k f -> DMap k f
forall k (k :: k -> *) (v :: k) (f :: k -> *).
k v -> f v -> DMap k f -> DMap k f
insertMax k v
kx f v
x DMap k f
l
combine k v
kx f v
x l :: DMap k f
l@(Bin Int
sizeL k v
ky f v
y DMap k f
ly DMap k f
ry) r :: DMap k f
r@(Bin Int
sizeR k v
kz f v
z DMap k f
lz DMap k f
rz)
  | Int
deltaInt -> Int -> Int
forall a. Num a => a -> a -> a
*Int
sizeL Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
sizeR  = k v -> f v -> DMap k f -> DMap k f -> DMap k f
forall k (k :: k -> *) (v :: k) (f :: k -> *).
k v -> f v -> DMap k f -> DMap k f -> DMap k f
balance k v
kz f v
z (k v -> f v -> DMap k f -> DMap k f -> DMap k f
forall k (k :: k -> *) (v :: k) (f :: k -> *).
GCompare k =>
k v -> f v -> DMap k f -> DMap k f -> DMap k f
combine k v
kx f v
x DMap k f
l DMap k f
lz) DMap k f
rz
  | Int
deltaInt -> Int -> Int
forall a. Num a => a -> a -> a
*Int
sizeR Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
sizeL  = k v -> f v -> DMap k f -> DMap k f -> DMap k f
forall k (k :: k -> *) (v :: k) (f :: k -> *).
k v -> f v -> DMap k f -> DMap k f -> DMap k f
balance k v
ky f v
y DMap k f
ly (k v -> f v -> DMap k f -> DMap k f -> DMap k f
forall k (k :: k -> *) (v :: k) (f :: k -> *).
GCompare k =>
k v -> f v -> DMap k f -> DMap k f -> DMap k f
combine k v
kx f v
x DMap k f
ry DMap k f
r)
  | Bool
otherwise             = k v -> f v -> DMap k f -> DMap k f -> DMap k f
forall k (k :: k -> *) (v :: k) (f :: k -> *).
k v -> f v -> DMap k f -> DMap k f -> DMap k f
bin k v
kx f v
x DMap k f
l DMap k f
r


-- insertMin and insertMax don't perform potentially expensive comparisons.
insertMax,insertMin :: k v -> f v -> DMap k f -> DMap k f
insertMax :: k v -> f v -> DMap k f -> DMap k f
insertMax k v
kx f v
x DMap k f
t
  = case DMap k f
t of
      DMap k f
Tip -> k v -> f v -> DMap k f
forall k (k :: k -> *) (v :: k) (f :: k -> *).
k v -> f v -> DMap k f
singleton k v
kx f v
x
      Bin Int
_ k v
ky f v
y DMap k f
l DMap k f
r
          -> k v -> f v -> DMap k f -> DMap k f -> DMap k f
forall k (k :: k -> *) (v :: k) (f :: k -> *).
k v -> f v -> DMap k f -> DMap k f -> DMap k f
balance k v
ky f v
y DMap k f
l (k v -> f v -> DMap k f -> DMap k f
forall k (k :: k -> *) (v :: k) (f :: k -> *).
k v -> f v -> DMap k f -> DMap k f
insertMax k v
kx f v
x DMap k f
r)

insertMin :: k v -> f v -> DMap k f -> DMap k f
insertMin k v
kx f v
x DMap k f
t
  = case DMap k f
t of
      DMap k f
Tip -> k v -> f v -> DMap k f
forall k (k :: k -> *) (v :: k) (f :: k -> *).
k v -> f v -> DMap k f
singleton k v
kx f v
x
      Bin Int
_ k v
ky f v
y DMap k f
l DMap k f
r
          -> k v -> f v -> DMap k f -> DMap k f -> DMap k f
forall k (k :: k -> *) (v :: k) (f :: k -> *).
k v -> f v -> DMap k f -> DMap k f -> DMap k f
balance k v
ky f v
y (k v -> f v -> DMap k f -> DMap k f
forall k (k :: k -> *) (v :: k) (f :: k -> *).
k v -> f v -> DMap k f -> DMap k f
insertMin k v
kx f v
x DMap k f
l) DMap k f
r

{--------------------------------------------------------------------
  [merge l r]: merges two trees.
--------------------------------------------------------------------}
merge :: DMap k f -> DMap k f -> DMap k f
merge :: DMap k f -> DMap k f -> DMap k f
merge DMap k f
Tip DMap k f
r   = DMap k f
r
merge DMap k f
l DMap k f
Tip   = DMap k f
l
merge l :: DMap k f
l@(Bin Int
sizeL k v
kx f v
x DMap k f
lx DMap k f
rx) r :: DMap k f
r@(Bin Int
sizeR k v
ky f v
y DMap k f
ly DMap k f
ry)
  | Int
deltaInt -> Int -> Int
forall a. Num a => a -> a -> a
*Int
sizeL Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
sizeR = k v -> f v -> DMap k f -> DMap k f -> DMap k f
forall k (k :: k -> *) (v :: k) (f :: k -> *).
k v -> f v -> DMap k f -> DMap k f -> DMap k f
balance k v
ky f v
y (DMap k f -> DMap k f -> DMap k f
forall k (k :: k -> *) (f :: k -> *).
DMap k f -> DMap k f -> DMap k f
merge DMap k f
l DMap k f
ly) DMap k f
ry
  | Int
deltaInt -> Int -> Int
forall a. Num a => a -> a -> a
*Int
sizeR Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
sizeL = k v -> f v -> DMap k f -> DMap k f -> DMap k f
forall k (k :: k -> *) (v :: k) (f :: k -> *).
k v -> f v -> DMap k f -> DMap k f -> DMap k f
balance k v
kx f v
x DMap k f
lx (DMap k f -> DMap k f -> DMap k f
forall k (k :: k -> *) (f :: k -> *).
DMap k f -> DMap k f -> DMap k f
merge DMap k f
rx DMap k f
r)
  | Bool
otherwise            = DMap k f -> DMap k f -> DMap k f
forall k (k :: k -> *) (f :: k -> *).
DMap k f -> DMap k f -> DMap k f
glue DMap k f
l DMap k f
r

{--------------------------------------------------------------------
  [glue l r]: glues two trees together.
  Assumes that [l] and [r] are already balanced with respect to each other.
--------------------------------------------------------------------}
glue :: DMap k f -> DMap k f -> DMap k f
glue :: DMap k f -> DMap k f -> DMap k f
glue DMap k f
Tip DMap k f
r = DMap k f
r
glue DMap k f
l DMap k f
Tip = DMap k f
l
glue DMap k f
l DMap k f
r
  | DMap k f -> Int
forall k (k :: k -> *) (f :: k -> *). DMap k f -> Int
size DMap k f
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> DMap k f -> Int
forall k (k :: k -> *) (f :: k -> *). DMap k f -> Int
size DMap k f
r = case DMap k f -> (DSum k f, DMap k f)
forall k (k :: k -> *) (f :: k -> *).
DMap k f -> (DSum k f, DMap k f)
deleteFindMax DMap k f
l of (k a
km :=> f a
m,DMap k f
l') -> k a -> f a -> DMap k f -> DMap k f -> DMap k f
forall k (k :: k -> *) (v :: k) (f :: k -> *).
k v -> f v -> DMap k f -> DMap k f -> DMap k f
balance k a
km f a
m DMap k f
l' DMap k f
r
  | Bool
otherwise       = case DMap k f -> (DSum k f, DMap k f)
forall k (k :: k -> *) (f :: k -> *).
DMap k f -> (DSum k f, DMap k f)
deleteFindMin DMap k f
r of (k a
km :=> f a
m,DMap k f
r') -> k a -> f a -> DMap k f -> DMap k f -> DMap k f
forall k (k :: k -> *) (v :: k) (f :: k -> *).
k v -> f v -> DMap k f -> DMap k f -> DMap k f
balance k a
km f a
m DMap k f
l DMap k f
r'

-- | /O(log n)/. Delete and find the minimal element.
--
-- > deleteFindMin (fromList [(5,"a"), (3,"b"), (10,"c")]) == ((3,"b"), fromList[(5,"a"), (10,"c")])
-- > deleteFindMin                                            Error: can not return the minimal element of an empty map

deleteFindMin :: DMap k f -> (DSum k f, DMap k f)
deleteFindMin :: DMap k f -> (DSum k f, DMap k f)
deleteFindMin DMap k f
t = case DMap k f -> Maybe (DSum k f, DMap k f)
forall k (k :: k -> *) (f :: k -> *).
DMap k f -> Maybe (DSum k f, DMap k f)
minViewWithKey DMap k f
t of
      Maybe (DSum k f, DMap k f)
Nothing -> ([Char] -> DSum k f
forall a. HasCallStack => [Char] -> a
error [Char]
"Map.deleteFindMin: can not return the minimal element of an empty map", DMap k f
forall k (k :: k -> *) (f :: k -> *). DMap k f
Tip)
      Just (DSum k f, DMap k f)
p -> (DSum k f, DMap k f)
p

-- | A strict pair.
data (:*:) a b = !a :*: !b
infixr 1 :*:

-- | Convert a strict pair to a pair.
toPair :: a :*: b -> (a, b)
toPair :: (a :*: b) -> (a, b)
toPair (a
a :*: b
b) = (a
a, b
b)
{-# INLINE toPair #-}

data Triple' a b c = Triple' !a !b !c

-- | Convert a strict triple to a triple.
toTriple :: Triple' a b c -> (a, b, c)
toTriple :: Triple' a b c -> (a, b, c)
toTriple (Triple' a
a b
b c
c) = (a
a, b
b, c
c)
{-# INLINE toTriple #-}

-- | /O(log n)/. Retrieves the minimal (key :=> value) entry of the map, and
-- the map stripped of that element, or 'Nothing' if passed an empty map.
minViewWithKey :: forall k f . DMap k f -> Maybe (DSum k f, DMap k f)
minViewWithKey :: DMap k f -> Maybe (DSum k f, DMap k f)
minViewWithKey DMap k f
Tip = Maybe (DSum k f, DMap k f)
forall a. Maybe a
Nothing
minViewWithKey (Bin Int
_ k v
k0 f v
x0 DMap k f
l0 DMap k f
r0) = (DSum k f, DMap k f) -> Maybe (DSum k f, DMap k f)
forall a. a -> Maybe a
Just ((DSum k f, DMap k f) -> Maybe (DSum k f, DMap k f))
-> (DSum k f, DMap k f) -> Maybe (DSum k f, DMap k f)
forall a b. (a -> b) -> a -> b
$! (DSum k f :*: DMap k f) -> (DSum k f, DMap k f)
forall a b. (a :*: b) -> (a, b)
toPair ((DSum k f :*: DMap k f) -> (DSum k f, DMap k f))
-> (DSum k f :*: DMap k f) -> (DSum k f, DMap k f)
forall a b. (a -> b) -> a -> b
$ k v -> f v -> DMap k f -> DMap k f -> DSum k f :*: DMap k f
forall (v :: k).
k v -> f v -> DMap k f -> DMap k f -> DSum k f :*: DMap k f
go k v
k0 f v
x0 DMap k f
l0 DMap k f
r0
  where
    go :: k v -> f v -> DMap k f -> DMap k f -> DSum k f :*: DMap k f
    go :: k v -> f v -> DMap k f -> DMap k f -> DSum k f :*: DMap k f
go k v
k f v
x DMap k f
Tip DMap k f
r = (k v
k k v -> f v -> DSum k f
forall k (tag :: k -> *) (f :: k -> *) (a :: k).
tag a -> f a -> DSum tag f
:=> f v
x) DSum k f -> DMap k f -> DSum k f :*: DMap k f
forall a b. a -> b -> a :*: b
:*: DMap k f
r
    go k v
k f v
x (Bin Int
_ k v
kl f v
xl DMap k f
ll DMap k f
lr) DMap k f
r =
      let !(DSum k f
km :*: DMap k f
l') = k v -> f v -> DMap k f -> DMap k f -> DSum k f :*: DMap k f
forall (v :: k).
k v -> f v -> DMap k f -> DMap k f -> DSum k f :*: DMap k f
go k v
kl f v
xl DMap k f
ll DMap k f
lr
      in (DSum k f
km DSum k f -> DMap k f -> DSum k f :*: DMap k f
forall a b. a -> b -> a :*: b
:*: k v -> f v -> DMap k f -> DMap k f -> DMap k f
forall k (k :: k -> *) (v :: k) (f :: k -> *).
k v -> f v -> DMap k f -> DMap k f -> DMap k f
balance k v
k f v
x DMap k f
l' DMap k f
r)

-- | /O(log n)/. Retrieves the maximal (key :=> value) entry of the map, and
-- the map stripped of that element, or 'Nothing' if passed an empty map.
maxViewWithKey :: forall k f . DMap k f -> Maybe (DSum k f, DMap k f)
maxViewWithKey :: DMap k f -> Maybe (DSum k f, DMap k f)
maxViewWithKey DMap k f
Tip = Maybe (DSum k f, DMap k f)
forall a. Maybe a
Nothing
maxViewWithKey (Bin Int
_ k v
k0 f v
x0 DMap k f
l0 DMap k f
r0) = (DSum k f, DMap k f) -> Maybe (DSum k f, DMap k f)
forall a. a -> Maybe a
Just ((DSum k f, DMap k f) -> Maybe (DSum k f, DMap k f))
-> (DSum k f, DMap k f) -> Maybe (DSum k f, DMap k f)
forall a b. (a -> b) -> a -> b
$! (DSum k f :*: DMap k f) -> (DSum k f, DMap k f)
forall a b. (a :*: b) -> (a, b)
toPair ((DSum k f :*: DMap k f) -> (DSum k f, DMap k f))
-> (DSum k f :*: DMap k f) -> (DSum k f, DMap k f)
forall a b. (a -> b) -> a -> b
$ k v -> f v -> DMap k f -> DMap k f -> DSum k f :*: DMap k f
forall (v :: k).
k v -> f v -> DMap k f -> DMap k f -> DSum k f :*: DMap k f
go k v
k0 f v
x0 DMap k f
l0 DMap k f
r0
  where
    go :: k v -> f v -> DMap k f -> DMap k f -> DSum k f :*: DMap k f
    go :: k v -> f v -> DMap k f -> DMap k f -> DSum k f :*: DMap k f
go k v
k f v
x DMap k f
l DMap k f
Tip = (k v
k k v -> f v -> DSum k f
forall k (tag :: k -> *) (f :: k -> *) (a :: k).
tag a -> f a -> DSum tag f
:=> f v
x) DSum k f -> DMap k f -> DSum k f :*: DMap k f
forall a b. a -> b -> a :*: b
:*: DMap k f
l
    go k v
k f v
x DMap k f
l (Bin Int
_ k v
kr f v
xr DMap k f
rl DMap k f
rr) =
      let !(DSum k f
km :*: DMap k f
r') = k v -> f v -> DMap k f -> DMap k f -> DSum k f :*: DMap k f
forall (v :: k).
k v -> f v -> DMap k f -> DMap k f -> DSum k f :*: DMap k f
go k v
kr f v
xr DMap k f
rl DMap k f
rr
      in (DSum k f
km DSum k f -> DMap k f -> DSum k f :*: DMap k f
forall a b. a -> b -> a :*: b
:*: k v -> f v -> DMap k f -> DMap k f -> DMap k f
forall k (k :: k -> *) (v :: k) (f :: k -> *).
k v -> f v -> DMap k f -> DMap k f -> DMap k f
balance k v
k f v
x DMap k f
l DMap k f
r')

-- | /O(log n)/. Delete and find the maximal element.
--
-- > deleteFindMax (fromList [(5,"a"), (3,"b"), (10,"c")]) == ((10,"c"), fromList [(3,"b"), (5,"a")])
-- > deleteFindMax empty                                      Error: can not return the maximal element of an empty map

deleteFindMax :: DMap k f -> (DSum k f, DMap k f)
deleteFindMax :: DMap k f -> (DSum k f, DMap k f)
deleteFindMax DMap k f
t
  = case DMap k f
t of
      Bin Int
_ k v
k f v
x DMap k f
l DMap k f
Tip -> (k v
k k v -> f v -> DSum k f
forall k (tag :: k -> *) (f :: k -> *) (a :: k).
tag a -> f a -> DSum tag f
:=> f v
x,DMap k f
l)
      Bin Int
_ k v
k f v
x DMap k f
l DMap k f
r   -> let (DSum k f
km,DMap k f
r') = DMap k f -> (DSum k f, DMap k f)
forall k (k :: k -> *) (f :: k -> *).
DMap k f -> (DSum k f, DMap k f)
deleteFindMax DMap k f
r in (DSum k f
km,k v -> f v -> DMap k f -> DMap k f -> DMap k f
forall k (k :: k -> *) (v :: k) (f :: k -> *).
k v -> f v -> DMap k f -> DMap k f -> DMap k f
balance k v
k f v
x DMap k f
l DMap k f
r')
      DMap k f
Tip             -> ([Char] -> DSum k f
forall a. HasCallStack => [Char] -> a
error [Char]
"Map.deleteFindMax: can not return the maximal element of an empty map", DMap k f
forall k (k :: k -> *) (f :: k -> *). DMap k f
Tip)


{--------------------------------------------------------------------
  [balance l x r] balances two trees with value x.
  The sizes of the trees should balance after decreasing the
  size of one of them. (a rotation).

  [delta] is the maximal relative difference between the sizes of
          two trees, it corresponds with the [w] in Adams' paper.
  [ratio] is the ratio between an outer and inner sibling of the
          heavier subtree in an unbalanced setting. It determines
          whether a double or single rotation should be performed
          to restore balance. It corresponds with the inverse
          of $\alpha$ in Adam's article.

  Note that:
  - [delta] should be larger than 4.646 with a [ratio] of 2.
  - [delta] should be larger than 3.745 with a [ratio] of 1.534.

  - A lower [delta] leads to a more 'perfectly' balanced tree.
  - A higher [delta] performs less rebalancing.

  - Balancing is automatic for random data and a balancing
    scheme is only necessary to avoid pathological worst cases.
    Almost any choice will do, and in practice, a rather large
    [delta] may perform better than smaller one.

  Note: in contrast to Adam's paper, we use a ratio of (at least) [2]
  to decide whether a single or double rotation is needed. Although
  he actually proves that this ratio is needed to maintain the
  invariants, his implementation uses an invalid ratio of [1].
--------------------------------------------------------------------}
delta,ratio :: Int
delta :: Int
delta = Int
4
ratio :: Int
ratio = Int
2

balance :: k v -> f v -> DMap k f -> DMap k f -> DMap k f
balance :: k v -> f v -> DMap k f -> DMap k f -> DMap k f
balance k v
k f v
x DMap k f
l DMap k f
r
  | Int
sizeL Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
sizeR Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
1    = Int -> k v -> f v -> DMap k f -> DMap k f -> DMap k f
forall k (k :: k -> *) (v :: k) (f :: k -> *).
Int -> k v -> f v -> DMap k f -> DMap k f -> DMap k f
Bin Int
sizeX k v
k f v
x DMap k f
l DMap k f
r
  | Int
sizeR Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
deltaInt -> Int -> Int
forall a. Num a => a -> a -> a
*Int
sizeL  = k v -> f v -> DMap k f -> DMap k f -> DMap k f
forall k (k :: k -> *) (v :: k) (f :: k -> *).
k v -> f v -> DMap k f -> DMap k f -> DMap k f
rotateL k v
k f v
x DMap k f
l DMap k f
r
  | Int
sizeL Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
deltaInt -> Int -> Int
forall a. Num a => a -> a -> a
*Int
sizeR  = k v -> f v -> DMap k f -> DMap k f -> DMap k f
forall k (k :: k -> *) (v :: k) (f :: k -> *).
k v -> f v -> DMap k f -> DMap k f -> DMap k f
rotateR k v
k f v
x DMap k f
l DMap k f
r
  | Bool
otherwise             = Int -> k v -> f v -> DMap k f -> DMap k f -> DMap k f
forall k (k :: k -> *) (v :: k) (f :: k -> *).
Int -> k v -> f v -> DMap k f -> DMap k f -> DMap k f
Bin Int
sizeX k v
k f v
x DMap k f
l DMap k f
r
  where
    sizeL :: Int
sizeL = DMap k f -> Int
forall k (k :: k -> *) (f :: k -> *). DMap k f -> Int
size DMap k f
l
    sizeR :: Int
sizeR = DMap k f -> Int
forall k (k :: k -> *) (f :: k -> *). DMap k f -> Int
size DMap k f
r
    sizeX :: Int
sizeX = Int
sizeL Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
sizeR Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1

-- rotate
rotateL :: k v -> f v -> DMap k f -> DMap k f -> DMap k f
rotateL :: k v -> f v -> DMap k f -> DMap k f -> DMap k f
rotateL k v
k f v
x DMap k f
l r :: DMap k f
r@(Bin Int
_ k v
_ f v
_ DMap k f
ly DMap k f
ry)
  | DMap k f -> Int
forall k (k :: k -> *) (f :: k -> *). DMap k f -> Int
size DMap k f
ly Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
ratioInt -> Int -> Int
forall a. Num a => a -> a -> a
*DMap k f -> Int
forall k (k :: k -> *) (f :: k -> *). DMap k f -> Int
size DMap k f
ry = k v -> f v -> DMap k f -> DMap k f -> DMap k f
forall k (k :: k -> *) (v :: k) (f :: k -> *).
k v -> f v -> DMap k f -> DMap k f -> DMap k f
singleL k v
k f v
x DMap k f
l DMap k f
r
  | Bool
otherwise               = k v -> f v -> DMap k f -> DMap k f -> DMap k f
forall k (k :: k -> *) (v :: k) (f :: k -> *).
k v -> f v -> DMap k f -> DMap k f -> DMap k f
doubleL k v
k f v
x DMap k f
l DMap k f
r
rotateL k v
_ f v
_ DMap k f
_ DMap k f
Tip = [Char] -> DMap k f
forall a. HasCallStack => [Char] -> a
error [Char]
"rotateL Tip"

rotateR :: k v -> f v -> DMap k f -> DMap k f -> DMap k f
rotateR :: k v -> f v -> DMap k f -> DMap k f -> DMap k f
rotateR k v
k f v
x l :: DMap k f
l@(Bin Int
_ k v
_ f v
_ DMap k f
ly DMap k f
ry) DMap k f
r
  | DMap k f -> Int
forall k (k :: k -> *) (f :: k -> *). DMap k f -> Int
size DMap k f
ry Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
ratioInt -> Int -> Int
forall a. Num a => a -> a -> a
*DMap k f -> Int
forall k (k :: k -> *) (f :: k -> *). DMap k f -> Int
size DMap k f
ly = k v -> f v -> DMap k f -> DMap k f -> DMap k f
forall k (k :: k -> *) (v :: k) (f :: k -> *).
k v -> f v -> DMap k f -> DMap k f -> DMap k f
singleR k v
k f v
x DMap k f
l DMap k f
r
  | Bool
otherwise               = k v -> f v -> DMap k f -> DMap k f -> DMap k f
forall k (k :: k -> *) (v :: k) (f :: k -> *).
k v -> f v -> DMap k f -> DMap k f -> DMap k f
doubleR k v
k f v
x DMap k f
l DMap k f
r
rotateR k v
_ f v
_ DMap k f
Tip DMap k f
_ = [Char] -> DMap k f
forall a. HasCallStack => [Char] -> a
error [Char]
"rotateR Tip"

-- basic rotations
singleL, singleR :: k v -> f v -> DMap k f -> DMap k f -> DMap k f
singleL :: k v -> f v -> DMap k f -> DMap k f -> DMap k f
singleL k v
k1 f v
x1 DMap k f
t1 (Bin Int
_ k v
k2 f v
x2 DMap k f
t2 DMap k f
t3)  = k v -> f v -> DMap k f -> DMap k f -> DMap k f
forall k (k :: k -> *) (v :: k) (f :: k -> *).
k v -> f v -> DMap k f -> DMap k f -> DMap k f
bin k v
k2 f v
x2 (k v -> f v -> DMap k f -> DMap k f -> DMap k f
forall k (k :: k -> *) (v :: k) (f :: k -> *).
k v -> f v -> DMap k f -> DMap k f -> DMap k f
bin k v
k1 f v
x1 DMap k f
t1 DMap k f
t2) DMap k f
t3
singleL k v
_ f v
_ DMap k f
_ DMap k f
Tip = [Char] -> DMap k f
forall a. HasCallStack => [Char] -> a
error [Char]
"singleL Tip"
singleR :: k v -> f v -> DMap k f -> DMap k f -> DMap k f
singleR k v
k1 f v
x1 (Bin Int
_ k v
k2 f v
x2 DMap k f
t1 DMap k f
t2) DMap k f
t3  = k v -> f v -> DMap k f -> DMap k f -> DMap k f
forall k (k :: k -> *) (v :: k) (f :: k -> *).
k v -> f v -> DMap k f -> DMap k f -> DMap k f
bin k v
k2 f v
x2 DMap k f
t1 (k v -> f v -> DMap k f -> DMap k f -> DMap k f
forall k (k :: k -> *) (v :: k) (f :: k -> *).
k v -> f v -> DMap k f -> DMap k f -> DMap k f
bin k v
k1 f v
x1 DMap k f
t2 DMap k f
t3)
singleR k v
_ f v
_ DMap k f
Tip DMap k f
_ = [Char] -> DMap k f
forall a. HasCallStack => [Char] -> a
error [Char]
"singleR Tip"

doubleL, doubleR :: k v -> f v -> DMap k f -> DMap k f -> DMap k f
doubleL :: k v -> f v -> DMap k f -> DMap k f -> DMap k f
doubleL k v
k1 f v
x1 DMap k f
t1 (Bin Int
_ k v
k2 f v
x2 (Bin Int
_ k v
k3 f v
x3 DMap k f
t2 DMap k f
t3) DMap k f
t4) = k v -> f v -> DMap k f -> DMap k f -> DMap k f
forall k (k :: k -> *) (v :: k) (f :: k -> *).
k v -> f v -> DMap k f -> DMap k f -> DMap k f
bin k v
k3 f v
x3 (k v -> f v -> DMap k f -> DMap k f -> DMap k f
forall k (k :: k -> *) (v :: k) (f :: k -> *).
k v -> f v -> DMap k f -> DMap k f -> DMap k f
bin k v
k1 f v
x1 DMap k f
t1 DMap k f
t2) (k v -> f v -> DMap k f -> DMap k f -> DMap k f
forall k (k :: k -> *) (v :: k) (f :: k -> *).
k v -> f v -> DMap k f -> DMap k f -> DMap k f
bin k v
k2 f v
x2 DMap k f
t3 DMap k f
t4)
doubleL k v
_ f v
_ DMap k f
_ DMap k f
_ = [Char] -> DMap k f
forall a. HasCallStack => [Char] -> a
error [Char]
"doubleL"
doubleR :: k v -> f v -> DMap k f -> DMap k f -> DMap k f
doubleR k v
k1 f v
x1 (Bin Int
_ k v
k2 f v
x2 DMap k f
t1 (Bin Int
_ k v
k3 f v
x3 DMap k f
t2 DMap k f
t3)) DMap k f
t4 = k v -> f v -> DMap k f -> DMap k f -> DMap k f
forall k (k :: k -> *) (v :: k) (f :: k -> *).
k v -> f v -> DMap k f -> DMap k f -> DMap k f
bin k v
k3 f v
x3 (k v -> f v -> DMap k f -> DMap k f -> DMap k f
forall k (k :: k -> *) (v :: k) (f :: k -> *).
k v -> f v -> DMap k f -> DMap k f -> DMap k f
bin k v
k2 f v
x2 DMap k f
t1 DMap k f
t2) (k v -> f v -> DMap k f -> DMap k f -> DMap k f
forall k (k :: k -> *) (v :: k) (f :: k -> *).
k v -> f v -> DMap k f -> DMap k f -> DMap k f
bin k v
k1 f v
x1 DMap k f
t3 DMap k f
t4)
doubleR k v
_ f v
_ DMap k f
_ DMap k f
_ = [Char] -> DMap k f
forall a. HasCallStack => [Char] -> a
error [Char]
"doubleR"

{--------------------------------------------------------------------
  The bin constructor maintains the size of the tree
--------------------------------------------------------------------}
bin :: k v -> f v -> DMap k f -> DMap k f -> DMap k f
bin :: k v -> f v -> DMap k f -> DMap k f -> DMap k f
bin k v
k f v
x DMap k f
l DMap k f
r
  = Int -> k v -> f v -> DMap k f -> DMap k f -> DMap k f
forall k (k :: k -> *) (v :: k) (f :: k -> *).
Int -> k v -> f v -> DMap k f -> DMap k f -> DMap k f
Bin (DMap k f -> Int
forall k (k :: k -> *) (f :: k -> *). DMap k f -> Int
size DMap k f
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ DMap k f -> Int
forall k (k :: k -> *) (f :: k -> *). DMap k f -> Int
size DMap k f
r Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) k v
k f v
x DMap k f
l DMap k f
r

{--------------------------------------------------------------------
  Utility functions that return sub-ranges of the original
  tree. Some functions take a comparison function as argument to
  allow comparisons against infinite values. A function [cmplo k]
  should be read as [compare lo k].

  [trim cmplo cmphi t]  A tree that is either empty or where [cmplo k == LT]
                        and [cmphi k == GT] for the key [k] of the root.
  [filterGt cmp t]      A tree where for all keys [k]. [cmp k == LT]
  [filterLt cmp t]      A tree where for all keys [k]. [cmp k == GT]

  [split k t]           Returns two trees [l] and [r] where all keys
                        in [l] are <[k] and all keys in [r] are >[k].
  [splitLookup k t]     Just like [split] but also returns whether [k]
                        was found in the tree.
--------------------------------------------------------------------}

{--------------------------------------------------------------------
  [trim lo hi t] trims away all subtrees that surely contain no
  values between the range [lo] to [hi]. The returned tree is either
  empty or the key of the root is between @lo@ and @hi@.
--------------------------------------------------------------------}
trim :: (Some k -> Ordering) -> (Some k -> Ordering) -> DMap k f -> DMap k f
trim :: (Some k -> Ordering)
-> (Some k -> Ordering) -> DMap k f -> DMap k f
trim Some k -> Ordering
_     Some k -> Ordering
_     DMap k f
Tip = DMap k f
forall k (k :: k -> *) (f :: k -> *). DMap k f
Tip
trim Some k -> Ordering
cmplo Some k -> Ordering
cmphi t :: DMap k f
t@(Bin Int
_ k v
kx f v
_ DMap k f
l DMap k f
r)
  = case Some k -> Ordering
cmplo (k v -> Some k
forall k (tag :: k -> *) (a :: k). tag a -> Some tag
mkSome k v
kx) of
      Ordering
LT -> case Some k -> Ordering
cmphi (k v -> Some k
forall k (tag :: k -> *) (a :: k). tag a -> Some tag
mkSome k v
kx) of
              Ordering
GT -> DMap k f
t
              Ordering
_  -> (Some k -> Ordering)
-> (Some k -> Ordering) -> DMap k f -> DMap k f
forall k (k :: k -> *) (f :: k -> *).
(Some k -> Ordering)
-> (Some k -> Ordering) -> DMap k f -> DMap k f
trim Some k -> Ordering
cmplo Some k -> Ordering
cmphi DMap k f
l
      Ordering
_  -> (Some k -> Ordering)
-> (Some k -> Ordering) -> DMap k f -> DMap k f
forall k (k :: k -> *) (f :: k -> *).
(Some k -> Ordering)
-> (Some k -> Ordering) -> DMap k f -> DMap k f
trim Some k -> Ordering
cmplo Some k -> Ordering
cmphi DMap k f
r

trimLookupLo :: GCompare k => Some k -> (Some k -> Ordering) -> DMap k f -> (Maybe (DSum k f), DMap k f)
trimLookupLo :: Some k
-> (Some k -> Ordering) -> DMap k f -> (Maybe (DSum k f), DMap k f)
trimLookupLo Some k
_  Some k -> Ordering
_     DMap k f
Tip = (Maybe (DSum k f)
forall a. Maybe a
Nothing,DMap k f
forall k (k :: k -> *) (f :: k -> *). DMap k f
Tip)
trimLookupLo Some k
lo Some k -> Ordering
cmphi t :: DMap k f
t@(Bin Int
_ k v
kx f v
x DMap k f
l DMap k f
r)
  = case Some k -> Some k -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Some k
lo (k v -> Some k
forall k (tag :: k -> *) (a :: k). tag a -> Some tag
mkSome k v
kx) of
      Ordering
LT -> case Some k -> Ordering
cmphi (k v -> Some k
forall k (tag :: k -> *) (a :: k). tag a -> Some tag
mkSome k v
kx) of
              Ordering
GT -> (Some k -> DMap k f -> Maybe (DSum k f)
forall k k (k :: k -> *) (f :: k -> *) (v :: k).
GCompare k =>
Some k -> DMap k f -> Maybe (DSum k f)
lookupAssoc Some k
lo DMap k f
t, DMap k f
t)
              Ordering
_  -> Some k
-> (Some k -> Ordering) -> DMap k f -> (Maybe (DSum k f), DMap k f)
forall k (k :: k -> *) (f :: k -> *).
GCompare k =>
Some k
-> (Some k -> Ordering) -> DMap k f -> (Maybe (DSum k f), DMap k f)
trimLookupLo Some k
lo Some k -> Ordering
cmphi DMap k f
l
      Ordering
GT -> Some k
-> (Some k -> Ordering) -> DMap k f -> (Maybe (DSum k f), DMap k f)
forall k (k :: k -> *) (f :: k -> *).
GCompare k =>
Some k
-> (Some k -> Ordering) -> DMap k f -> (Maybe (DSum k f), DMap k f)
trimLookupLo Some k
lo Some k -> Ordering
cmphi DMap k f
r
      Ordering
EQ -> (DSum k f -> Maybe (DSum k f)
forall a. a -> Maybe a
Just (k v
kx k v -> f v -> DSum k f
forall k (tag :: k -> *) (f :: k -> *) (a :: k).
tag a -> f a -> DSum tag f
:=> f v
x),(Some k -> Ordering)
-> (Some k -> Ordering) -> DMap k f -> DMap k f
forall k (k :: k -> *) (f :: k -> *).
(Some k -> Ordering)
-> (Some k -> Ordering) -> DMap k f -> DMap k f
trim (Some k -> Some k -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Some k
lo) Some k -> Ordering
cmphi DMap k f
r)


{--------------------------------------------------------------------
  [filterGt k t] filter all keys >[k] from tree [t]
  [filterLt k t] filter all keys <[k] from tree [t]
--------------------------------------------------------------------}
filterGt :: GCompare k => (Some k -> Ordering) -> DMap k f -> DMap k f
filterGt :: (Some k -> Ordering) -> DMap k f -> DMap k f
filterGt Some k -> Ordering
cmp = DMap k f -> DMap k f
go
  where
    go :: DMap k f -> DMap k f
go DMap k f
Tip              = DMap k f
forall k (k :: k -> *) (f :: k -> *). DMap k f
Tip
    go (Bin Int
_ k v
kx f v
x DMap k f
l DMap k f
r) = case Some k -> Ordering
cmp (k v -> Some k
forall k (tag :: k -> *) (a :: k). tag a -> Some tag
mkSome k v
kx) of
              Ordering
LT -> k v -> f v -> DMap k f -> DMap k f -> DMap k f
forall k (k :: k -> *) (v :: k) (f :: k -> *).
GCompare k =>
k v -> f v -> DMap k f -> DMap k f -> DMap k f
combine k v
kx f v
x (DMap k f -> DMap k f
go DMap k f
l) DMap k f
r
              Ordering
GT -> DMap k f -> DMap k f
go DMap k f
r
              Ordering
EQ -> DMap k f
r

filterLt :: GCompare k => (Some k -> Ordering) -> DMap k f -> DMap k f
filterLt :: (Some k -> Ordering) -> DMap k f -> DMap k f
filterLt Some k -> Ordering
cmp = DMap k f -> DMap k f
go
  where
    go :: DMap k f -> DMap k f
go DMap k f
Tip              = DMap k f
forall k (k :: k -> *) (f :: k -> *). DMap k f
Tip
    go (Bin Int
_ k v
kx f v
x DMap k f
l DMap k f
r) = case Some k -> Ordering
cmp (k v -> Some k
forall k (tag :: k -> *) (a :: k). tag a -> Some tag
mkSome k v
kx) of
          Ordering
LT -> DMap k f -> DMap k f
go DMap k f
l
          Ordering
GT -> k v -> f v -> DMap k f -> DMap k f -> DMap k f
forall k (k :: k -> *) (v :: k) (f :: k -> *).
GCompare k =>
k v -> f v -> DMap k f -> DMap k f -> DMap k f
combine k v
kx f v
x DMap k f
l (DMap k f -> DMap k f
go DMap k f
r)
          Ordering
EQ -> DMap k f
l