-- | (Import this module qualified)
module Unimap.Find
  ( Changed (..)
  , Equiv (..)
  , UnionFind
  , empty
  , singleton
  , size
  , member
  , InsertRes (..)
  , InsertVal (..)
  , insert
  , insertLM
  , insertM
  , LookupRes (..)
  , LookupVal (..)
  , lookup
  , lookupLM
  , lookupM
  , equiv
  , equivLM
  , equivM
  , compact
  , compactLM
  , compactM
  , MergeRes (..)
  , MergeVal (..)
  , mergeOne
  , mergeOneLM
  , mergeOneM
  , mergeMany
  , mergeManyLM
  , mergeManyM
  )
where

import Control.Monad.State.Strict (MonadState)
import Data.Bifunctor (second)
import Data.Coerce (Coercible)
import Data.Functor ((<&>))
import Data.Void (Void)
import IntLike.Map (IntLikeMap)
import IntLike.Set (IntLikeSet)
import IntLike.Set qualified as ILS
import Optics (Lens')
import Optics.Coerce (coerceA, coerceB)
import Optics.Lens (equality')
import Unimap (Changed (..), Equiv, UnionMap)
import Unimap qualified as UM
import Prelude hiding (lookup)

coerceUFL :: UnionFindLens s k -> UM.UnionMapLens s k (IntLikeSet k)
coerceUFL :: forall s k. UnionFindLens s k -> UnionMapLens s k (IntLikeSet k)
coerceUFL = Optic A_Lens NoIx s s (UnionFind k) (UnionMap k (IntLikeSet k))
-> Optic
     A_Lens
     NoIx
     s
     s
     (UnionMap k (IntLikeSet k))
     (UnionMap k (IntLikeSet k))
forall a a' k (is :: IxList) s t b.
Coercible a a' =>
Optic k is s t a b -> Optic k is s t a' b
coerceA (Optic A_Lens NoIx s s (UnionFind k) (UnionMap k (IntLikeSet k))
 -> Optic
      A_Lens
      NoIx
      s
      s
      (UnionMap k (IntLikeSet k))
      (UnionMap k (IntLikeSet k)))
-> (UnionFindLens s k
    -> Optic A_Lens NoIx s s (UnionFind k) (UnionMap k (IntLikeSet k)))
-> UnionFindLens s k
-> Optic
     A_Lens
     NoIx
     s
     s
     (UnionMap k (IntLikeSet k))
     (UnionMap k (IntLikeSet k))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. UnionFindLens s k
-> Optic A_Lens NoIx s s (UnionFind k) (UnionMap k (IntLikeSet k))
forall b b' k (is :: IxList) s t a.
Coercible b b' =>
Optic k is s t a b -> Optic k is s t a b'
coerceB
{-# INLINE coerceUFL #-}

newtype UnionFind k = UnionFind {forall k. UnionFind k -> UnionMap k (IntLikeSet k)
unUnionFind :: UnionMap k (IntLikeSet k)}
  deriving stock (UnionFind k -> UnionFind k -> Bool
(UnionFind k -> UnionFind k -> Bool)
-> (UnionFind k -> UnionFind k -> Bool) -> Eq (UnionFind k)
forall k. Eq k => UnionFind k -> UnionFind k -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall k. Eq k => UnionFind k -> UnionFind k -> Bool
== :: UnionFind k -> UnionFind k -> Bool
$c/= :: forall k. Eq k => UnionFind k -> UnionFind k -> Bool
/= :: UnionFind k -> UnionFind k -> Bool
Eq, Int -> UnionFind k -> ShowS
[UnionFind k] -> ShowS
UnionFind k -> String
(Int -> UnionFind k -> ShowS)
-> (UnionFind k -> String)
-> ([UnionFind k] -> ShowS)
-> Show (UnionFind k)
forall k. Show k => Int -> UnionFind k -> ShowS
forall k. Show k => [UnionFind k] -> ShowS
forall k. Show k => UnionFind k -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall k. Show k => Int -> UnionFind k -> ShowS
showsPrec :: Int -> UnionFind k -> ShowS
$cshow :: forall k. Show k => UnionFind k -> String
show :: UnionFind k -> String
$cshowList :: forall k. Show k => [UnionFind k] -> ShowS
showList :: [UnionFind k] -> ShowS
Show)

type UnionFindLens s k = Lens' s (UnionFind k)

empty :: UnionFind k
empty :: forall k. UnionFind k
empty = UnionMap k (IntLikeSet k) -> UnionFind k
forall k. UnionMap k (IntLikeSet k) -> UnionFind k
UnionFind UnionMap k (IntLikeSet k)
forall k v. UnionMap k v
UM.empty

singleton :: (Coercible k Int) => k -> UnionFind k
singleton :: forall k. Coercible k Int => k -> UnionFind k
singleton k
k = UnionMap k (IntLikeSet k) -> UnionFind k
forall k. UnionMap k (IntLikeSet k) -> UnionFind k
UnionFind (k -> IntLikeSet k -> UnionMap k (IntLikeSet k)
forall k v. Coercible k Int => k -> v -> UnionMap k v
UM.singleton k
k (k -> IntLikeSet k
forall x. Coercible x Int => x -> IntLikeSet x
ILS.singleton k
k))

size :: UnionFind k -> Int
size :: forall k. UnionFind k -> Int
size = UnionMap k (IntLikeSet k) -> Int
forall k v. UnionMap k v -> Int
UM.size (UnionMap k (IntLikeSet k) -> Int)
-> (UnionFind k -> UnionMap k (IntLikeSet k)) -> UnionFind k -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. UnionFind k -> UnionMap k (IntLikeSet k)
forall k. UnionFind k -> UnionMap k (IntLikeSet k)
unUnionFind

member :: (Coercible k Int) => k -> UnionFind k -> Bool
member :: forall k. Coercible k Int => k -> UnionFind k -> Bool
member k
k = k -> UnionMap k (IntLikeSet k) -> Bool
forall k v. Coercible k Int => k -> UnionMap k v -> Bool
UM.member k
k (UnionMap k (IntLikeSet k) -> Bool)
-> (UnionFind k -> UnionMap k (IntLikeSet k))
-> UnionFind k
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. UnionFind k -> UnionMap k (IntLikeSet k)
forall k. UnionFind k -> UnionMap k (IntLikeSet k)
unUnionFind

data InsertRes k
  = InsertResAdded !(UnionFind k)
  | InsertResDuplicate
  deriving stock (InsertRes k -> InsertRes k -> Bool
(InsertRes k -> InsertRes k -> Bool)
-> (InsertRes k -> InsertRes k -> Bool) -> Eq (InsertRes k)
forall k. Eq k => InsertRes k -> InsertRes k -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall k. Eq k => InsertRes k -> InsertRes k -> Bool
== :: InsertRes k -> InsertRes k -> Bool
$c/= :: forall k. Eq k => InsertRes k -> InsertRes k -> Bool
/= :: InsertRes k -> InsertRes k -> Bool
Eq, Int -> InsertRes k -> ShowS
[InsertRes k] -> ShowS
InsertRes k -> String
(Int -> InsertRes k -> ShowS)
-> (InsertRes k -> String)
-> ([InsertRes k] -> ShowS)
-> Show (InsertRes k)
forall k. Show k => Int -> InsertRes k -> ShowS
forall k. Show k => [InsertRes k] -> ShowS
forall k. Show k => InsertRes k -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall k. Show k => Int -> InsertRes k -> ShowS
showsPrec :: Int -> InsertRes k -> ShowS
$cshow :: forall k. Show k => InsertRes k -> String
show :: InsertRes k -> String
$cshowList :: forall k. Show k => [InsertRes k] -> ShowS
showList :: [InsertRes k] -> ShowS
Show)

insert :: (Coercible k Int) => k -> UnionFind k -> InsertRes k
insert :: forall k. Coercible k Int => k -> UnionFind k -> InsertRes k
insert k
k (UnionFind UnionMap k (IntLikeSet k)
um) =
  case k
-> IntLikeSet k
-> UnionMap k (IntLikeSet k)
-> AddRes k (IntLikeSet k)
forall k v. Coercible k Int => k -> v -> UnionMap k v -> AddRes k v
UM.add k
k (k -> IntLikeSet k
forall x. Coercible x Int => x -> IntLikeSet x
ILS.singleton k
k) UnionMap k (IntLikeSet k)
um of
    UM.AddResAdded UnionMap k (IntLikeSet k)
um' -> UnionFind k -> InsertRes k
forall k. UnionFind k -> InsertRes k
InsertResAdded (UnionMap k (IntLikeSet k) -> UnionFind k
forall k. UnionMap k (IntLikeSet k) -> UnionFind k
UnionFind UnionMap k (IntLikeSet k)
um')
    AddRes k (IntLikeSet k)
UM.AddResDuplicate -> InsertRes k
forall k. InsertRes k
InsertResDuplicate

data InsertVal
  = InsertValInserted
  | InsertValDuplicate
  deriving stock (InsertVal -> InsertVal -> Bool
(InsertVal -> InsertVal -> Bool)
-> (InsertVal -> InsertVal -> Bool) -> Eq InsertVal
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: InsertVal -> InsertVal -> Bool
== :: InsertVal -> InsertVal -> Bool
$c/= :: InsertVal -> InsertVal -> Bool
/= :: InsertVal -> InsertVal -> Bool
Eq, Int -> InsertVal -> ShowS
[InsertVal] -> ShowS
InsertVal -> String
(Int -> InsertVal -> ShowS)
-> (InsertVal -> String)
-> ([InsertVal] -> ShowS)
-> Show InsertVal
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> InsertVal -> ShowS
showsPrec :: Int -> InsertVal -> ShowS
$cshow :: InsertVal -> String
show :: InsertVal -> String
$cshowList :: [InsertVal] -> ShowS
showList :: [InsertVal] -> ShowS
Show)

insertLM :: (Coercible k Int, MonadState s m) => UnionFindLens s k -> k -> m InsertVal
insertLM :: forall k s (m :: * -> *).
(Coercible k Int, MonadState s m) =>
UnionFindLens s k -> k -> m InsertVal
insertLM UnionFindLens s k
l k
k =
  UnionMapLens s k (IntLikeSet k) -> k -> IntLikeSet k -> m AddVal
forall k s (m :: * -> *) v.
(Coercible k Int, MonadState s m) =>
UnionMapLens s k v -> k -> v -> m AddVal
UM.addLM (UnionFindLens s k -> UnionMapLens s k (IntLikeSet k)
forall s k. UnionFindLens s k -> UnionMapLens s k (IntLikeSet k)
coerceUFL UnionFindLens s k
l) k
k (k -> IntLikeSet k
forall x. Coercible x Int => x -> IntLikeSet x
ILS.singleton k
k) m AddVal -> (AddVal -> InsertVal) -> m InsertVal
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \case
    AddVal
UM.AddValAdded -> InsertVal
InsertValInserted
    AddVal
UM.AddValDuplicate -> InsertVal
InsertValDuplicate

insertM :: (Coercible k Int, MonadState (UnionFind k) m) => k -> m InsertVal
insertM :: forall k (m :: * -> *).
(Coercible k Int, MonadState (UnionFind k) m) =>
k -> m InsertVal
insertM = UnionFindLens (UnionFind k) k -> k -> m InsertVal
forall k s (m :: * -> *).
(Coercible k Int, MonadState s m) =>
UnionFindLens s k -> k -> m InsertVal
insertLM UnionFindLens (UnionFind k) k
forall a b. Lens a b a b
equality'

data LookupRes k
  = LookupResMissing !k
  | LookupResFound !k !(IntLikeSet k) !(Maybe (UnionFind k))
  deriving stock (LookupRes k -> LookupRes k -> Bool
(LookupRes k -> LookupRes k -> Bool)
-> (LookupRes k -> LookupRes k -> Bool) -> Eq (LookupRes k)
forall k. Eq k => LookupRes k -> LookupRes k -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall k. Eq k => LookupRes k -> LookupRes k -> Bool
== :: LookupRes k -> LookupRes k -> Bool
$c/= :: forall k. Eq k => LookupRes k -> LookupRes k -> Bool
/= :: LookupRes k -> LookupRes k -> Bool
Eq, Int -> LookupRes k -> ShowS
[LookupRes k] -> ShowS
LookupRes k -> String
(Int -> LookupRes k -> ShowS)
-> (LookupRes k -> String)
-> ([LookupRes k] -> ShowS)
-> Show (LookupRes k)
forall k. Show k => Int -> LookupRes k -> ShowS
forall k. Show k => [LookupRes k] -> ShowS
forall k. Show k => LookupRes k -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall k. Show k => Int -> LookupRes k -> ShowS
showsPrec :: Int -> LookupRes k -> ShowS
$cshow :: forall k. Show k => LookupRes k -> String
show :: LookupRes k -> String
$cshowList :: forall k. Show k => [LookupRes k] -> ShowS
showList :: [LookupRes k] -> ShowS
Show)

lookup :: (Coercible k Int) => k -> UnionFind k -> LookupRes k
lookup :: forall k. Coercible k Int => k -> UnionFind k -> LookupRes k
lookup k
k (UnionFind UnionMap k (IntLikeSet k)
um) =
  case k -> UnionMap k (IntLikeSet k) -> LookupRes k (IntLikeSet k)
forall k v. Coercible k Int => k -> UnionMap k v -> LookupRes k v
UM.lookup k
k UnionMap k (IntLikeSet k)
um of
    UM.LookupResMissing k
k' -> k -> LookupRes k
forall k. k -> LookupRes k
LookupResMissing k
k'
    UM.LookupResFound k
k' IntLikeSet k
x Maybe (UnionMap k (IntLikeSet k))
y -> k -> IntLikeSet k -> Maybe (UnionFind k) -> LookupRes k
forall k. k -> IntLikeSet k -> Maybe (UnionFind k) -> LookupRes k
LookupResFound k
k' IntLikeSet k
x ((UnionMap k (IntLikeSet k) -> UnionFind k)
-> Maybe (UnionMap k (IntLikeSet k)) -> Maybe (UnionFind k)
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap UnionMap k (IntLikeSet k) -> UnionFind k
forall k. UnionMap k (IntLikeSet k) -> UnionFind k
UnionFind Maybe (UnionMap k (IntLikeSet k))
y)

data LookupVal k
  = LookupValMissing !k
  | LookupValOk !k !(IntLikeSet k) !Changed
  deriving stock (LookupVal k -> LookupVal k -> Bool
(LookupVal k -> LookupVal k -> Bool)
-> (LookupVal k -> LookupVal k -> Bool) -> Eq (LookupVal k)
forall k. Eq k => LookupVal k -> LookupVal k -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall k. Eq k => LookupVal k -> LookupVal k -> Bool
== :: LookupVal k -> LookupVal k -> Bool
$c/= :: forall k. Eq k => LookupVal k -> LookupVal k -> Bool
/= :: LookupVal k -> LookupVal k -> Bool
Eq, Int -> LookupVal k -> ShowS
[LookupVal k] -> ShowS
LookupVal k -> String
(Int -> LookupVal k -> ShowS)
-> (LookupVal k -> String)
-> ([LookupVal k] -> ShowS)
-> Show (LookupVal k)
forall k. Show k => Int -> LookupVal k -> ShowS
forall k. Show k => [LookupVal k] -> ShowS
forall k. Show k => LookupVal k -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall k. Show k => Int -> LookupVal k -> ShowS
showsPrec :: Int -> LookupVal k -> ShowS
$cshow :: forall k. Show k => LookupVal k -> String
show :: LookupVal k -> String
$cshowList :: forall k. Show k => [LookupVal k] -> ShowS
showList :: [LookupVal k] -> ShowS
Show)

lookupLM :: (Coercible k Int, MonadState s m) => UnionFindLens s k -> k -> m (LookupVal k)
lookupLM :: forall k s (m :: * -> *).
(Coercible k Int, MonadState s m) =>
UnionFindLens s k -> k -> m (LookupVal k)
lookupLM UnionFindLens s k
l k
k =
  UnionMapLens s k (IntLikeSet k)
-> k -> m (LookupVal k (IntLikeSet k))
forall k s (m :: * -> *) v.
(Coercible k Int, MonadState s m) =>
UnionMapLens s k v -> k -> m (LookupVal k v)
UM.lookupLM (UnionFindLens s k -> UnionMapLens s k (IntLikeSet k)
forall s k. UnionFindLens s k -> UnionMapLens s k (IntLikeSet k)
coerceUFL UnionFindLens s k
l) k
k m (LookupVal k (IntLikeSet k))
-> (LookupVal k (IntLikeSet k) -> LookupVal k) -> m (LookupVal k)
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \case
    UM.LookupValMissing k
k' -> k -> LookupVal k
forall k. k -> LookupVal k
LookupValMissing k
k'
    UM.LookupValOk k
x IntLikeSet k
y Changed
c -> k -> IntLikeSet k -> Changed -> LookupVal k
forall k. k -> IntLikeSet k -> Changed -> LookupVal k
LookupValOk k
x IntLikeSet k
y Changed
c

lookupM :: (Coercible k Int, MonadState (UnionFind k) m) => k -> m (LookupVal k)
lookupM :: forall k (m :: * -> *).
(Coercible k Int, MonadState (UnionFind k) m) =>
k -> m (LookupVal k)
lookupM = UnionFindLens (UnionFind k) k -> k -> m (LookupVal k)
forall k s (m :: * -> *).
(Coercible k Int, MonadState s m) =>
UnionFindLens s k -> k -> m (LookupVal k)
lookupLM UnionFindLens (UnionFind k) k
forall a b. Lens a b a b
equality'

equiv :: (Coercible k Int) => UnionFind k -> (Equiv k, Maybe (UnionFind k))
equiv :: forall k.
Coercible k Int =>
UnionFind k -> (Equiv k, Maybe (UnionFind k))
equiv = (Maybe (UnionMap k (IntLikeSet k)) -> Maybe (UnionFind k))
-> (Equiv k, Maybe (UnionMap k (IntLikeSet k)))
-> (Equiv k, Maybe (UnionFind k))
forall b c a. (b -> c) -> (a, b) -> (a, c)
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second ((UnionMap k (IntLikeSet k) -> UnionFind k)
-> Maybe (UnionMap k (IntLikeSet k)) -> Maybe (UnionFind k)
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap UnionMap k (IntLikeSet k) -> UnionFind k
forall k. UnionMap k (IntLikeSet k) -> UnionFind k
UnionFind) ((Equiv k, Maybe (UnionMap k (IntLikeSet k)))
 -> (Equiv k, Maybe (UnionFind k)))
-> (UnionFind k -> (Equiv k, Maybe (UnionMap k (IntLikeSet k))))
-> UnionFind k
-> (Equiv k, Maybe (UnionFind k))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. UnionMap k (IntLikeSet k)
-> (Equiv k, Maybe (UnionMap k (IntLikeSet k)))
forall k v.
Coercible k Int =>
UnionMap k v -> (Equiv k, Maybe (UnionMap k v))
UM.equiv (UnionMap k (IntLikeSet k)
 -> (Equiv k, Maybe (UnionMap k (IntLikeSet k))))
-> (UnionFind k -> UnionMap k (IntLikeSet k))
-> UnionFind k
-> (Equiv k, Maybe (UnionMap k (IntLikeSet k)))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. UnionFind k -> UnionMap k (IntLikeSet k)
forall k. UnionFind k -> UnionMap k (IntLikeSet k)
unUnionFind

equivLM :: (Coercible k Int, MonadState s m) => UnionFindLens s k -> m (Equiv k)
equivLM :: forall k s (m :: * -> *).
(Coercible k Int, MonadState s m) =>
UnionFindLens s k -> m (Equiv k)
equivLM UnionFindLens s k
l = UnionMapLens s k (IntLikeSet k) -> m (Equiv k)
forall k s (m :: * -> *) v.
(Coercible k Int, MonadState s m) =>
UnionMapLens s k v -> m (Equiv k)
UM.equivLM (UnionFindLens s k -> UnionMapLens s k (IntLikeSet k)
forall s k. UnionFindLens s k -> UnionMapLens s k (IntLikeSet k)
coerceUFL UnionFindLens s k
l)

equivM :: (Coercible k Int, MonadState (UnionFind k) m) => m (Equiv k)
equivM :: forall k (m :: * -> *).
(Coercible k Int, MonadState (UnionFind k) m) =>
m (Equiv k)
equivM = UnionFindLens (UnionFind k) k -> m (Equiv k)
forall k s (m :: * -> *).
(Coercible k Int, MonadState s m) =>
UnionFindLens s k -> m (Equiv k)
equivLM UnionFindLens (UnionFind k) k
forall a b. Lens a b a b
equality'

compact :: (Coercible k Int) => UnionFind k -> (IntLikeMap k k, UnionFind k)
compact :: forall k.
Coercible k Int =>
UnionFind k -> (IntLikeMap k k, UnionFind k)
compact = (UnionMap k (IntLikeSet k) -> UnionFind k)
-> (IntLikeMap k k, UnionMap k (IntLikeSet k))
-> (IntLikeMap k k, UnionFind k)
forall b c a. (b -> c) -> (a, b) -> (a, c)
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second UnionMap k (IntLikeSet k) -> UnionFind k
forall k. UnionMap k (IntLikeSet k) -> UnionFind k
UnionFind ((IntLikeMap k k, UnionMap k (IntLikeSet k))
 -> (IntLikeMap k k, UnionFind k))
-> (UnionFind k -> (IntLikeMap k k, UnionMap k (IntLikeSet k)))
-> UnionFind k
-> (IntLikeMap k k, UnionFind k)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. UnionMap k (IntLikeSet k)
-> (IntLikeMap k k, UnionMap k (IntLikeSet k))
forall k v.
Coercible k Int =>
UnionMap k v -> (IntLikeMap k k, UnionMap k v)
UM.compact (UnionMap k (IntLikeSet k)
 -> (IntLikeMap k k, UnionMap k (IntLikeSet k)))
-> (UnionFind k -> UnionMap k (IntLikeSet k))
-> UnionFind k
-> (IntLikeMap k k, UnionMap k (IntLikeSet k))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. UnionFind k -> UnionMap k (IntLikeSet k)
forall k. UnionFind k -> UnionMap k (IntLikeSet k)
unUnionFind

compactLM :: (Coercible k Int, MonadState s m) => UnionFindLens s k -> m (IntLikeMap k k)
compactLM :: forall k s (m :: * -> *).
(Coercible k Int, MonadState s m) =>
UnionFindLens s k -> m (IntLikeMap k k)
compactLM UnionFindLens s k
l = UnionMapLens s k (IntLikeSet k) -> m (IntLikeMap k k)
forall k s (m :: * -> *) v.
(Coercible k Int, MonadState s m) =>
UnionMapLens s k v -> m (IntLikeMap k k)
UM.compactLM (UnionFindLens s k -> UnionMapLens s k (IntLikeSet k)
forall s k. UnionFindLens s k -> UnionMapLens s k (IntLikeSet k)
coerceUFL UnionFindLens s k
l)

compactM :: (Coercible k Int, MonadState (UnionFind k) m) => m (IntLikeMap k k)
compactM :: forall k (m :: * -> *).
(Coercible k Int, MonadState (UnionFind k) m) =>
m (IntLikeMap k k)
compactM = UnionFindLens (UnionFind k) k -> m (IntLikeMap k k)
forall k s (m :: * -> *).
(Coercible k Int, MonadState s m) =>
UnionFindLens s k -> m (IntLikeMap k k)
compactLM UnionFindLens (UnionFind k) k
forall a b. Lens a b a b
equality'

data MergeRes k
  = MergeResMissing !k
  | MergeResMerged !k !(IntLikeSet k) !(UnionFind k)
  deriving stock (MergeRes k -> MergeRes k -> Bool
(MergeRes k -> MergeRes k -> Bool)
-> (MergeRes k -> MergeRes k -> Bool) -> Eq (MergeRes k)
forall k. Eq k => MergeRes k -> MergeRes k -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall k. Eq k => MergeRes k -> MergeRes k -> Bool
== :: MergeRes k -> MergeRes k -> Bool
$c/= :: forall k. Eq k => MergeRes k -> MergeRes k -> Bool
/= :: MergeRes k -> MergeRes k -> Bool
Eq, Int -> MergeRes k -> ShowS
[MergeRes k] -> ShowS
MergeRes k -> String
(Int -> MergeRes k -> ShowS)
-> (MergeRes k -> String)
-> ([MergeRes k] -> ShowS)
-> Show (MergeRes k)
forall k. Show k => Int -> MergeRes k -> ShowS
forall k. Show k => [MergeRes k] -> ShowS
forall k. Show k => MergeRes k -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall k. Show k => Int -> MergeRes k -> ShowS
showsPrec :: Int -> MergeRes k -> ShowS
$cshow :: forall k. Show k => MergeRes k -> String
show :: MergeRes k -> String
$cshowList :: forall k. Show k => [MergeRes k] -> ShowS
showList :: [MergeRes k] -> ShowS
Show)

data MergeVal k
  = MergeValMissing !k
  | MergeValMerged !k !(IntLikeSet k)
  deriving stock (MergeVal k -> MergeVal k -> Bool
(MergeVal k -> MergeVal k -> Bool)
-> (MergeVal k -> MergeVal k -> Bool) -> Eq (MergeVal k)
forall k. Eq k => MergeVal k -> MergeVal k -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall k. Eq k => MergeVal k -> MergeVal k -> Bool
== :: MergeVal k -> MergeVal k -> Bool
$c/= :: forall k. Eq k => MergeVal k -> MergeVal k -> Bool
/= :: MergeVal k -> MergeVal k -> Bool
Eq, Int -> MergeVal k -> ShowS
[MergeVal k] -> ShowS
MergeVal k -> String
(Int -> MergeVal k -> ShowS)
-> (MergeVal k -> String)
-> ([MergeVal k] -> ShowS)
-> Show (MergeVal k)
forall k. Show k => Int -> MergeVal k -> ShowS
forall k. Show k => [MergeVal k] -> ShowS
forall k. Show k => MergeVal k -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall k. Show k => Int -> MergeVal k -> ShowS
showsPrec :: Int -> MergeVal k -> ShowS
$cshow :: forall k. Show k => MergeVal k -> String
show :: MergeVal k -> String
$cshowList :: forall k. Show k => [MergeVal k] -> ShowS
showList :: [MergeVal k] -> ShowS
Show)

type MergeOne e v r = Maybe v -> v -> Either e (r, v)

type MergeMany f e v r = Maybe v -> f v -> Either e (r, v)

oneFn :: UM.MergeOne Void (IntLikeSet k) ()
oneFn :: forall k. MergeOne Void (IntLikeSet k) ()
oneFn = (IntLikeSet k -> IntLikeSet k -> Either Void ((), IntLikeSet k))
-> MergeOne Void (IntLikeSet k) ()
forall r v e.
Monoid r =>
(v -> v -> Either e (r, v)) -> MergeOne e v r
UM.foldMergeOne (\IntLikeSet k
x IntLikeSet k
y -> ((), IntLikeSet k) -> Either Void ((), IntLikeSet k)
forall a b. b -> Either a b
Right ((), IntLikeSet k
x IntLikeSet k -> IntLikeSet k -> IntLikeSet k
forall a. Semigroup a => a -> a -> a
<> IntLikeSet k
y))

manyFn :: (Foldable f) => UM.MergeMany f Void (IntLikeSet k) ()
manyFn :: forall (f :: * -> *) k.
Foldable f =>
MergeMany f Void (IntLikeSet k) ()
manyFn = Either Void (IntLikeSet k)
-> (IntLikeSet k -> IntLikeSet k -> Either Void ((), IntLikeSet k))
-> MergeMany f Void (IntLikeSet k) ()
forall (f :: * -> *) r e v.
(Foldable f, Monoid r) =>
Either e v -> (v -> v -> Either e (r, v)) -> MergeMany f e v r
UM.foldMergeMany (IntLikeSet k -> Either Void (IntLikeSet k)
forall a b. b -> Either a b
Right IntLikeSet k
forall a. Monoid a => a
mempty) (\IntLikeSet k
x IntLikeSet k
y -> ((), IntLikeSet k) -> Either Void ((), IntLikeSet k)
forall a b. b -> Either a b
Right ((), IntLikeSet k
x IntLikeSet k -> IntLikeSet k -> IntLikeSet k
forall a. Semigroup a => a -> a -> a
<> IntLikeSet k
y))

mergeOne :: (Coercible k Int, Eq k) => k -> k -> UnionFind k -> MergeRes k
mergeOne :: forall k.
(Coercible k Int, Eq k) =>
k -> k -> UnionFind k -> MergeRes k
mergeOne k
k k
j (UnionFind UnionMap k (IntLikeSet k)
um) =
  case MergeOne Void (IntLikeSet k) ()
-> k
-> k
-> UnionMap k (IntLikeSet k)
-> MergeRes Void k (IntLikeSet k) ()
forall k e v r.
(Coercible k Int, Eq k) =>
MergeOne e v r -> k -> k -> UnionMap k v -> MergeRes e k v r
UM.mergeOne MergeOne Void (IntLikeSet k) ()
forall k. MergeOne Void (IntLikeSet k) ()
oneFn k
k k
j UnionMap k (IntLikeSet k)
um of
    UM.MergeResMissing k
k' -> k -> MergeRes k
forall k. k -> MergeRes k
MergeResMissing k
k'
    UM.MergeResMerged k
k' IntLikeSet k
x ()
_ UnionMap k (IntLikeSet k)
y -> k -> IntLikeSet k -> UnionFind k -> MergeRes k
forall k. k -> IntLikeSet k -> UnionFind k -> MergeRes k
MergeResMerged k
k' IntLikeSet k
x (UnionMap k (IntLikeSet k) -> UnionFind k
forall k. UnionMap k (IntLikeSet k) -> UnionFind k
UnionFind UnionMap k (IntLikeSet k)
y)

mergeOneLM
  :: (Coercible k Int, Eq k, MonadState s m) => UnionFindLens s k -> k -> k -> m (MergeVal k)
mergeOneLM :: forall k s (m :: * -> *).
(Coercible k Int, Eq k, MonadState s m) =>
UnionFindLens s k -> k -> k -> m (MergeVal k)
mergeOneLM UnionFindLens s k
l k
k k
j =
  UnionMapLens s k (IntLikeSet k)
-> MergeOne Void (IntLikeSet k) ()
-> k
-> k
-> m (MergeVal Void k (IntLikeSet k) ())
forall k s (m :: * -> *) v e r.
(Coercible k Int, Eq k, MonadState s m) =>
UnionMapLens s k v
-> MergeOne e v r -> k -> k -> m (MergeVal e k v r)
UM.mergeOneLM (UnionFindLens s k -> UnionMapLens s k (IntLikeSet k)
forall s k. UnionFindLens s k -> UnionMapLens s k (IntLikeSet k)
coerceUFL UnionFindLens s k
l) MergeOne Void (IntLikeSet k) ()
forall k. MergeOne Void (IntLikeSet k) ()
oneFn k
k k
j m (MergeVal Void k (IntLikeSet k) ())
-> (MergeVal Void k (IntLikeSet k) () -> MergeVal k)
-> m (MergeVal k)
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \case
    UM.MergeValMissing k
k' -> k -> MergeVal k
forall k. k -> MergeVal k
MergeValMissing k
k'
    UM.MergeValMerged k
k' IntLikeSet k
x ()
_ -> k -> IntLikeSet k -> MergeVal k
forall k. k -> IntLikeSet k -> MergeVal k
MergeValMerged k
k' IntLikeSet k
x

mergeOneM :: (Coercible k Int, Eq k, MonadState (UnionFind k) m) => k -> k -> m (MergeVal k)
mergeOneM :: forall k (m :: * -> *).
(Coercible k Int, Eq k, MonadState (UnionFind k) m) =>
k -> k -> m (MergeVal k)
mergeOneM = UnionFindLens (UnionFind k) k -> k -> k -> m (MergeVal k)
forall k s (m :: * -> *).
(Coercible k Int, Eq k, MonadState s m) =>
UnionFindLens s k -> k -> k -> m (MergeVal k)
mergeOneLM UnionFindLens (UnionFind k) k
forall a b. Lens a b a b
equality'

mergeMany :: (Traversable f, Coercible k Int, Eq k) => k -> f k -> UnionFind k -> MergeRes k
mergeMany :: forall (f :: * -> *) k.
(Traversable f, Coercible k Int, Eq k) =>
k -> f k -> UnionFind k -> MergeRes k
mergeMany k
k f k
js (UnionFind UnionMap k (IntLikeSet k)
um) =
  case MergeMany f Void (IntLikeSet k) ()
-> k
-> f k
-> UnionMap k (IntLikeSet k)
-> MergeRes Void k (IntLikeSet k) ()
forall (f :: * -> *) k e v r.
(Traversable f, Coercible k Int, Eq k) =>
MergeMany f e v r -> k -> f k -> UnionMap k v -> MergeRes e k v r
UM.mergeMany MergeMany f Void (IntLikeSet k) ()
forall (f :: * -> *) k.
Foldable f =>
MergeMany f Void (IntLikeSet k) ()
manyFn k
k f k
js UnionMap k (IntLikeSet k)
um of
    UM.MergeResMissing k
k' -> k -> MergeRes k
forall k. k -> MergeRes k
MergeResMissing k
k'
    UM.MergeResMerged k
k' IntLikeSet k
x ()
_ UnionMap k (IntLikeSet k)
y -> k -> IntLikeSet k -> UnionFind k -> MergeRes k
forall k. k -> IntLikeSet k -> UnionFind k -> MergeRes k
MergeResMerged k
k' IntLikeSet k
x (UnionMap k (IntLikeSet k) -> UnionFind k
forall k. UnionMap k (IntLikeSet k) -> UnionFind k
UnionFind UnionMap k (IntLikeSet k)
y)

mergeManyLM
  :: (Traversable f, Coercible k Int, Eq k, MonadState s m)
  => UnionFindLens s k
  -> k
  -> f k
  -> m (MergeVal k)
mergeManyLM :: forall (f :: * -> *) k s (m :: * -> *).
(Traversable f, Coercible k Int, Eq k, MonadState s m) =>
UnionFindLens s k -> k -> f k -> m (MergeVal k)
mergeManyLM UnionFindLens s k
l k
k f k
js =
  UnionMapLens s k (IntLikeSet k)
-> MergeMany f Void (IntLikeSet k) ()
-> k
-> f k
-> m (MergeVal Void k (IntLikeSet k) ())
forall (f :: * -> *) k s (m :: * -> *) v e r.
(Traversable f, Coercible k Int, Eq k, MonadState s m) =>
UnionMapLens s k v
-> MergeMany f e v r -> k -> f k -> m (MergeVal e k v r)
UM.mergeManyLM (UnionFindLens s k -> UnionMapLens s k (IntLikeSet k)
forall s k. UnionFindLens s k -> UnionMapLens s k (IntLikeSet k)
coerceUFL UnionFindLens s k
l) MergeMany f Void (IntLikeSet k) ()
forall (f :: * -> *) k.
Foldable f =>
MergeMany f Void (IntLikeSet k) ()
manyFn k
k f k
js m (MergeVal Void k (IntLikeSet k) ())
-> (MergeVal Void k (IntLikeSet k) () -> MergeVal k)
-> m (MergeVal k)
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \case
    UM.MergeValMissing k
k' -> k -> MergeVal k
forall k. k -> MergeVal k
MergeValMissing k
k'
    UM.MergeValMerged k
k' IntLikeSet k
x ()
_ -> k -> IntLikeSet k -> MergeVal k
forall k. k -> IntLikeSet k -> MergeVal k
MergeValMerged k
k' IntLikeSet k
x

mergeManyM
  :: (Traversable f, Coercible k Int, Eq k, MonadState (UnionFind k) m)
  => k
  -> f k
  -> m (MergeVal k)
mergeManyM :: forall (f :: * -> *) k (m :: * -> *).
(Traversable f, Coercible k Int, Eq k,
 MonadState (UnionFind k) m) =>
k -> f k -> m (MergeVal k)
mergeManyM = UnionFindLens (UnionFind k) k -> k -> f k -> m (MergeVal k)
forall (f :: * -> *) k s (m :: * -> *).
(Traversable f, Coercible k Int, Eq k, MonadState s m) =>
UnionFindLens s k -> k -> f k -> m (MergeVal k)
mergeManyLM UnionFindLens (UnionFind k) k
forall a b. Lens a b a b
equality'