{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DeriveFoldable #-}
{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE CPP #-}

{-# OPTIONS_GHC -Wall #-}

{-|
Maps with disjoint sets as the key. The type in this module can be
roughly understood as:

> DisjointMap k v ≈ Map (Set k) v

Internally, @DisjointMap@ is implemented like a disjoint set
but the data structure that maps representatives to their rank also holds the value
associated with that representative element. Additionally, it holds the set
of all keys that in the same equivalence class as the representative.
This makes it possible to implementat functions like @foldlWithKeys\'@
efficiently.

-}

module Data.DisjointMap
  ( DisjointMap
    -- * Construction
  , empty
  , singleton
  , singletons
  , insert
  , union
  , unionWeakly
    -- * Query
  , lookup
  , lookup'
  , representative
  , representative'
    -- * Conversion
  , toLists
  , toSets
  , fromSets
  , pretty
  , prettyList
  , foldlWithKeys'
    -- * Tutorial
    -- $tutorial
  ) where

import Prelude hiding (lookup)
import Control.Monad.Trans.State.Strict
import Control.Monad.Trans.Maybe
import Control.Monad.Trans.Class
import Control.Monad

import Data.Map (Map)
import Data.Set (Set)
import Data.Bifunctor (first)
import Data.Foldable (Foldable)
import Data.Maybe (fromMaybe)
import Data.Foldable (foldlM)
import qualified Data.Semigroup as SG
import qualified Data.Map.Strict as M
import qualified Data.Set as S
import qualified GHC.OldList as L
import qualified Data.Foldable as F

-- | A map having disjoints sets of @k@ as keys and
--   @v@ as values.
data DisjointMap k v = DisjointMap
  !(Map k k) -- parents and values
  !(Map k (Ranked k v)) -- ranks
  deriving (forall a b. a -> DisjointMap k b -> DisjointMap k a
forall a b. (a -> b) -> DisjointMap k a -> DisjointMap k b
forall k a b. a -> DisjointMap k b -> DisjointMap k a
forall k a b. (a -> b) -> DisjointMap k a -> DisjointMap k b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: forall a b. a -> DisjointMap k b -> DisjointMap k a
$c<$ :: forall k a b. a -> DisjointMap k b -> DisjointMap k a
fmap :: forall a b. (a -> b) -> DisjointMap k a -> DisjointMap k b
$cfmap :: forall k a b. (a -> b) -> DisjointMap k a -> DisjointMap k b
Functor,forall a. DisjointMap k a -> Bool
forall k a. Eq a => a -> DisjointMap k a -> Bool
forall k a. Num a => DisjointMap k a -> a
forall k a. Ord a => DisjointMap k a -> a
forall m a. Monoid m => (a -> m) -> DisjointMap k a -> m
forall k m. Monoid m => DisjointMap k m -> m
forall k a. DisjointMap k a -> Bool
forall k a. DisjointMap k a -> Int
forall k a. DisjointMap k a -> [a]
forall a b. (a -> b -> b) -> b -> DisjointMap k a -> b
forall k a. (a -> a -> a) -> DisjointMap k a -> a
forall k m a. Monoid m => (a -> m) -> DisjointMap k a -> m
forall k b a. (b -> a -> b) -> b -> DisjointMap k a -> b
forall k a b. (a -> b -> b) -> b -> DisjointMap k a -> b
forall (t :: * -> *).
(forall m. Monoid m => t m -> m)
-> (forall m a. Monoid m => (a -> m) -> t a -> m)
-> (forall m a. Monoid m => (a -> m) -> t a -> m)
-> (forall a b. (a -> b -> b) -> b -> t a -> b)
-> (forall a b. (a -> b -> b) -> b -> t a -> b)
-> (forall b a. (b -> a -> b) -> b -> t a -> b)
-> (forall b a. (b -> a -> b) -> b -> t a -> b)
-> (forall a. (a -> a -> a) -> t a -> a)
-> (forall a. (a -> a -> a) -> t a -> a)
-> (forall a. t a -> [a])
-> (forall a. t a -> Bool)
-> (forall a. t a -> Int)
-> (forall a. Eq a => a -> t a -> Bool)
-> (forall a. Ord a => t a -> a)
-> (forall a. Ord a => t a -> a)
-> (forall a. Num a => t a -> a)
-> (forall a. Num a => t a -> a)
-> Foldable t
product :: forall a. Num a => DisjointMap k a -> a
$cproduct :: forall k a. Num a => DisjointMap k a -> a
sum :: forall a. Num a => DisjointMap k a -> a
$csum :: forall k a. Num a => DisjointMap k a -> a
minimum :: forall a. Ord a => DisjointMap k a -> a
$cminimum :: forall k a. Ord a => DisjointMap k a -> a
maximum :: forall a. Ord a => DisjointMap k a -> a
$cmaximum :: forall k a. Ord a => DisjointMap k a -> a
elem :: forall a. Eq a => a -> DisjointMap k a -> Bool
$celem :: forall k a. Eq a => a -> DisjointMap k a -> Bool
length :: forall a. DisjointMap k a -> Int
$clength :: forall k a. DisjointMap k a -> Int
null :: forall a. DisjointMap k a -> Bool
$cnull :: forall k a. DisjointMap k a -> Bool
toList :: forall a. DisjointMap k a -> [a]
$ctoList :: forall k a. DisjointMap k a -> [a]
foldl1 :: forall a. (a -> a -> a) -> DisjointMap k a -> a
$cfoldl1 :: forall k a. (a -> a -> a) -> DisjointMap k a -> a
foldr1 :: forall a. (a -> a -> a) -> DisjointMap k a -> a
$cfoldr1 :: forall k a. (a -> a -> a) -> DisjointMap k a -> a
foldl' :: forall b a. (b -> a -> b) -> b -> DisjointMap k a -> b
$cfoldl' :: forall k b a. (b -> a -> b) -> b -> DisjointMap k a -> b
foldl :: forall b a. (b -> a -> b) -> b -> DisjointMap k a -> b
$cfoldl :: forall k b a. (b -> a -> b) -> b -> DisjointMap k a -> b
foldr' :: forall a b. (a -> b -> b) -> b -> DisjointMap k a -> b
$cfoldr' :: forall k a b. (a -> b -> b) -> b -> DisjointMap k a -> b
foldr :: forall a b. (a -> b -> b) -> b -> DisjointMap k a -> b
$cfoldr :: forall k a b. (a -> b -> b) -> b -> DisjointMap k a -> b
foldMap' :: forall m a. Monoid m => (a -> m) -> DisjointMap k a -> m
$cfoldMap' :: forall k m a. Monoid m => (a -> m) -> DisjointMap k a -> m
foldMap :: forall m a. Monoid m => (a -> m) -> DisjointMap k a -> m
$cfoldMap :: forall k m a. Monoid m => (a -> m) -> DisjointMap k a -> m
fold :: forall m. Monoid m => DisjointMap k m -> m
$cfold :: forall k m. Monoid m => DisjointMap k m -> m
Foldable,forall k. Functor (DisjointMap k)
forall k. Foldable (DisjointMap k)
forall k (m :: * -> *) a.
Monad m =>
DisjointMap k (m a) -> m (DisjointMap k a)
forall k (f :: * -> *) a.
Applicative f =>
DisjointMap k (f a) -> f (DisjointMap k a)
forall k (m :: * -> *) a b.
Monad m =>
(a -> m b) -> DisjointMap k a -> m (DisjointMap k b)
forall k (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> DisjointMap k a -> f (DisjointMap k b)
forall (t :: * -> *).
Functor t
-> Foldable t
-> (forall (f :: * -> *) a b.
    Applicative f =>
    (a -> f b) -> t a -> f (t b))
-> (forall (f :: * -> *) a. Applicative f => t (f a) -> f (t a))
-> (forall (m :: * -> *) a b.
    Monad m =>
    (a -> m b) -> t a -> m (t b))
-> (forall (m :: * -> *) a. Monad m => t (m a) -> m (t a))
-> Traversable t
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> DisjointMap k a -> f (DisjointMap k b)
sequence :: forall (m :: * -> *) a.
Monad m =>
DisjointMap k (m a) -> m (DisjointMap k a)
$csequence :: forall k (m :: * -> *) a.
Monad m =>
DisjointMap k (m a) -> m (DisjointMap k a)
mapM :: forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> DisjointMap k a -> m (DisjointMap k b)
$cmapM :: forall k (m :: * -> *) a b.
Monad m =>
(a -> m b) -> DisjointMap k a -> m (DisjointMap k b)
sequenceA :: forall (f :: * -> *) a.
Applicative f =>
DisjointMap k (f a) -> f (DisjointMap k a)
$csequenceA :: forall k (f :: * -> *) a.
Applicative f =>
DisjointMap k (f a) -> f (DisjointMap k a)
traverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> DisjointMap k a -> f (DisjointMap k b)
$ctraverse :: forall k (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> DisjointMap k a -> f (DisjointMap k b)
Traversable)

-- the name ranked is no longer totally appropriate since
-- a set of keys has been added in here as well.
data Ranked k v = Ranked {-# UNPACK #-} !Int !(Set k) !v
  deriving (forall a b. a -> Ranked k b -> Ranked k a
forall a b. (a -> b) -> Ranked k a -> Ranked k b
forall k a b. a -> Ranked k b -> Ranked k a
forall k a b. (a -> b) -> Ranked k a -> Ranked k b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: forall a b. a -> Ranked k b -> Ranked k a
$c<$ :: forall k a b. a -> Ranked k b -> Ranked k a
fmap :: forall a b. (a -> b) -> Ranked k a -> Ranked k b
$cfmap :: forall k a b. (a -> b) -> Ranked k a -> Ranked k b
Functor,forall a. Ranked k a -> Bool
forall k a. Eq a => a -> Ranked k a -> Bool
forall k a. Num a => Ranked k a -> a
forall k a. Ord a => Ranked k a -> a
forall m a. Monoid m => (a -> m) -> Ranked k a -> m
forall k m. Monoid m => Ranked k m -> m
forall k a. Ranked k a -> Bool
forall k a. Ranked k a -> Int
forall k a. Ranked k a -> [a]
forall a b. (a -> b -> b) -> b -> Ranked k a -> b
forall k a. (a -> a -> a) -> Ranked k a -> a
forall k m a. Monoid m => (a -> m) -> Ranked k a -> m
forall k b a. (b -> a -> b) -> b -> Ranked k a -> b
forall k a b. (a -> b -> b) -> b -> Ranked k a -> b
forall (t :: * -> *).
(forall m. Monoid m => t m -> m)
-> (forall m a. Monoid m => (a -> m) -> t a -> m)
-> (forall m a. Monoid m => (a -> m) -> t a -> m)
-> (forall a b. (a -> b -> b) -> b -> t a -> b)
-> (forall a b. (a -> b -> b) -> b -> t a -> b)
-> (forall b a. (b -> a -> b) -> b -> t a -> b)
-> (forall b a. (b -> a -> b) -> b -> t a -> b)
-> (forall a. (a -> a -> a) -> t a -> a)
-> (forall a. (a -> a -> a) -> t a -> a)
-> (forall a. t a -> [a])
-> (forall a. t a -> Bool)
-> (forall a. t a -> Int)
-> (forall a. Eq a => a -> t a -> Bool)
-> (forall a. Ord a => t a -> a)
-> (forall a. Ord a => t a -> a)
-> (forall a. Num a => t a -> a)
-> (forall a. Num a => t a -> a)
-> Foldable t
product :: forall a. Num a => Ranked k a -> a
$cproduct :: forall k a. Num a => Ranked k a -> a
sum :: forall a. Num a => Ranked k a -> a
$csum :: forall k a. Num a => Ranked k a -> a
minimum :: forall a. Ord a => Ranked k a -> a
$cminimum :: forall k a. Ord a => Ranked k a -> a
maximum :: forall a. Ord a => Ranked k a -> a
$cmaximum :: forall k a. Ord a => Ranked k a -> a
elem :: forall a. Eq a => a -> Ranked k a -> Bool
$celem :: forall k a. Eq a => a -> Ranked k a -> Bool
length :: forall a. Ranked k a -> Int
$clength :: forall k a. Ranked k a -> Int
null :: forall a. Ranked k a -> Bool
$cnull :: forall k a. Ranked k a -> Bool
toList :: forall a. Ranked k a -> [a]
$ctoList :: forall k a. Ranked k a -> [a]
foldl1 :: forall a. (a -> a -> a) -> Ranked k a -> a
$cfoldl1 :: forall k a. (a -> a -> a) -> Ranked k a -> a
foldr1 :: forall a. (a -> a -> a) -> Ranked k a -> a
$cfoldr1 :: forall k a. (a -> a -> a) -> Ranked k a -> a
foldl' :: forall b a. (b -> a -> b) -> b -> Ranked k a -> b
$cfoldl' :: forall k b a. (b -> a -> b) -> b -> Ranked k a -> b
foldl :: forall b a. (b -> a -> b) -> b -> Ranked k a -> b
$cfoldl :: forall k b a. (b -> a -> b) -> b -> Ranked k a -> b
foldr' :: forall a b. (a -> b -> b) -> b -> Ranked k a -> b
$cfoldr' :: forall k a b. (a -> b -> b) -> b -> Ranked k a -> b
foldr :: forall a b. (a -> b -> b) -> b -> Ranked k a -> b
$cfoldr :: forall k a b. (a -> b -> b) -> b -> Ranked k a -> b
foldMap' :: forall m a. Monoid m => (a -> m) -> Ranked k a -> m
$cfoldMap' :: forall k m a. Monoid m => (a -> m) -> Ranked k a -> m
foldMap :: forall m a. Monoid m => (a -> m) -> Ranked k a -> m
$cfoldMap :: forall k m a. Monoid m => (a -> m) -> Ranked k a -> m
fold :: forall m. Monoid m => Ranked k m -> m
$cfold :: forall k m. Monoid m => Ranked k m -> m
Foldable,forall k. Functor (Ranked k)
forall k. Foldable (Ranked k)
forall k (m :: * -> *) a.
Monad m =>
Ranked k (m a) -> m (Ranked k a)
forall k (f :: * -> *) a.
Applicative f =>
Ranked k (f a) -> f (Ranked k a)
forall k (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Ranked k a -> m (Ranked k b)
forall k (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Ranked k a -> f (Ranked k b)
forall (t :: * -> *).
Functor t
-> Foldable t
-> (forall (f :: * -> *) a b.
    Applicative f =>
    (a -> f b) -> t a -> f (t b))
-> (forall (f :: * -> *) a. Applicative f => t (f a) -> f (t a))
-> (forall (m :: * -> *) a b.
    Monad m =>
    (a -> m b) -> t a -> m (t b))
-> (forall (m :: * -> *) a. Monad m => t (m a) -> m (t a))
-> Traversable t
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Ranked k a -> f (Ranked k b)
sequence :: forall (m :: * -> *) a. Monad m => Ranked k (m a) -> m (Ranked k a)
$csequence :: forall k (m :: * -> *) a.
Monad m =>
Ranked k (m a) -> m (Ranked k a)
mapM :: forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Ranked k a -> m (Ranked k b)
$cmapM :: forall k (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Ranked k a -> m (Ranked k b)
sequenceA :: forall (f :: * -> *) a.
Applicative f =>
Ranked k (f a) -> f (Ranked k a)
$csequenceA :: forall k (f :: * -> *) a.
Applicative f =>
Ranked k (f a) -> f (Ranked k a)
traverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Ranked k a -> f (Ranked k b)
$ctraverse :: forall k (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Ranked k a -> f (Ranked k b)
Traversable)

instance (Ord k, Monoid v) => Monoid (DisjointMap k v) where
  mempty :: DisjointMap k v
mempty = forall k v. DisjointMap k v
empty

-- | This only satisfies the associativity law when the 'Monoid'
--   instance for @v@ is commutative.
instance (Ord k, Semigroup v) => SG.Semigroup (DisjointMap k v) where
  <> :: DisjointMap k v -> DisjointMap k v -> DisjointMap k v
(<>) = forall k v.
(Ord k, Semigroup v) =>
DisjointMap k v -> DisjointMap k v -> DisjointMap k v
append

-- technically, it should be possible to weaken the Ord constraint on v to
-- an Eq constraint
instance (Ord k, Ord v) => Eq (DisjointMap k v) where
  DisjointMap k v
a == :: DisjointMap k v -> DisjointMap k v -> Bool
== DisjointMap k v
b = forall a. Ord a => [a] -> Set a
S.fromList (forall k v. DisjointMap k v -> [(Set k, v)]
toSets DisjointMap k v
a) forall a. Eq a => a -> a -> Bool
== forall a. Ord a => [a] -> Set a
S.fromList (forall k v. DisjointMap k v -> [(Set k, v)]
toSets DisjointMap k v
b)

instance (Ord k, Ord v) => Ord (DisjointMap k v) where
  compare :: DisjointMap k v -> DisjointMap k v -> Ordering
compare DisjointMap k v
a DisjointMap k v
b = forall a. Ord a => a -> a -> Ordering
compare (forall a. Ord a => [a] -> Set a
S.fromList (forall k v. DisjointMap k v -> [(Set k, v)]
toSets DisjointMap k v
a)) (forall a. Ord a => [a] -> Set a
S.fromList (forall k v. DisjointMap k v -> [(Set k, v)]
toSets DisjointMap k v
b))

instance (Show k, Ord k, Show v) => Show (DisjointMap k v) where
  show :: DisjointMap k v -> String
show = forall k v. (Show k, Ord k, Show v) => DisjointMap k v -> String
showDisjointSet

fromSets :: Ord k => [(Set k,v)] -> Maybe (DisjointMap k v)
fromSets :: forall k v. Ord k => [(Set k, v)] -> Maybe (DisjointMap k v)
fromSets [(Set k, v)]
xs = case forall a. Ord a => [Set a] -> Maybe (Set a)
unionDistinctAll (forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> a
fst [(Set k, v)]
xs) of
  Maybe (Set k)
Nothing -> forall a. Maybe a
Nothing
  Just Set k
_ -> forall a. a -> Maybe a
Just (forall k v.
Ord k =>
[(Set k, v)] -> DisjointMap k v -> DisjointMap k v
unsafeFromSets [(Set k, v)]
xs forall k v. DisjointMap k v
empty)

unsafeFromSets :: Ord k => [(Set k,v)] -> DisjointMap k v -> DisjointMap k v
unsafeFromSets :: forall k v.
Ord k =>
[(Set k, v)] -> DisjointMap k v -> DisjointMap k v
unsafeFromSets [(Set k, v)]
ys !ds :: DisjointMap k v
ds@(DisjointMap Map k k
p Map k (Ranked k v)
r) = case [(Set k, v)]
ys of
  [] -> DisjointMap k v
ds
  (Set k
x,v
v) : [(Set k, v)]
xs -> case forall a. Set a -> Maybe a
setLookupMin Set k
x of
    Maybe k
Nothing -> forall k v.
Ord k =>
[(Set k, v)] -> DisjointMap k v -> DisjointMap k v
unsafeFromSets [(Set k, v)]
xs DisjointMap k v
ds
    Just k
m -> forall k v.
Ord k =>
[(Set k, v)] -> DisjointMap k v -> DisjointMap k v
unsafeFromSets [(Set k, v)]
xs forall a b. (a -> b) -> a -> b
$ forall k v. Map k k -> Map k (Ranked k v) -> DisjointMap k v
DisjointMap
      (forall k a. Ord k => Map k a -> Map k a -> Map k a
M.union (forall k a. (k -> a) -> Set k -> Map k a
M.fromSet (\k
_ -> k
m) Set k
x) Map k k
p)
      (forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert k
m (forall k v. Int -> Set k -> v -> Ranked k v
Ranked Int
0 Set k
x v
v) Map k (Ranked k v)
r)
  

unionDistinctAll :: Ord a => [Set a] -> Maybe (Set a)
unionDistinctAll :: forall a. Ord a => [Set a] -> Maybe (Set a)
unionDistinctAll = forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldlM forall a. Ord a => Set a -> Set a -> Maybe (Set a)
unionDistinct forall a. Set a
S.empty

unionDistinct :: Ord a => Set a -> Set a -> Maybe (Set a)
unionDistinct :: forall a. Ord a => Set a -> Set a -> Maybe (Set a)
unionDistinct Set a
a Set a
b = 
  let s :: Set a
s = forall a. Ord a => Set a -> Set a -> Set a
S.union Set a
a Set a
b
   in if forall a. Set a -> Int
S.size Set a
a forall a. Num a => a -> a -> a
+ forall a. Set a -> Int
S.size Set a
b forall a. Eq a => a -> a -> Bool
== forall a. Set a -> Int
S.size Set a
s
        then forall a. a -> Maybe a
Just Set a
s
        else forall a. Maybe a
Nothing

showDisjointSet :: (Show k, Ord k, Show v) => DisjointMap k v -> String
showDisjointSet :: forall k v. (Show k, Ord k, Show v) => DisjointMap k v -> String
showDisjointSet = forall a. Show a => a -> String
show forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k v. DisjointMap k v -> [([k], v)]
toLists

toLists :: DisjointMap k v -> [([k],v)]
toLists :: forall k v. DisjointMap k v -> [([k], v)]
toLists = (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmapforall b c a. (b -> c) -> (a -> b) -> a -> c
.forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first) forall a. Set a -> [a]
S.toList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k v. DisjointMap k v -> [(Set k, v)]
toSets

toSets :: DisjointMap k v -> [(Set k,v)]
toSets :: forall k v. DisjointMap k v -> [(Set k, v)]
toSets (DisjointMap Map k k
_ Map k (Ranked k v)
r) = forall a b k. (a -> b -> b) -> b -> Map k a -> b
M.foldr
  (\(Ranked Int
_ Set k
s v
v) [(Set k, v)]
xs -> (Set k
s,v
v) forall a. a -> [a] -> [a]
: [(Set k, v)]
xs) [] Map k (Ranked k v)
r

pretty :: (Show k, Show v) => DisjointMap k v -> String
pretty :: forall k v. (Show k, Show v) => DisjointMap k v -> String
pretty DisjointMap k v
dm = String
"{" forall a. [a] -> [a] -> [a]
++ forall a. [a] -> [[a]] -> [a]
L.intercalate String
", " (forall k v. (Show k, Show v) => DisjointMap k v -> [String]
prettyList DisjointMap k v
dm) forall a. [a] -> [a] -> [a]
++ String
"}"

prettyList :: (Show k, Show v) => DisjointMap k v -> [String]
prettyList :: forall k v. (Show k, Show v) => DisjointMap k v -> [String]
prettyList DisjointMap k v
dm = forall a b. (a -> b) -> [a] -> [b]
L.map (\(Set k
ks,v
v) -> String
"{" forall a. [a] -> [a] -> [a]
++ forall k. Show k => Set k -> String
commafied Set k
ks forall a. [a] -> [a] -> [a]
++ String
"} -> " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show v
v) (forall k v. DisjointMap k v -> [(Set k, v)]
toSets DisjointMap k v
dm)

commafied :: Show k => Set k -> String
commafied :: forall k. Show k => Set k -> String
commafied = forall (m :: * -> *) a. Monad m => m (m a) -> m a
join forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. a -> [a] -> [a]
L.intersperse String
"," forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map forall a. Show a => a -> String
show forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Set a -> [a]
S.toList

foldlWithKeys' :: (a -> Set k -> v -> a) -> a -> DisjointMap k v -> a
foldlWithKeys' :: forall a k v. (a -> Set k -> v -> a) -> a -> DisjointMap k v -> a
foldlWithKeys' a -> Set k -> v -> a
f a
a0 (DisjointMap Map k k
_ Map k (Ranked k v)
r) =
  forall a b k. (a -> b -> a) -> a -> Map k b -> a
M.foldl' (\a
a (Ranked Int
_ Set k
ks v
v) -> a -> Set k -> v -> a
f a
a Set k
ks v
v) a
a0 Map k (Ranked k v)
r

{-|
Create an equivalence relation between x and y. If either x or y
are not already in the disjoint set, they are first created
as singletons with a value that is 'mempty'.
-}
union :: (Ord k, Monoid v) => k -> k -> DisjointMap k v -> DisjointMap k v
union :: forall k v.
(Ord k, Monoid v) =>
k -> k -> DisjointMap k v -> DisjointMap k v
union !k
x !k
y DisjointMap k v
set = forall a b c. (a -> b -> c) -> b -> a -> c
flip forall s a. State s a -> s -> s
execState DisjointMap k v
set forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. MaybeT m a -> m (Maybe a)
runMaybeT forall a b. (a -> b) -> a -> b
$ do
  k
repx <- forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) s a. Monad m => (s -> (a, s)) -> StateT s m a
state forall a b. (a -> b) -> a -> b
$ forall k v.
(Ord k, Monoid v) =>
k -> DisjointMap k v -> (k, DisjointMap k v)
lookupCompressAdd k
x
  k
repy <- forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) s a. Monad m => (s -> (a, s)) -> StateT s m a
state forall a b. (a -> b) -> a -> b
$ forall k v.
(Ord k, Monoid v) =>
k -> DisjointMap k v -> (k, DisjointMap k v)
lookupCompressAdd k
y
  forall (f :: * -> *). Alternative f => Bool -> f ()
guard forall a b. (a -> b) -> a -> b
$ k
repx forall a. Eq a => a -> a -> Bool
/= k
repy
  DisjointMap Map k k
p Map k (Ranked k v)
r <- forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall (m :: * -> *) s. Monad m => StateT s m s
get
  let Ranked Int
rankx Set k
keysx v
valx = Map k (Ranked k v)
r forall k a. Ord k => Map k a -> k -> a
M.! k
repx
  let Ranked Int
ranky Set k
keysy v
valy = Map k (Ranked k v)
r forall k a. Ord k => Map k a -> k -> a
M.! k
repy
  let val :: v
val = forall a. Monoid a => a -> a -> a
mappend v
valx v
valy
      keys :: Set k
keys = forall a. Monoid a => a -> a -> a
mappend Set k
keysx Set k
keysy
  forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) s. Monad m => s -> StateT s m ()
put forall a b. (a -> b) -> a -> b
$! case forall a. Ord a => a -> a -> Ordering
compare Int
rankx Int
ranky of
    Ordering
LT -> let p' :: Map k k
p' = forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert k
repx k
repy Map k k
p
              r' :: Map k (Ranked k v)
r' = forall k a. Ord k => k -> Map k a -> Map k a
M.delete k
repx forall a b. (a -> b) -> a -> b
$! forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert k
repy (forall k v. Int -> Set k -> v -> Ranked k v
Ranked Int
ranky Set k
keys v
val) Map k (Ranked k v)
r
          in  forall k v. Map k k -> Map k (Ranked k v) -> DisjointMap k v
DisjointMap Map k k
p' Map k (Ranked k v)
r'
    Ordering
GT -> let p' :: Map k k
p' = forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert k
repy k
repx Map k k
p
              r' :: Map k (Ranked k v)
r' = forall k a. Ord k => k -> Map k a -> Map k a
M.delete k
repy forall a b. (a -> b) -> a -> b
$! forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert k
repx (forall k v. Int -> Set k -> v -> Ranked k v
Ranked Int
rankx Set k
keys v
val) Map k (Ranked k v)
r
          in  forall k v. Map k k -> Map k (Ranked k v) -> DisjointMap k v
DisjointMap Map k k
p' Map k (Ranked k v)
r'
    Ordering
EQ -> let p' :: Map k k
p' = forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert k
repx k
repy Map k k
p
              r' :: Map k (Ranked k v)
r' = forall k a. Ord k => k -> Map k a -> Map k a
M.delete k
repx forall a b. (a -> b) -> a -> b
$! forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert k
repy (forall k v. Int -> Set k -> v -> Ranked k v
Ranked (Int
ranky forall a. Num a => a -> a -> a
+ Int
1) Set k
keys v
val) Map k (Ranked k v)
r
          in  forall k v. Map k k -> Map k (Ranked k v) -> DisjointMap k v
DisjointMap Map k k
p' Map k (Ranked k v)
r'
{-|
Create an equivalence relation between x and y. If both x and y
are missing, do not create either of them. Otherwise, they will
both exist in the map.
-}
unionWeakly :: (Ord k, Semigroup v) => k -> k -> DisjointMap k v -> DisjointMap k v
unionWeakly :: forall k v.
(Ord k, Semigroup v) =>
k -> k -> DisjointMap k v -> DisjointMap k v
unionWeakly !k
x !k
y DisjointMap k v
set = forall a b c. (a -> b -> c) -> b -> a -> c
flip forall s a. State s a -> s -> s
execState DisjointMap k v
set forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. MaybeT m a -> m (Maybe a)
runMaybeT forall a b. (a -> b) -> a -> b
$ do
  Maybe k
mx <- forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) s a. Monad m => (s -> (a, s)) -> StateT s m a
state forall a b. (a -> b) -> a -> b
$ forall k v.
Ord k =>
k -> DisjointMap k v -> (Maybe k, DisjointMap k v)
lookupCompress k
x
  Maybe k
my <- forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) s a. Monad m => (s -> (a, s)) -> StateT s m a
state forall a b. (a -> b) -> a -> b
$ forall k v.
Ord k =>
k -> DisjointMap k v -> (Maybe k, DisjointMap k v)
lookupCompress k
y
  case Maybe k
mx of
    Maybe k
Nothing -> case Maybe k
my of
      Maybe k
Nothing -> forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
      Just k
repy -> do
        DisjointMap Map k k
p Map k (Ranked k v)
r <- forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall (m :: * -> *) s. Monad m => StateT s m s
get
        forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) s. Monad m => s -> StateT s m ()
put forall a b. (a -> b) -> a -> b
$
          let p' :: Map k k
p' = forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert k
x k
repy Map k k
p
              Ranked Int
ranky Set k
keys v
val = forall a. a -> Maybe a -> a
fromMaybe (forall a. HasCallStack => String -> a
error String
"Data.DisjointMap.unionWeakly") (forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup k
repy Map k (Ranked k v)
r)
              r' :: Map k (Ranked k v)
r' = forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert k
repy (forall k v. Int -> Set k -> v -> Ranked k v
Ranked Int
ranky (forall a. Ord a => a -> Set a -> Set a
S.insert k
x Set k
keys) v
val) Map k (Ranked k v)
r
           in forall k v. Map k k -> Map k (Ranked k v) -> DisjointMap k v
DisjointMap Map k k
p' Map k (Ranked k v)
r'
    Just k
repx -> case Maybe k
my of
      Maybe k
Nothing -> do
        DisjointMap Map k k
p Map k (Ranked k v)
r <- forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall (m :: * -> *) s. Monad m => StateT s m s
get
        forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) s. Monad m => s -> StateT s m ()
put forall a b. (a -> b) -> a -> b
$
          let p' :: Map k k
p' = forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert k
y k
repx Map k k
p
              Ranked Int
rankx Set k
keys v
val = forall a. a -> Maybe a -> a
fromMaybe (forall a. HasCallStack => String -> a
error String
"Data.DisjointMap.unionWeakly") (forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup k
repx Map k (Ranked k v)
r)
              r' :: Map k (Ranked k v)
r' = forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert k
repx (forall k v. Int -> Set k -> v -> Ranked k v
Ranked Int
rankx (forall a. Ord a => a -> Set a -> Set a
S.insert k
y Set k
keys) v
val) Map k (Ranked k v)
r
           in forall k v. Map k k -> Map k (Ranked k v) -> DisjointMap k v
DisjointMap Map k k
p' Map k (Ranked k v)
r'
      Just k
repy -> do
        forall (f :: * -> *). Alternative f => Bool -> f ()
guard forall a b. (a -> b) -> a -> b
$ k
repx forall a. Eq a => a -> a -> Bool
/= k
repy
        DisjointMap Map k k
p Map k (Ranked k v)
r <- forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall (m :: * -> *) s. Monad m => StateT s m s
get
        let Ranked Int
rankx Set k
keysx v
valx = Map k (Ranked k v)
r forall k a. Ord k => Map k a -> k -> a
M.! k
repx
        let Ranked Int
ranky Set k
keysy v
valy = Map k (Ranked k v)
r forall k a. Ord k => Map k a -> k -> a
M.! k
repy
        let val :: v
val = v
valx forall a. Semigroup a => a -> a -> a
<> v
valy
        let keys :: Set k
keys = forall a. Monoid a => a -> a -> a
mappend Set k
keysx Set k
keysy
        forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) s. Monad m => s -> StateT s m ()
put forall a b. (a -> b) -> a -> b
$! case forall a. Ord a => a -> a -> Ordering
compare Int
rankx Int
ranky of
          Ordering
LT -> let p' :: Map k k
p' = forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert k
repx k
repy Map k k
p
                    r' :: Map k (Ranked k v)
r' = forall k a. Ord k => k -> Map k a -> Map k a
M.delete k
repx forall a b. (a -> b) -> a -> b
$! forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert k
repy (forall k v. Int -> Set k -> v -> Ranked k v
Ranked Int
ranky Set k
keys v
val) Map k (Ranked k v)
r
                in  forall k v. Map k k -> Map k (Ranked k v) -> DisjointMap k v
DisjointMap Map k k
p' Map k (Ranked k v)
r'
          Ordering
GT -> let p' :: Map k k
p' = forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert k
repy k
repx Map k k
p
                    r' :: Map k (Ranked k v)
r' = forall k a. Ord k => k -> Map k a -> Map k a
M.delete k
repy forall a b. (a -> b) -> a -> b
$! forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert k
repx (forall k v. Int -> Set k -> v -> Ranked k v
Ranked Int
rankx Set k
keys v
val) Map k (Ranked k v)
r
                in  forall k v. Map k k -> Map k (Ranked k v) -> DisjointMap k v
DisjointMap Map k k
p' Map k (Ranked k v)
r'
          Ordering
EQ -> let p' :: Map k k
p' = forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert k
repx k
repy Map k k
p
                    r' :: Map k (Ranked k v)
r' = forall k a. Ord k => k -> Map k a -> Map k a
M.delete k
repx forall a b. (a -> b) -> a -> b
$! forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert k
repy (forall k v. Int -> Set k -> v -> Ranked k v
Ranked (Int
ranky forall a. Num a => a -> a -> a
+ Int
1) Set k
keys v
val) Map k (Ranked k v)
r
                in  forall k v. Map k k -> Map k (Ranked k v) -> DisjointMap k v
DisjointMap Map k k
p' Map k (Ranked k v)
r'

{-|
Find the set representative for this input. This function
ignores the values in the map.
-}
representative :: Ord k => k -> DisjointMap k v -> Maybe k
representative :: forall k v. Ord k => k -> DisjointMap k v -> Maybe k
representative = forall k v. Ord k => k -> DisjointMap k v -> Maybe k
find

{-| Insert a key-value pair into the disjoint map. If the key
    is is already present in another set, combine the value
    monoidally with the value belonging to it. The new value
    is on the left side of the append, and the old value is
    on the right.
    Otherwise, this creates a singleton set as a new key and
    associates it with the value.
-}
insert :: (Ord k, Semigroup v) => k -> v -> DisjointMap k v -> DisjointMap k v
insert :: forall k v.
(Ord k, Semigroup v) =>
k -> v -> DisjointMap k v -> DisjointMap k v
insert !k
x = forall k v.
(Ord k, Semigroup v) =>
k -> Set k -> v -> DisjointMap k v -> DisjointMap k v
insertInternal k
x (forall a. a -> Set a
S.singleton k
x)

-- Precondition: Nothing in ks already exists in the disjoint map.
-- This function should only be used by insert.
insertInternal :: (Ord k, Semigroup v) => k -> Set k -> v -> DisjointMap k v -> DisjointMap k v
insertInternal :: forall k v.
(Ord k, Semigroup v) =>
k -> Set k -> v -> DisjointMap k v -> DisjointMap k v
insertInternal !k
x !Set k
ks !v
v set :: DisjointMap k v
set@(DisjointMap Map k k
p Map k (Ranked k v)
r) =
  let (Maybe k
l, Map k k
p') = forall k a.
Ord k =>
(k -> a -> a -> a) -> k -> a -> Map k a -> (Maybe a, Map k a)
M.insertLookupWithKey (\k
_ k
_ k
old -> k
old) k
x k
x Map k k
p
   in case Maybe k
l of
        Just k
_ ->
          let (Maybe k
m,DisjointMap Map k k
p2 Map k (Ranked k v)
r') = forall k v.
Ord k =>
k -> DisjointMap k v -> (Maybe k, DisjointMap k v)
representative' k
x DisjointMap k v
set
           in case Maybe k
m of
                Maybe k
Nothing -> forall a. HasCallStack => String -> a
error String
"DisjointMap insert: invariant violated"
                Just k
root -> forall k v. Map k k -> Map k (Ranked k v) -> DisjointMap k v
DisjointMap Map k k
p2 (forall k a. Ord k => (a -> a) -> k -> Map k a -> Map k a
M.adjust (\(Ranked Int
rank Set k
oldKs v
vOld) -> forall k v. Int -> Set k -> v -> Ranked k v
Ranked Int
rank (forall a. Monoid a => a -> a -> a
mappend Set k
oldKs Set k
ks) (v
v forall a. Semigroup a => a -> a -> a
<> v
vOld)) k
root Map k (Ranked k v)
r')
        Maybe k
Nothing ->
          let r' :: Map k (Ranked k v)
r' = forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert k
x (forall k v. Int -> Set k -> v -> Ranked k v
Ranked Int
0 Set k
ks v
v) Map k (Ranked k v)
r
          in  forall k v. Map k k -> Map k (Ranked k v) -> DisjointMap k v
DisjointMap Map k k
p' Map k (Ranked k v)
r'


{-| Create a disjoint map with one key: a singleton set. O(1). -}
singleton :: k -> v -> DisjointMap k v
singleton :: forall k v. k -> v -> DisjointMap k v
singleton !k
x !v
v =
  let p :: Map k k
p = forall k a. k -> a -> Map k a
M.singleton k
x k
x
      r :: Map k (Ranked k v)
r = forall k a. k -> a -> Map k a
M.singleton k
x (forall k v. Int -> Set k -> v -> Ranked k v
Ranked Int
0 (forall a. a -> Set a
S.singleton k
x) v
v)
   in forall k v. Map k k -> Map k (Ranked k v) -> DisjointMap k v
DisjointMap Map k k
p Map k (Ranked k v)
r

{-| The empty map -}
empty :: DisjointMap k v
empty :: forall k v. DisjointMap k v
empty = forall k v. Map k k -> Map k (Ranked k v) -> DisjointMap k v
DisjointMap forall k a. Map k a
M.empty forall k a. Map k a
M.empty

append :: (Ord k, Semigroup v) => DisjointMap k v -> DisjointMap k v -> DisjointMap k v
append :: forall k v.
(Ord k, Semigroup v) =>
DisjointMap k v -> DisjointMap k v -> DisjointMap k v
append s1 :: DisjointMap k v
s1@(DisjointMap Map k k
m1 Map k (Ranked k v)
r1) s2 :: DisjointMap k v
s2@(DisjointMap Map k k
m2 Map k (Ranked k v)
r2) = if forall k a. Map k a -> Int
M.size Map k k
m1 forall a. Ord a => a -> a -> Bool
> forall k a. Map k a -> Int
M.size Map k k
m2
  then forall k v.
(Ord k, Semigroup v) =>
DisjointMap k v -> Map k k -> DisjointMap k v
appendPhase2 (forall k v.
(Ord k, Semigroup v) =>
Map k (Ranked k v) -> DisjointMap k v -> Map k k -> DisjointMap k v
appendPhase1 Map k (Ranked k v)
r2 DisjointMap k v
s1 Map k k
m2) Map k k
m2
  else forall k v.
(Ord k, Semigroup v) =>
DisjointMap k v -> Map k k -> DisjointMap k v
appendPhase2 (forall k v.
(Ord k, Semigroup v) =>
Map k (Ranked k v) -> DisjointMap k v -> Map k k -> DisjointMap k v
appendPhase1 Map k (Ranked k v)
r1 DisjointMap k v
s2 Map k k
m1) Map k k
m1

appendPhase1 :: (Ord k, Semigroup v)
  => Map k (Ranked k v)
  -> DisjointMap k v
  -> Map k k
  -> DisjointMap k v
appendPhase1 :: forall k v.
(Ord k, Semigroup v) =>
Map k (Ranked k v) -> DisjointMap k v -> Map k k -> DisjointMap k v
appendPhase1 !Map k (Ranked k v)
ranks = forall a k b. (a -> k -> b -> a) -> a -> Map k b -> a
M.foldlWithKey' forall a b. (a -> b) -> a -> b
$ \DisjointMap k v
ds k
x k
y -> if k
x forall a. Eq a => a -> a -> Bool
== k
y
  then case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup k
x Map k (Ranked k v)
ranks of
    Maybe (Ranked k v)
Nothing -> forall a. HasCallStack => String -> a
error String
"Data.DisjointMap.appendParents: invariant violated"
    Just (Ranked Int
_ Set k
ks v
v) -> forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
F.foldl' (\DisjointMap k v
dm k
k -> forall k v.
(Ord k, Semigroup v) =>
k -> k -> DisjointMap k v -> DisjointMap k v
unionWeakly k
k k
x DisjointMap k v
dm) (forall k v.
(Ord k, Semigroup v) =>
k -> v -> DisjointMap k v -> DisjointMap k v
insert k
x v
v DisjointMap k v
ds) Set k
ks
  else DisjointMap k v
ds

appendPhase2 :: (Ord k, Semigroup v) => DisjointMap k v -> Map k k -> DisjointMap k v
appendPhase2 :: forall k v.
(Ord k, Semigroup v) =>
DisjointMap k v -> Map k k -> DisjointMap k v
appendPhase2 = forall a k b. (a -> k -> b -> a) -> a -> Map k b -> a
M.foldlWithKey' forall a b. (a -> b) -> a -> b
$ \DisjointMap k v
ds k
x k
y -> if k
x forall a. Eq a => a -> a -> Bool
== k
y
  then DisjointMap k v
ds
  else forall k v.
(Ord k, Semigroup v) =>
k -> k -> DisjointMap k v -> DisjointMap k v
unionWeakly k
x k
y DisjointMap k v
ds

{-| Create a disjoint map with one key. Everything in the
    'Set' argument is consider part of the same equivalence
    class.
-}
singletons :: Eq k => Set k -> v -> DisjointMap k v
singletons :: forall k v. Eq k => Set k -> v -> DisjointMap k v
singletons Set k
s v
v = case forall a. Set a -> Maybe a
setLookupMin Set k
s of
  Maybe k
Nothing -> forall k v. DisjointMap k v
empty
  Just k
x ->
    let p :: Map k k
p = forall k a. (k -> a) -> Set k -> Map k a
M.fromSet (\k
_ -> k
x) Set k
s
        rank :: Int
rank = if forall a. Set a -> Int
S.size Set k
s forall a. Eq a => a -> a -> Bool
== Int
1 then Int
0 else Int
1
        r :: Map k (Ranked k v)
r = forall k a. k -> a -> Map k a
M.singleton k
x (forall k v. Int -> Set k -> v -> Ranked k v
Ranked Int
rank Set k
s v
v)
    in forall k v. Map k k -> Map k (Ranked k v) -> DisjointMap k v
DisjointMap Map k k
p Map k (Ranked k v)
r

setLookupMin :: Set a -> Maybe a
#if MIN_VERSION_containers(0,5,9) 
setLookupMin :: forall a. Set a -> Maybe a
setLookupMin = forall a. Set a -> Maybe a
S.lookupMin
#else
setLookupMin s = if S.size s > 0 then Just (S.findMin s) else Nothing
#endif

{-|
Find the set representative for this input. Returns a new disjoint
set in which the path has been compressed.
-}
representative' :: Ord k => k -> DisjointMap k v -> (Maybe k, DisjointMap k v)
representative' :: forall k v.
Ord k =>
k -> DisjointMap k v -> (Maybe k, DisjointMap k v)
representative' !k
x DisjointMap k v
set =
  case forall k v. Ord k => k -> DisjointMap k v -> Maybe k
find k
x DisjointMap k v
set of
    Maybe k
Nothing  -> (forall a. Maybe a
Nothing, DisjointMap k v
set)
    Just k
rep -> let set' :: DisjointMap k v
set' = forall k v. Ord k => k -> k -> DisjointMap k v -> DisjointMap k v
compress k
rep k
x DisjointMap k v
set
                in  DisjointMap k v
set' seq :: forall a b. a -> b -> b
`seq` (forall a. a -> Maybe a
Just k
rep, DisjointMap k v
set')

lookupCompressAdd :: (Ord k, Monoid v) => k -> DisjointMap k v -> (k, DisjointMap k v)
lookupCompressAdd :: forall k v.
(Ord k, Monoid v) =>
k -> DisjointMap k v -> (k, DisjointMap k v)
lookupCompressAdd !k
x DisjointMap k v
set =
  case forall k v. Ord k => k -> DisjointMap k v -> Maybe k
find k
x DisjointMap k v
set of
    Maybe k
Nothing -> (k
x, forall k v.
(Ord k, Semigroup v) =>
k -> v -> DisjointMap k v -> DisjointMap k v
insert k
x forall a. Monoid a => a
mempty DisjointMap k v
set)
    Just k
rep -> let !set' :: DisjointMap k v
set' = forall k v. Ord k => k -> k -> DisjointMap k v -> DisjointMap k v
compress k
rep k
x DisjointMap k v
set
                 in (k
rep, DisjointMap k v
set')

lookupCompress :: Ord k => k -> DisjointMap k v -> (Maybe k, DisjointMap k v)
lookupCompress :: forall k v.
Ord k =>
k -> DisjointMap k v -> (Maybe k, DisjointMap k v)
lookupCompress !k
x DisjointMap k v
set =
  case forall k v. Ord k => k -> DisjointMap k v -> Maybe k
find k
x DisjointMap k v
set of
    Maybe k
Nothing -> (forall a. Maybe a
Nothing, DisjointMap k v
set)
    Just k
rep ->
      let !set' :: DisjointMap k v
set' = forall k v. Ord k => k -> k -> DisjointMap k v -> DisjointMap k v
compress k
rep k
x DisjointMap k v
set
       in (forall a. a -> Maybe a
Just k
rep, DisjointMap k v
set')

find :: Ord k => k -> DisjointMap k v -> Maybe k
find :: forall k v. Ord k => k -> DisjointMap k v -> Maybe k
find !k
x (DisjointMap Map k k
p Map k (Ranked k v)
_) = do
  k
x' <- forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup k
x Map k k
p
  forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! if k
x forall a. Eq a => a -> a -> Bool
== k
x' then k
x' else k -> k
find' k
x'
  where
  find' :: k -> k
find' k
y =
    let y' :: k
y' = Map k k
p forall k a. Ord k => Map k a -> k -> a
M.! k
y
     in if k
y forall a. Eq a => a -> a -> Bool
== k
y' then k
y' else k -> k
find' k
y'

{-| Find the value associated with the set containing
    the provided key. If the key is not found, use 'mempty'.
-}
lookup :: (Ord k, Monoid v) => k -> DisjointMap k v -> v
lookup :: forall k v. (Ord k, Monoid v) => k -> DisjointMap k v -> v
lookup k
k = forall a. a -> Maybe a -> a
fromMaybe forall a. Monoid a => a
mempty forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k v. Ord k => k -> DisjointMap k v -> Maybe v
lookup' k
k

{-| Find the value associated with the set containing
    the provided key.
-}
lookup' :: Ord k => k -> DisjointMap k v -> Maybe v
lookup' :: forall k v. Ord k => k -> DisjointMap k v -> Maybe v
lookup' !k
x (DisjointMap Map k k
p Map k (Ranked k v)
r) = do
  k
x' <- forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup k
x Map k k
p
  if k
x forall a. Eq a => a -> a -> Bool
== k
x'
    then case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup k
x Map k (Ranked k v)
r of
      Maybe (Ranked k v)
Nothing -> forall a. Maybe a
Nothing
      Just (Ranked Int
_ Set k
_ v
v) -> forall a. a -> Maybe a
Just v
v
    else k -> Maybe v
find' k
x'
  where
  find' :: k -> Maybe v
find' k
y =
    let y' :: k
y' = Map k k
p forall k a. Ord k => Map k a -> k -> a
M.! k
y
     in if k
y forall a. Eq a => a -> a -> Bool
== k
y'
          then case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup k
y Map k (Ranked k v)
r of
            Maybe (Ranked k v)
Nothing -> forall a. Maybe a
Nothing
            Just (Ranked Int
_ Set k
_ v
v) -> forall a. a -> Maybe a
Just v
v
          else k -> Maybe v
find' k
y'

-- TODO: make this smarter about recreating the parents Map.
-- Currently, it will recreate it more often than needed.
compress :: Ord k => k -> k -> DisjointMap k v -> DisjointMap k v
compress :: forall k v. Ord k => k -> k -> DisjointMap k v -> DisjointMap k v
compress !k
rep = forall {v}. k -> DisjointMap k v -> DisjointMap k v
helper
  where
  helper :: k -> DisjointMap k v -> DisjointMap k v
helper !k
x set :: DisjointMap k v
set@(DisjointMap Map k k
p Map k (Ranked k v)
r)
    | k
x forall a. Eq a => a -> a -> Bool
== k
rep  = DisjointMap k v
set
    | Bool
otherwise = k -> DisjointMap k v -> DisjointMap k v
helper k
x' DisjointMap k v
set'
    where
    x' :: k
x' = Map k k
p forall k a. Ord k => Map k a -> k -> a
M.! k
x
    set' :: DisjointMap k v
set' = let !p' :: Map k k
p' = forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert k
x k
rep Map k k
p
            in forall k v. Map k k -> Map k (Ranked k v) -> DisjointMap k v
DisjointMap Map k k
p' Map k (Ranked k v)
r

{- $tutorial

The disjoint map data structure can be used to model scenarios where
the key of a map is a disjoint set. Let us consider trying to find
the lowest rating of our by town. Due to differing subcultures,
inhabitants do not necessarily agree on a canonical name for each town.
Consequently, the survey allows participants to write in their town
name as they would call it. We begin with a rating data type:

>>> import Data.Function ((&))
>>> data Rating = Lowest | Low | Medium | High | Highest deriving (Eq,Ord,Show)
>>> instance Semigroup Rating where (<>) = min
>>> instance Monoid Rating where mempty = Highest; mappend = min

Notice that the 'Monoid' instance combines ratings by choosing
the lower one. Now, we consider the data from the survey:

>>> let resA = [("Big Lake",High),("Newport Lake",Medium),("Dustboro",Low)]
>>> let resB = [("Sand Town",Medium),("Sand Town",High),("Mount Lucy",High)]
>>> let resC = [("Lucy Hill",Highest),("Dustboro",High),("Dustville",High)]
>>> let m1 = foldMap (uncurry singleton) (resA ++ resB ++ resC)
>>> :t m1
m1 :: DisjointMap [Char] Rating
>>> mapM_ putStrLn (prettyList m1)
{"Big Lake"} -> High
{"Dustboro"} -> Low
{"Dustville"} -> High
{"Lucy Hill"} -> Highest
{"Mount Lucy"} -> High
{"Newport Lake"} -> Medium
{"Sand Town"} -> Medium

Since we haven't defined any equivalence classes for the town names yet,
what we have so far is no different from an ordinary 'Map'. Now we
will introduce some equivalences:

>>> let m2 = m1 & union "Big Lake" "Newport Lake" & union "Sand Town" "Dustboro"
>>> let m3 = m2 & union "Dustboro" "Dustville" & union "Lucy Hill" "Mount Lucy"
>>> mapM_ putStrLn (prettyList m3)
{"Dustboro","Dustville","Sand Town"} -> Low
{"Lucy Hill","Mount Lucy"} -> High
{"Big Lake","Newport Lake"} -> Medium

We can now query the map to find the lowest rating in a given town:

>>> lookup "Dustville" m3
Low
>>> lookup "Lucy Hill" m3
High

-}