{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE DeriveTraversable #-}
{-# 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 Control.Monad
import Control.Monad.Trans.Class
import Control.Monad.Trans.Maybe
import Control.Monad.Trans.State.Strict
import Prelude hiding (lookup)

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

{- | 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 -> b) -> DisjointMap k a -> DisjointMap k b)
-> (forall a b. a -> DisjointMap k b -> DisjointMap k a)
-> Functor (DisjointMap k)
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
$cfmap :: forall k a b. (a -> b) -> DisjointMap k a -> DisjointMap k b
fmap :: forall a b. (a -> b) -> DisjointMap k a -> DisjointMap k b
$c<$ :: forall k a b. a -> DisjointMap k b -> DisjointMap k a
<$ :: forall a b. a -> DisjointMap k b -> DisjointMap k a
Functor, (forall m. Monoid m => DisjointMap k m -> m)
-> (forall m a. Monoid m => (a -> m) -> DisjointMap k a -> m)
-> (forall m a. Monoid m => (a -> m) -> DisjointMap k a -> m)
-> (forall a b. (a -> b -> b) -> b -> DisjointMap k a -> b)
-> (forall a b. (a -> b -> b) -> b -> DisjointMap k a -> b)
-> (forall b a. (b -> a -> b) -> b -> DisjointMap k a -> b)
-> (forall b a. (b -> a -> b) -> b -> DisjointMap k a -> b)
-> (forall a. (a -> a -> a) -> DisjointMap k a -> a)
-> (forall a. (a -> a -> a) -> DisjointMap k a -> a)
-> (forall a. DisjointMap k a -> [a])
-> (forall a. DisjointMap k a -> Bool)
-> (forall a. DisjointMap k a -> Int)
-> (forall a. Eq a => a -> DisjointMap k a -> Bool)
-> (forall a. Ord a => DisjointMap k a -> a)
-> (forall a. Ord a => DisjointMap k a -> a)
-> (forall a. Num a => DisjointMap k a -> a)
-> (forall a. Num a => DisjointMap k a -> a)
-> Foldable (DisjointMap k)
forall a. Eq a => a -> DisjointMap k a -> Bool
forall a. Num a => DisjointMap k a -> a
forall a. Ord a => DisjointMap k a -> a
forall m. Monoid m => DisjointMap k m -> m
forall a. DisjointMap k a -> Bool
forall a. DisjointMap k a -> Int
forall a. DisjointMap k a -> [a]
forall a. (a -> a -> a) -> DisjointMap k a -> a
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 b a. (b -> a -> b) -> b -> DisjointMap k a -> b
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
$cfold :: forall k m. Monoid m => DisjointMap k m -> m
fold :: forall m. Monoid m => DisjointMap k m -> 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
foldMap' :: forall m a. Monoid m => (a -> m) -> DisjointMap k a -> m
$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
foldr' :: forall a b. (a -> b -> 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
foldl' :: forall b a. (b -> a -> b) -> b -> DisjointMap k a -> b
$cfoldr1 :: forall k a. (a -> a -> a) -> DisjointMap k a -> a
foldr1 :: forall a. (a -> a -> a) -> DisjointMap k a -> a
$cfoldl1 :: forall k a. (a -> a -> a) -> DisjointMap k a -> a
foldl1 :: forall a. (a -> a -> a) -> DisjointMap k a -> a
$ctoList :: forall k a. DisjointMap k a -> [a]
toList :: forall a. DisjointMap k a -> [a]
$cnull :: forall k a. DisjointMap k a -> Bool
null :: forall a. DisjointMap k a -> Bool
$clength :: forall k a. DisjointMap k a -> Int
length :: forall a. DisjointMap k a -> Int
$celem :: forall k a. Eq a => a -> DisjointMap k a -> Bool
elem :: forall a. Eq a => a -> DisjointMap k a -> Bool
$cmaximum :: forall k a. Ord a => DisjointMap k a -> a
maximum :: forall a. Ord a => DisjointMap k a -> a
$cminimum :: forall k a. Ord a => DisjointMap k a -> a
minimum :: forall a. Ord a => DisjointMap k a -> a
$csum :: forall k a. Num a => DisjointMap k a -> a
sum :: forall a. Num a => DisjointMap k a -> a
$cproduct :: forall k a. Num a => DisjointMap k a -> a
product :: forall a. Num a => DisjointMap k a -> a
Foldable, Functor (DisjointMap k)
Foldable (DisjointMap k)
(Functor (DisjointMap k), Foldable (DisjointMap k)) =>
(forall (f :: * -> *) a b.
 Applicative f =>
 (a -> f b) -> DisjointMap k a -> f (DisjointMap k b))
-> (forall (f :: * -> *) a.
    Applicative f =>
    DisjointMap k (f a) -> f (DisjointMap k a))
-> (forall (m :: * -> *) a b.
    Monad m =>
    (a -> m b) -> DisjointMap k a -> m (DisjointMap k b))
-> (forall (m :: * -> *) a.
    Monad m =>
    DisjointMap k (m a) -> m (DisjointMap k a))
-> Traversable (DisjointMap k)
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 (m :: * -> *) a.
Monad m =>
DisjointMap k (m a) -> m (DisjointMap k a)
forall (f :: * -> *) a.
Applicative f =>
DisjointMap k (f a) -> f (DisjointMap k a)
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> DisjointMap k a -> m (DisjointMap k b)
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)
traverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> DisjointMap k a -> f (DisjointMap k b)
$csequenceA :: forall k (f :: * -> *) a.
Applicative f =>
DisjointMap k (f a) -> f (DisjointMap k a)
sequenceA :: forall (f :: * -> *) a.
Applicative f =>
DisjointMap k (f a) -> f (DisjointMap k a)
$cmapM :: forall k (m :: * -> *) a b.
Monad m =>
(a -> m b) -> DisjointMap k a -> m (DisjointMap k b)
mapM :: forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> DisjointMap k a -> m (DisjointMap k b)
$csequence :: forall k (m :: * -> *) a.
Monad m =>
DisjointMap k (m a) -> m (DisjointMap k a)
sequence :: forall (m :: * -> *) a.
Monad m =>
DisjointMap k (m a) -> m (DisjointMap k a)
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 -> b) -> Ranked k a -> Ranked k b)
-> (forall a b. a -> Ranked k b -> Ranked k a)
-> Functor (Ranked k)
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
$cfmap :: forall k a b. (a -> b) -> Ranked k a -> Ranked k b
fmap :: forall a b. (a -> b) -> Ranked k a -> Ranked k b
$c<$ :: forall k a b. a -> Ranked k b -> Ranked k a
<$ :: forall a b. a -> Ranked k b -> Ranked k a
Functor, (forall m. Monoid m => Ranked k m -> m)
-> (forall m a. Monoid m => (a -> m) -> Ranked k a -> m)
-> (forall m a. Monoid m => (a -> m) -> Ranked k a -> m)
-> (forall a b. (a -> b -> b) -> b -> Ranked k a -> b)
-> (forall a b. (a -> b -> b) -> b -> Ranked k a -> b)
-> (forall b a. (b -> a -> b) -> b -> Ranked k a -> b)
-> (forall b a. (b -> a -> b) -> b -> Ranked k a -> b)
-> (forall a. (a -> a -> a) -> Ranked k a -> a)
-> (forall a. (a -> a -> a) -> Ranked k a -> a)
-> (forall a. Ranked k a -> [a])
-> (forall a. Ranked k a -> Bool)
-> (forall a. Ranked k a -> Int)
-> (forall a. Eq a => a -> Ranked k a -> Bool)
-> (forall a. Ord a => Ranked k a -> a)
-> (forall a. Ord a => Ranked k a -> a)
-> (forall a. Num a => Ranked k a -> a)
-> (forall a. Num a => Ranked k a -> a)
-> Foldable (Ranked k)
forall a. Eq a => a -> Ranked k a -> Bool
forall a. Num a => Ranked k a -> a
forall a. Ord a => Ranked k a -> a
forall m. Monoid m => Ranked k m -> m
forall a. Ranked k a -> Bool
forall a. Ranked k a -> Int
forall a. Ranked k a -> [a]
forall a. (a -> a -> a) -> Ranked k a -> a
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 b a. (b -> a -> b) -> b -> Ranked k a -> b
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
$cfold :: forall k m. Monoid m => Ranked k m -> m
fold :: forall m. Monoid m => Ranked k m -> 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
foldMap' :: forall m a. Monoid m => (a -> m) -> Ranked k a -> m
$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
foldr' :: forall a b. (a -> b -> 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
foldl' :: forall b a. (b -> a -> b) -> b -> Ranked k a -> b
$cfoldr1 :: forall k a. (a -> a -> a) -> Ranked k a -> a
foldr1 :: forall a. (a -> a -> a) -> Ranked k a -> a
$cfoldl1 :: forall k a. (a -> a -> a) -> Ranked k a -> a
foldl1 :: forall a. (a -> a -> a) -> Ranked k a -> a
$ctoList :: forall k a. Ranked k a -> [a]
toList :: forall a. Ranked k a -> [a]
$cnull :: forall k a. Ranked k a -> Bool
null :: forall a. Ranked k a -> Bool
$clength :: forall k a. Ranked k a -> Int
length :: forall a. Ranked k a -> Int
$celem :: forall k a. Eq a => a -> Ranked k a -> Bool
elem :: forall a. Eq a => a -> Ranked k a -> Bool
$cmaximum :: forall k a. Ord a => Ranked k a -> a
maximum :: forall a. Ord a => Ranked k a -> a
$cminimum :: forall k a. Ord a => Ranked k a -> a
minimum :: forall a. Ord a => Ranked k a -> a
$csum :: forall k a. Num a => Ranked k a -> a
sum :: forall a. Num a => Ranked k a -> a
$cproduct :: forall k a. Num a => Ranked k a -> a
product :: forall a. Num a => Ranked k a -> a
Foldable, Functor (Ranked k)
Foldable (Ranked k)
(Functor (Ranked k), Foldable (Ranked k)) =>
(forall (f :: * -> *) a b.
 Applicative f =>
 (a -> f b) -> Ranked k a -> f (Ranked k b))
-> (forall (f :: * -> *) a.
    Applicative f =>
    Ranked k (f a) -> f (Ranked k a))
-> (forall (m :: * -> *) a b.
    Monad m =>
    (a -> m b) -> Ranked k a -> m (Ranked k b))
-> (forall (m :: * -> *) a.
    Monad m =>
    Ranked k (m a) -> m (Ranked k a))
-> Traversable (Ranked k)
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 (m :: * -> *) a. Monad m => Ranked k (m a) -> m (Ranked k a)
forall (f :: * -> *) a.
Applicative f =>
Ranked k (f a) -> f (Ranked k a)
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Ranked k a -> m (Ranked k b)
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)
traverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Ranked k a -> f (Ranked k b)
$csequenceA :: forall k (f :: * -> *) a.
Applicative f =>
Ranked k (f a) -> f (Ranked k a)
sequenceA :: forall (f :: * -> *) a.
Applicative f =>
Ranked k (f a) -> f (Ranked k a)
$cmapM :: forall k (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Ranked k a -> m (Ranked k b)
mapM :: forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Ranked k a -> m (Ranked k b)
$csequence :: forall k (m :: * -> *) a.
Monad m =>
Ranked k (m a) -> m (Ranked k a)
sequence :: forall (m :: * -> *) a. Monad m => Ranked k (m a) -> m (Ranked k a)
Traversable)

instance (Ord k, Monoid v) => Monoid (DisjointMap k v) where
  mempty :: DisjointMap k v
mempty = DisjointMap k v
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
(<>) = 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 = [(Set k, v)] -> Set (Set k, v)
forall a. Ord a => [a] -> Set a
S.fromList (DisjointMap k v -> [(Set k, v)]
forall k v. DisjointMap k v -> [(Set k, v)]
toSets DisjointMap k v
a) Set (Set k, v) -> Set (Set k, v) -> Bool
forall a. Eq a => a -> a -> Bool
== [(Set k, v)] -> Set (Set k, v)
forall a. Ord a => [a] -> Set a
S.fromList (DisjointMap k v -> [(Set k, v)]
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 = Set (Set k, v) -> Set (Set k, v) -> Ordering
forall a. Ord a => a -> a -> Ordering
compare ([(Set k, v)] -> Set (Set k, v)
forall a. Ord a => [a] -> Set a
S.fromList (DisjointMap k v -> [(Set k, v)]
forall k v. DisjointMap k v -> [(Set k, v)]
toSets DisjointMap k v
a)) ([(Set k, v)] -> Set (Set k, v)
forall a. Ord a => [a] -> Set a
S.fromList (DisjointMap k v -> [(Set k, v)]
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 = DisjointMap k v -> String
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 [Set k] -> Maybe (Set k)
forall a. Ord a => [Set a] -> Maybe (Set a)
unionDistinctAll (((Set k, v) -> Set k) -> [(Set k, v)] -> [Set k]
forall a b. (a -> b) -> [a] -> [b]
map (Set k, v) -> Set k
forall a b. (a, b) -> a
fst [(Set k, v)]
xs) of
  Maybe (Set k)
Nothing -> Maybe (DisjointMap k v)
forall a. Maybe a
Nothing
  Just Set k
_ -> DisjointMap k v -> Maybe (DisjointMap k v)
forall a. a -> Maybe a
Just ([(Set k, v)] -> DisjointMap k v -> DisjointMap k v
forall k v.
Ord k =>
[(Set k, v)] -> DisjointMap k v -> DisjointMap k v
unsafeFromSets [(Set k, v)]
xs DisjointMap k v
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 Set k -> Maybe k
forall a. Set a -> Maybe a
setLookupMin Set k
x of
    Maybe k
Nothing -> [(Set k, v)] -> DisjointMap k v -> DisjointMap k v
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 ->
      [(Set k, v)] -> DisjointMap k v -> DisjointMap k v
forall k v.
Ord k =>
[(Set k, v)] -> DisjointMap k v -> DisjointMap k v
unsafeFromSets [(Set k, v)]
xs (DisjointMap k v -> DisjointMap k v)
-> DisjointMap k v -> DisjointMap k v
forall a b. (a -> b) -> a -> b
$
        Map k k -> Map k (Ranked k v) -> DisjointMap k v
forall k v. Map k k -> Map k (Ranked k v) -> DisjointMap k v
DisjointMap
          (Map k k -> Map k k -> Map k k
forall k a. Ord k => Map k a -> Map k a -> Map k a
M.union ((k -> k) -> Set k -> Map k k
forall k a. (k -> a) -> Set k -> Map k a
M.fromSet (\k
_ -> k
m) Set k
x) Map k k
p)
          (k -> Ranked k v -> Map k (Ranked k v) -> Map k (Ranked k v)
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert k
m (Int -> Set k -> v -> Ranked k v
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 = (Set a -> Set a -> Maybe (Set a))
-> Set a -> [Set a] -> Maybe (Set a)
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldlM Set a -> Set a -> Maybe (Set a)
forall a. Ord a => Set a -> Set a -> Maybe (Set a)
unionDistinct Set a
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 = Set a -> Set a -> Set a
forall a. Ord a => Set a -> Set a -> Set a
S.union Set a
a Set a
b
   in if Set a -> Int
forall a. Set a -> Int
S.size Set a
a Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Set a -> Int
forall a. Set a -> Int
S.size Set a
b Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Set a -> Int
forall a. Set a -> Int
S.size Set a
s
        then Set a -> Maybe (Set a)
forall a. a -> Maybe a
Just Set a
s
        else Maybe (Set a)
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 = [([k], v)] -> String
forall a. Show a => a -> String
show ([([k], v)] -> String)
-> (DisjointMap k v -> [([k], v)]) -> DisjointMap k v -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DisjointMap k v -> [([k], v)]
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 = (((Set k, v) -> ([k], v)) -> [(Set k, v)] -> [([k], v)]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (((Set k, v) -> ([k], v)) -> [(Set k, v)] -> [([k], v)])
-> ((Set k -> [k]) -> (Set k, v) -> ([k], v))
-> (Set k -> [k])
-> [(Set k, v)]
-> [([k], v)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Set k -> [k]) -> (Set k, v) -> ([k], v)
forall a b c. (a -> b) -> (a, c) -> (b, c)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first) Set k -> [k]
forall a. Set a -> [a]
S.toList ([(Set k, v)] -> [([k], v)])
-> (DisjointMap k v -> [(Set k, v)])
-> DisjointMap k v
-> [([k], v)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DisjointMap k v -> [(Set k, v)]
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) =
  (Ranked k v -> [(Set k, v)] -> [(Set k, v)])
-> [(Set k, v)] -> Map k (Ranked k v) -> [(Set k, v)]
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) (Set k, v) -> [(Set k, v)] -> [(Set k, 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
"{" String -> ShowS
forall a. [a] -> [a] -> [a]
++ String -> [String] -> String
forall a. [a] -> [[a]] -> [a]
L.intercalate String
", " (DisjointMap k v -> [String]
forall k v. (Show k, Show v) => DisjointMap k v -> [String]
prettyList DisjointMap k v
dm) String -> ShowS
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 = ((Set k, v) -> String) -> [(Set k, v)] -> [String]
forall a b. (a -> b) -> [a] -> [b]
L.map (\(Set k
ks, v
v) -> String
"{" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Set k -> String
forall k. Show k => Set k -> String
commafied Set k
ks String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"} -> " String -> ShowS
forall a. [a] -> [a] -> [a]
++ v -> String
forall a. Show a => a -> String
show v
v) (DisjointMap k v -> [(Set k, 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 = [String] -> String
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join ([String] -> String) -> (Set k -> [String]) -> Set k -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> [String] -> [String]
forall a. a -> [a] -> [a]
L.intersperse String
"," ([String] -> [String]) -> (Set k -> [String]) -> Set k -> [String]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (k -> String) -> [k] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map k -> String
forall a. Show a => a -> String
show ([k] -> [String]) -> (Set k -> [k]) -> Set k -> [String]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Set k -> [k]
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) =
  (a -> Ranked k v -> a) -> a -> Map k (Ranked k v) -> a
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 = (State (DisjointMap k v) (Maybe ())
 -> DisjointMap k v -> DisjointMap k v)
-> DisjointMap k v
-> State (DisjointMap k v) (Maybe ())
-> DisjointMap k v
forall a b c. (a -> b -> c) -> b -> a -> c
flip State (DisjointMap k v) (Maybe ())
-> DisjointMap k v -> DisjointMap k v
forall s a. State s a -> s -> s
execState DisjointMap k v
set (State (DisjointMap k v) (Maybe ()) -> DisjointMap k v)
-> State (DisjointMap k v) (Maybe ()) -> DisjointMap k v
forall a b. (a -> b) -> a -> b
$ MaybeT (StateT (DisjointMap k v) Identity) ()
-> State (DisjointMap k v) (Maybe ())
forall (m :: * -> *) a. MaybeT m a -> m (Maybe a)
runMaybeT (MaybeT (StateT (DisjointMap k v) Identity) ()
 -> State (DisjointMap k v) (Maybe ()))
-> MaybeT (StateT (DisjointMap k v) Identity) ()
-> State (DisjointMap k v) (Maybe ())
forall a b. (a -> b) -> a -> b
$ do
  k
repx <- StateT (DisjointMap k v) Identity k
-> MaybeT (StateT (DisjointMap k v) Identity) k
forall (m :: * -> *) a. Monad m => m a -> MaybeT m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (StateT (DisjointMap k v) Identity k
 -> MaybeT (StateT (DisjointMap k v) Identity) k)
-> StateT (DisjointMap k v) Identity k
-> MaybeT (StateT (DisjointMap k v) Identity) k
forall a b. (a -> b) -> a -> b
$ (DisjointMap k v -> (k, DisjointMap k v))
-> StateT (DisjointMap k v) Identity k
forall (m :: * -> *) s a. Monad m => (s -> (a, s)) -> StateT s m a
state ((DisjointMap k v -> (k, DisjointMap k v))
 -> StateT (DisjointMap k v) Identity k)
-> (DisjointMap k v -> (k, DisjointMap k v))
-> StateT (DisjointMap k v) Identity k
forall a b. (a -> b) -> a -> b
$ k -> DisjointMap k v -> (k, DisjointMap k v)
forall k v.
(Ord k, Monoid v) =>
k -> DisjointMap k v -> (k, DisjointMap k v)
lookupCompressAdd k
x
  k
repy <- StateT (DisjointMap k v) Identity k
-> MaybeT (StateT (DisjointMap k v) Identity) k
forall (m :: * -> *) a. Monad m => m a -> MaybeT m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (StateT (DisjointMap k v) Identity k
 -> MaybeT (StateT (DisjointMap k v) Identity) k)
-> StateT (DisjointMap k v) Identity k
-> MaybeT (StateT (DisjointMap k v) Identity) k
forall a b. (a -> b) -> a -> b
$ (DisjointMap k v -> (k, DisjointMap k v))
-> StateT (DisjointMap k v) Identity k
forall (m :: * -> *) s a. Monad m => (s -> (a, s)) -> StateT s m a
state ((DisjointMap k v -> (k, DisjointMap k v))
 -> StateT (DisjointMap k v) Identity k)
-> (DisjointMap k v -> (k, DisjointMap k v))
-> StateT (DisjointMap k v) Identity k
forall a b. (a -> b) -> a -> b
$ k -> DisjointMap k v -> (k, DisjointMap k v)
forall k v.
(Ord k, Monoid v) =>
k -> DisjointMap k v -> (k, DisjointMap k v)
lookupCompressAdd k
y
  Bool -> MaybeT (StateT (DisjointMap k v) Identity) ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> MaybeT (StateT (DisjointMap k v) Identity) ())
-> Bool -> MaybeT (StateT (DisjointMap k v) Identity) ()
forall a b. (a -> b) -> a -> b
$ k
repx k -> k -> Bool
forall a. Eq a => a -> a -> Bool
/= k
repy
  DisjointMap Map k k
p Map k (Ranked k v)
r <- StateT (DisjointMap k v) Identity (DisjointMap k v)
-> MaybeT (StateT (DisjointMap k v) Identity) (DisjointMap k v)
forall (m :: * -> *) a. Monad m => m a -> MaybeT m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift StateT (DisjointMap k v) Identity (DisjointMap k v)
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 Map k (Ranked k v) -> k -> Ranked k v
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 Map k (Ranked k v) -> k -> Ranked k v
forall k a. Ord k => Map k a -> k -> a
M.! k
repy
  let val :: v
val = v -> v -> v
forall a. Monoid a => a -> a -> a
mappend v
valx v
valy
      keys :: Set k
keys = Set k -> Set k -> Set k
forall a. Monoid a => a -> a -> a
mappend Set k
keysx Set k
keysy
  StateT (DisjointMap k v) Identity ()
-> MaybeT (StateT (DisjointMap k v) Identity) ()
forall (m :: * -> *) a. Monad m => m a -> MaybeT m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (StateT (DisjointMap k v) Identity ()
 -> MaybeT (StateT (DisjointMap k v) Identity) ())
-> StateT (DisjointMap k v) Identity ()
-> MaybeT (StateT (DisjointMap k v) Identity) ()
forall a b. (a -> b) -> a -> b
$ DisjointMap k v -> StateT (DisjointMap k v) Identity ()
forall (m :: * -> *) s. Monad m => s -> StateT s m ()
put (DisjointMap k v -> StateT (DisjointMap k v) Identity ())
-> DisjointMap k v -> StateT (DisjointMap k v) Identity ()
forall a b. (a -> b) -> a -> b
$! case Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Int
rankx Int
ranky of
    Ordering
LT ->
      let p' :: Map k k
p' = k -> k -> Map k k -> Map k k
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' = k -> Map k (Ranked k v) -> Map k (Ranked k v)
forall k a. Ord k => k -> Map k a -> Map k a
M.delete k
repx (Map k (Ranked k v) -> Map k (Ranked k v))
-> Map k (Ranked k v) -> Map k (Ranked k v)
forall a b. (a -> b) -> a -> b
$! k -> Ranked k v -> Map k (Ranked k v) -> Map k (Ranked k v)
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert k
repy (Int -> Set k -> v -> Ranked k v
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 Map k k -> Map k (Ranked k v) -> DisjointMap k v
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' = k -> k -> Map k k -> Map k k
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' = k -> Map k (Ranked k v) -> Map k (Ranked k v)
forall k a. Ord k => k -> Map k a -> Map k a
M.delete k
repy (Map k (Ranked k v) -> Map k (Ranked k v))
-> Map k (Ranked k v) -> Map k (Ranked k v)
forall a b. (a -> b) -> a -> b
$! k -> Ranked k v -> Map k (Ranked k v) -> Map k (Ranked k v)
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert k
repx (Int -> Set k -> v -> Ranked k v
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 Map k k -> Map k (Ranked k v) -> DisjointMap k v
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' = k -> k -> Map k k -> Map k k
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' = k -> Map k (Ranked k v) -> Map k (Ranked k v)
forall k a. Ord k => k -> Map k a -> Map k a
M.delete k
repx (Map k (Ranked k v) -> Map k (Ranked k v))
-> Map k (Ranked k v) -> Map k (Ranked k v)
forall a b. (a -> b) -> a -> b
$! k -> Ranked k v -> Map k (Ranked k v) -> Map k (Ranked k v)
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert k
repy (Int -> Set k -> v -> Ranked k v
forall k v. Int -> Set k -> v -> Ranked k v
Ranked (Int
ranky Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Set k
keys v
val) Map k (Ranked k v)
r
       in Map k k -> Map k (Ranked k v) -> DisjointMap k v
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 = (State (DisjointMap k v) (Maybe ())
 -> DisjointMap k v -> DisjointMap k v)
-> DisjointMap k v
-> State (DisjointMap k v) (Maybe ())
-> DisjointMap k v
forall a b c. (a -> b -> c) -> b -> a -> c
flip State (DisjointMap k v) (Maybe ())
-> DisjointMap k v -> DisjointMap k v
forall s a. State s a -> s -> s
execState DisjointMap k v
set (State (DisjointMap k v) (Maybe ()) -> DisjointMap k v)
-> State (DisjointMap k v) (Maybe ()) -> DisjointMap k v
forall a b. (a -> b) -> a -> b
$ MaybeT (StateT (DisjointMap k v) Identity) ()
-> State (DisjointMap k v) (Maybe ())
forall (m :: * -> *) a. MaybeT m a -> m (Maybe a)
runMaybeT (MaybeT (StateT (DisjointMap k v) Identity) ()
 -> State (DisjointMap k v) (Maybe ()))
-> MaybeT (StateT (DisjointMap k v) Identity) ()
-> State (DisjointMap k v) (Maybe ())
forall a b. (a -> b) -> a -> b
$ do
  Maybe k
mx <- StateT (DisjointMap k v) Identity (Maybe k)
-> MaybeT (StateT (DisjointMap k v) Identity) (Maybe k)
forall (m :: * -> *) a. Monad m => m a -> MaybeT m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (StateT (DisjointMap k v) Identity (Maybe k)
 -> MaybeT (StateT (DisjointMap k v) Identity) (Maybe k))
-> StateT (DisjointMap k v) Identity (Maybe k)
-> MaybeT (StateT (DisjointMap k v) Identity) (Maybe k)
forall a b. (a -> b) -> a -> b
$ (DisjointMap k v -> (Maybe k, DisjointMap k v))
-> StateT (DisjointMap k v) Identity (Maybe k)
forall (m :: * -> *) s a. Monad m => (s -> (a, s)) -> StateT s m a
state ((DisjointMap k v -> (Maybe k, DisjointMap k v))
 -> StateT (DisjointMap k v) Identity (Maybe k))
-> (DisjointMap k v -> (Maybe k, DisjointMap k v))
-> StateT (DisjointMap k v) Identity (Maybe k)
forall a b. (a -> b) -> a -> b
$ k -> DisjointMap k v -> (Maybe k, DisjointMap k v)
forall k v.
Ord k =>
k -> DisjointMap k v -> (Maybe k, DisjointMap k v)
lookupCompress k
x
  Maybe k
my <- StateT (DisjointMap k v) Identity (Maybe k)
-> MaybeT (StateT (DisjointMap k v) Identity) (Maybe k)
forall (m :: * -> *) a. Monad m => m a -> MaybeT m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (StateT (DisjointMap k v) Identity (Maybe k)
 -> MaybeT (StateT (DisjointMap k v) Identity) (Maybe k))
-> StateT (DisjointMap k v) Identity (Maybe k)
-> MaybeT (StateT (DisjointMap k v) Identity) (Maybe k)
forall a b. (a -> b) -> a -> b
$ (DisjointMap k v -> (Maybe k, DisjointMap k v))
-> StateT (DisjointMap k v) Identity (Maybe k)
forall (m :: * -> *) s a. Monad m => (s -> (a, s)) -> StateT s m a
state ((DisjointMap k v -> (Maybe k, DisjointMap k v))
 -> StateT (DisjointMap k v) Identity (Maybe k))
-> (DisjointMap k v -> (Maybe k, DisjointMap k v))
-> StateT (DisjointMap k v) Identity (Maybe k)
forall a b. (a -> b) -> a -> b
$ k -> DisjointMap k v -> (Maybe k, DisjointMap k v)
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 -> () -> MaybeT (StateT (DisjointMap k v) Identity) ()
forall a. a -> MaybeT (StateT (DisjointMap k v) Identity) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
      Just k
repy -> do
        DisjointMap Map k k
p Map k (Ranked k v)
r <- StateT (DisjointMap k v) Identity (DisjointMap k v)
-> MaybeT (StateT (DisjointMap k v) Identity) (DisjointMap k v)
forall (m :: * -> *) a. Monad m => m a -> MaybeT m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift StateT (DisjointMap k v) Identity (DisjointMap k v)
forall (m :: * -> *) s. Monad m => StateT s m s
get
        StateT (DisjointMap k v) Identity ()
-> MaybeT (StateT (DisjointMap k v) Identity) ()
forall (m :: * -> *) a. Monad m => m a -> MaybeT m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (StateT (DisjointMap k v) Identity ()
 -> MaybeT (StateT (DisjointMap k v) Identity) ())
-> StateT (DisjointMap k v) Identity ()
-> MaybeT (StateT (DisjointMap k v) Identity) ()
forall a b. (a -> b) -> a -> b
$
          DisjointMap k v -> StateT (DisjointMap k v) Identity ()
forall (m :: * -> *) s. Monad m => s -> StateT s m ()
put (DisjointMap k v -> StateT (DisjointMap k v) Identity ())
-> DisjointMap k v -> StateT (DisjointMap k v) Identity ()
forall a b. (a -> b) -> a -> b
$
            let p' :: Map k k
p' = k -> k -> Map k k -> Map k k
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 = Ranked k v -> Maybe (Ranked k v) -> Ranked k v
forall a. a -> Maybe a -> a
fromMaybe (String -> Ranked k v
forall a. HasCallStack => String -> a
error String
"Data.DisjointMap.unionWeakly") (k -> Map k (Ranked k v) -> Maybe (Ranked k v)
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' = k -> Ranked k v -> Map k (Ranked k v) -> Map k (Ranked k v)
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert k
repy (Int -> Set k -> v -> Ranked k v
forall k v. Int -> Set k -> v -> Ranked k v
Ranked Int
ranky (k -> Set k -> Set k
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 Map k k -> Map k (Ranked k v) -> DisjointMap k v
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 <- StateT (DisjointMap k v) Identity (DisjointMap k v)
-> MaybeT (StateT (DisjointMap k v) Identity) (DisjointMap k v)
forall (m :: * -> *) a. Monad m => m a -> MaybeT m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift StateT (DisjointMap k v) Identity (DisjointMap k v)
forall (m :: * -> *) s. Monad m => StateT s m s
get
        StateT (DisjointMap k v) Identity ()
-> MaybeT (StateT (DisjointMap k v) Identity) ()
forall (m :: * -> *) a. Monad m => m a -> MaybeT m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (StateT (DisjointMap k v) Identity ()
 -> MaybeT (StateT (DisjointMap k v) Identity) ())
-> StateT (DisjointMap k v) Identity ()
-> MaybeT (StateT (DisjointMap k v) Identity) ()
forall a b. (a -> b) -> a -> b
$
          DisjointMap k v -> StateT (DisjointMap k v) Identity ()
forall (m :: * -> *) s. Monad m => s -> StateT s m ()
put (DisjointMap k v -> StateT (DisjointMap k v) Identity ())
-> DisjointMap k v -> StateT (DisjointMap k v) Identity ()
forall a b. (a -> b) -> a -> b
$
            let p' :: Map k k
p' = k -> k -> Map k k -> Map k k
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 = Ranked k v -> Maybe (Ranked k v) -> Ranked k v
forall a. a -> Maybe a -> a
fromMaybe (String -> Ranked k v
forall a. HasCallStack => String -> a
error String
"Data.DisjointMap.unionWeakly") (k -> Map k (Ranked k v) -> Maybe (Ranked k v)
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' = k -> Ranked k v -> Map k (Ranked k v) -> Map k (Ranked k v)
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert k
repx (Int -> Set k -> v -> Ranked k v
forall k v. Int -> Set k -> v -> Ranked k v
Ranked Int
rankx (k -> Set k -> Set k
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 Map k k -> Map k (Ranked k v) -> DisjointMap k v
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
        Bool -> MaybeT (StateT (DisjointMap k v) Identity) ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> MaybeT (StateT (DisjointMap k v) Identity) ())
-> Bool -> MaybeT (StateT (DisjointMap k v) Identity) ()
forall a b. (a -> b) -> a -> b
$ k
repx k -> k -> Bool
forall a. Eq a => a -> a -> Bool
/= k
repy
        DisjointMap Map k k
p Map k (Ranked k v)
r <- StateT (DisjointMap k v) Identity (DisjointMap k v)
-> MaybeT (StateT (DisjointMap k v) Identity) (DisjointMap k v)
forall (m :: * -> *) a. Monad m => m a -> MaybeT m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift StateT (DisjointMap k v) Identity (DisjointMap k v)
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 Map k (Ranked k v) -> k -> Ranked k v
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 Map k (Ranked k v) -> k -> Ranked k v
forall k a. Ord k => Map k a -> k -> a
M.! k
repy
        let val :: v
val = v
valx v -> v -> v
forall a. Semigroup a => a -> a -> a
<> v
valy
        let keys :: Set k
keys = Set k -> Set k -> Set k
forall a. Monoid a => a -> a -> a
mappend Set k
keysx Set k
keysy
        StateT (DisjointMap k v) Identity ()
-> MaybeT (StateT (DisjointMap k v) Identity) ()
forall (m :: * -> *) a. Monad m => m a -> MaybeT m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (StateT (DisjointMap k v) Identity ()
 -> MaybeT (StateT (DisjointMap k v) Identity) ())
-> StateT (DisjointMap k v) Identity ()
-> MaybeT (StateT (DisjointMap k v) Identity) ()
forall a b. (a -> b) -> a -> b
$ DisjointMap k v -> StateT (DisjointMap k v) Identity ()
forall (m :: * -> *) s. Monad m => s -> StateT s m ()
put (DisjointMap k v -> StateT (DisjointMap k v) Identity ())
-> DisjointMap k v -> StateT (DisjointMap k v) Identity ()
forall a b. (a -> b) -> a -> b
$! case Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Int
rankx Int
ranky of
          Ordering
LT ->
            let p' :: Map k k
p' = k -> k -> Map k k -> Map k k
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' = k -> Map k (Ranked k v) -> Map k (Ranked k v)
forall k a. Ord k => k -> Map k a -> Map k a
M.delete k
repx (Map k (Ranked k v) -> Map k (Ranked k v))
-> Map k (Ranked k v) -> Map k (Ranked k v)
forall a b. (a -> b) -> a -> b
$! k -> Ranked k v -> Map k (Ranked k v) -> Map k (Ranked k v)
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert k
repy (Int -> Set k -> v -> Ranked k v
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 Map k k -> Map k (Ranked k v) -> DisjointMap k v
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' = k -> k -> Map k k -> Map k k
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' = k -> Map k (Ranked k v) -> Map k (Ranked k v)
forall k a. Ord k => k -> Map k a -> Map k a
M.delete k
repy (Map k (Ranked k v) -> Map k (Ranked k v))
-> Map k (Ranked k v) -> Map k (Ranked k v)
forall a b. (a -> b) -> a -> b
$! k -> Ranked k v -> Map k (Ranked k v) -> Map k (Ranked k v)
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert k
repx (Int -> Set k -> v -> Ranked k v
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 Map k k -> Map k (Ranked k v) -> DisjointMap k v
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' = k -> k -> Map k k -> Map k k
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' = k -> Map k (Ranked k v) -> Map k (Ranked k v)
forall k a. Ord k => k -> Map k a -> Map k a
M.delete k
repx (Map k (Ranked k v) -> Map k (Ranked k v))
-> Map k (Ranked k v) -> Map k (Ranked k v)
forall a b. (a -> b) -> a -> b
$! k -> Ranked k v -> Map k (Ranked k v) -> Map k (Ranked k v)
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert k
repy (Int -> Set k -> v -> Ranked k v
forall k v. Int -> Set k -> v -> Ranked k v
Ranked (Int
ranky Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Set k
keys v
val) Map k (Ranked k v)
r
             in Map k k -> Map k (Ranked k v) -> DisjointMap k v
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 = k -> DisjointMap k v -> Maybe k
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 = k -> Set k -> v -> DisjointMap k v -> DisjointMap k v
forall k v.
(Ord k, Semigroup v) =>
k -> Set k -> v -> DisjointMap k v -> DisjointMap k v
insertInternal k
x (k -> Set k
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') = (k -> k -> k -> k) -> k -> k -> Map k k -> (Maybe k, Map k k)
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') = k -> DisjointMap k v -> (Maybe k, DisjointMap k v)
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 -> String -> DisjointMap k v
forall a. HasCallStack => String -> a
error String
"DisjointMap insert: invariant violated"
                Just k
root -> Map k k -> Map k (Ranked k v) -> DisjointMap k v
forall k v. Map k k -> Map k (Ranked k v) -> DisjointMap k v
DisjointMap Map k k
p2 ((Ranked k v -> Ranked k v)
-> k -> Map k (Ranked k v) -> Map k (Ranked k v)
forall k a. Ord k => (a -> a) -> k -> Map k a -> Map k a
M.adjust (\(Ranked Int
rank Set k
oldKs v
vOld) -> Int -> Set k -> v -> Ranked k v
forall k v. Int -> Set k -> v -> Ranked k v
Ranked Int
rank (Set k -> Set k -> Set k
forall a. Monoid a => a -> a -> a
mappend Set k
oldKs Set k
ks) (v
v v -> 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' = k -> Ranked k v -> Map k (Ranked k v) -> Map k (Ranked k v)
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert k
x (Int -> Set k -> v -> Ranked k v
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 Map k k -> Map k (Ranked k v) -> DisjointMap k v
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 = k -> k -> Map k k
forall k a. k -> a -> Map k a
M.singleton k
x k
x
      r :: Map k (Ranked k v)
r = k -> Ranked k v -> Map k (Ranked k v)
forall k a. k -> a -> Map k a
M.singleton k
x (Int -> Set k -> v -> Ranked k v
forall k v. Int -> Set k -> v -> Ranked k v
Ranked Int
0 (k -> Set k
forall a. a -> Set a
S.singleton k
x) v
v)
   in Map k k -> Map k (Ranked k v) -> DisjointMap k v
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 = Map k k -> Map k (Ranked k v) -> DisjointMap k v
forall k v. Map k k -> Map k (Ranked k v) -> DisjointMap k v
DisjointMap Map k k
forall k a. Map k a
M.empty Map k (Ranked k v)
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 Map k k -> Int
forall k a. Map k a -> Int
M.size Map k k
m1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Map k k -> Int
forall k a. Map k a -> Int
M.size Map k k
m2
    then DisjointMap k v -> Map k k -> DisjointMap k v
forall k v.
(Ord k, Semigroup v) =>
DisjointMap k v -> Map k k -> DisjointMap k v
appendPhase2 (Map k (Ranked k v) -> DisjointMap k v -> Map k k -> DisjointMap k v
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 DisjointMap k v -> Map k k -> DisjointMap k v
forall k v.
(Ord k, Semigroup v) =>
DisjointMap k v -> Map k k -> DisjointMap k v
appendPhase2 (Map k (Ranked k v) -> DisjointMap k v -> Map k k -> DisjointMap k v
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 = (DisjointMap k v -> k -> k -> DisjointMap k v)
-> DisjointMap k v -> Map k k -> DisjointMap k v
forall a k b. (a -> k -> b -> a) -> a -> Map k b -> a
M.foldlWithKey' ((DisjointMap k v -> k -> k -> DisjointMap k v)
 -> DisjointMap k v -> Map k k -> DisjointMap k v)
-> (DisjointMap k v -> k -> k -> DisjointMap k v)
-> DisjointMap k v
-> Map k k
-> DisjointMap k v
forall a b. (a -> b) -> a -> b
$ \DisjointMap k v
ds k
x k
y ->
  if k
x k -> k -> Bool
forall a. Eq a => a -> a -> Bool
== k
y
    then case k -> Map k (Ranked k v) -> Maybe (Ranked k v)
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 -> String -> DisjointMap k v
forall a. HasCallStack => String -> a
error String
"Data.DisjointMap.appendParents: invariant violated"
      Just (Ranked Int
_ Set k
ks v
v) -> (DisjointMap k v -> k -> DisjointMap k v)
-> DisjointMap k v -> Set k -> DisjointMap k v
forall b a. (b -> a -> b) -> b -> Set a -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
F.foldl' (\DisjointMap k v
dm k
k -> k -> k -> DisjointMap k v -> DisjointMap k v
forall k v.
(Ord k, Semigroup v) =>
k -> k -> DisjointMap k v -> DisjointMap k v
unionWeakly k
k k
x DisjointMap k v
dm) (k -> v -> DisjointMap k v -> DisjointMap k v
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 = (DisjointMap k v -> k -> k -> DisjointMap k v)
-> DisjointMap k v -> Map k k -> DisjointMap k v
forall a k b. (a -> k -> b -> a) -> a -> Map k b -> a
M.foldlWithKey' ((DisjointMap k v -> k -> k -> DisjointMap k v)
 -> DisjointMap k v -> Map k k -> DisjointMap k v)
-> (DisjointMap k v -> k -> k -> DisjointMap k v)
-> DisjointMap k v
-> Map k k
-> DisjointMap k v
forall a b. (a -> b) -> a -> b
$ \DisjointMap k v
ds k
x k
y ->
  if k
x k -> k -> Bool
forall a. Eq a => a -> a -> Bool
== k
y
    then DisjointMap k v
ds
    else k -> k -> DisjointMap k v -> DisjointMap k v
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 Set k -> Maybe k
forall a. Set a -> Maybe a
setLookupMin Set k
s of
  Maybe k
Nothing -> DisjointMap k v
forall k v. DisjointMap k v
empty
  Just k
x ->
    let p :: Map k k
p = (k -> k) -> Set k -> Map k k
forall k a. (k -> a) -> Set k -> Map k a
M.fromSet (\k
_ -> k
x) Set k
s
        rank :: Int
rank = if Set k -> Int
forall a. Set a -> Int
S.size Set k
s Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1 then Int
0 else Int
1
        r :: Map k (Ranked k v)
r = k -> Ranked k v -> Map k (Ranked k v)
forall k a. k -> a -> Map k a
M.singleton k
x (Int -> Set k -> v -> Ranked k v
forall k v. Int -> Set k -> v -> Ranked k v
Ranked Int
rank Set k
s v
v)
     in Map k k -> Map k (Ranked k v) -> DisjointMap k v
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 = Set a -> Maybe a
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 k -> DisjointMap k v -> Maybe k
forall k v. Ord k => k -> DisjointMap k v -> Maybe k
find k
x DisjointMap k v
set of
    Maybe k
Nothing -> (Maybe k
forall a. Maybe a
Nothing, DisjointMap k v
set)
    Just k
rep ->
      let set' :: DisjointMap k v
set' = k -> k -> DisjointMap k v -> DisjointMap k v
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' DisjointMap k v
-> (Maybe k, DisjointMap k v) -> (Maybe k, DisjointMap k v)
forall a b. a -> b -> b
`seq` (k -> Maybe k
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 k -> DisjointMap k v -> Maybe k
forall k v. Ord k => k -> DisjointMap k v -> Maybe k
find k
x DisjointMap k v
set of
    Maybe k
Nothing -> (k
x, k -> v -> DisjointMap k v -> DisjointMap k v
forall k v.
(Ord k, Semigroup v) =>
k -> v -> DisjointMap k v -> DisjointMap k v
insert k
x v
forall a. Monoid a => a
mempty DisjointMap k v
set)
    Just k
rep ->
      let !set' :: DisjointMap k v
set' = k -> k -> DisjointMap k v -> DisjointMap k v
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 k -> DisjointMap k v -> Maybe k
forall k v. Ord k => k -> DisjointMap k v -> Maybe k
find k
x DisjointMap k v
set of
    Maybe k
Nothing -> (Maybe k
forall a. Maybe a
Nothing, DisjointMap k v
set)
    Just k
rep ->
      let !set' :: DisjointMap k v
set' = k -> k -> DisjointMap k v -> DisjointMap k v
forall k v. Ord k => k -> k -> DisjointMap k v -> DisjointMap k v
compress k
rep k
x DisjointMap k v
set
       in (k -> Maybe k
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' <- k -> Map k k -> Maybe k
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup k
x Map k k
p
  k -> Maybe k
forall a. a -> Maybe a
forall (m :: * -> *) a. Monad m => a -> m a
return (k -> Maybe k) -> k -> Maybe k
forall a b. (a -> b) -> a -> b
$! if k
x k -> k -> Bool
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 Map k k -> k -> k
forall k a. Ord k => Map k a -> k -> a
M.! k
y
     in if k
y k -> k -> Bool
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 = v -> Maybe v -> v
forall a. a -> Maybe a -> a
fromMaybe v
forall a. Monoid a => a
mempty (Maybe v -> v)
-> (DisjointMap k v -> Maybe v) -> DisjointMap k v -> v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. k -> DisjointMap k v -> Maybe v
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' <- k -> Map k k -> Maybe k
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup k
x Map k k
p
  if k
x k -> k -> Bool
forall a. Eq a => a -> a -> Bool
== k
x'
    then case k -> Map k (Ranked k v) -> Maybe (Ranked k v)
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 -> Maybe v
forall a. Maybe a
Nothing
      Just (Ranked Int
_ Set k
_ v
v) -> v -> Maybe 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 Map k k -> k -> k
forall k a. Ord k => Map k a -> k -> a
M.! k
y
     in if k
y k -> k -> Bool
forall a. Eq a => a -> a -> Bool
== k
y'
          then case k -> Map k (Ranked k v) -> Maybe (Ranked k v)
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 -> Maybe v
forall a. Maybe a
Nothing
            Just (Ranked Int
_ Set k
_ v
v) -> v -> Maybe 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 = k -> DisjointMap k v -> DisjointMap k v
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 k -> k -> Bool
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 Map k k -> k -> k
forall k a. Ord k => Map k a -> k -> a
M.! k
x
    set' :: DisjointMap k v
set' =
      let !p' :: Map k k
p' = k -> k -> Map k k -> Map k k
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert k
x k
rep Map k k
p
       in Map k k -> Map k (Ranked k v) -> DisjointMap k v
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 = (<>)

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 String 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
-}