-- | (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 (Iso', Lens', coerced, (%))
import Optics.Lens (equality')
import Unimap (Changed (..), Equiv, UnionMap)
import Unimap qualified as UM
import Prelude hiding (lookup)

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)

ufToUmIso :: Iso' (UnionFind k) (UnionMap k (IntLikeSet k))
ufToUmIso :: forall k. Iso' (UnionFind k) (UnionMap k (IntLikeSet k))
ufToUmIso = Iso
  (UnionFind k)
  (UnionFind k)
  (UnionMap k (IntLikeSet k))
  (UnionMap k (IntLikeSet k))
forall s a t b. (Coercible s a, Coercible t b) => Iso s t a b
coerced

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
l UnionFindLens s k
-> Optic
     An_Iso
     NoIx
     (UnionFind k)
     (UnionFind k)
     (UnionMap k (IntLikeSet k))
     (UnionMap k (IntLikeSet k))
-> UnionMapLens s k (IntLikeSet k)
forall k l m (is :: IxList) (js :: IxList) (ks :: IxList) s t u v a
       b.
(JoinKinds k l m, AppendIndices is js ks) =>
Optic k is s t u v -> Optic l js u v a b -> Optic m ks s t a b
% Optic
  An_Iso
  NoIx
  (UnionFind k)
  (UnionFind k)
  (UnionMap k (IntLikeSet k))
  (UnionMap k (IntLikeSet k))
forall k. Iso' (UnionFind k) (UnionMap k (IntLikeSet k))
ufToUmIso) 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
l UnionFindLens s k
-> Optic
     An_Iso
     NoIx
     (UnionFind k)
     (UnionFind k)
     (UnionMap k (IntLikeSet k))
     (UnionMap k (IntLikeSet k))
-> UnionMapLens s k (IntLikeSet k)
forall k l m (is :: IxList) (js :: IxList) (ks :: IxList) s t u v a
       b.
(JoinKinds k l m, AppendIndices is js ks) =>
Optic k is s t u v -> Optic l js u v a b -> Optic m ks s t a b
% Optic
  An_Iso
  NoIx
  (UnionFind k)
  (UnionFind k)
  (UnionMap k (IntLikeSet k))
  (UnionMap k (IntLikeSet k))
forall k. Iso' (UnionFind k) (UnionMap k (IntLikeSet k))
ufToUmIso) 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
l UnionFindLens s k
-> Optic
     An_Iso
     NoIx
     (UnionFind k)
     (UnionFind k)
     (UnionMap k (IntLikeSet k))
     (UnionMap k (IntLikeSet k))
-> UnionMapLens s k (IntLikeSet k)
forall k l m (is :: IxList) (js :: IxList) (ks :: IxList) s t u v a
       b.
(JoinKinds k l m, AppendIndices is js ks) =>
Optic k is s t u v -> Optic l js u v a b -> Optic m ks s t a b
% Optic
  An_Iso
  NoIx
  (UnionFind k)
  (UnionFind k)
  (UnionMap k (IntLikeSet k))
  (UnionMap k (IntLikeSet k))
forall k. Iso' (UnionFind k) (UnionMap k (IntLikeSet k))
ufToUmIso)

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
l UnionFindLens s k
-> Optic
     An_Iso
     NoIx
     (UnionFind k)
     (UnionFind k)
     (UnionMap k (IntLikeSet k))
     (UnionMap k (IntLikeSet k))
-> UnionMapLens s k (IntLikeSet k)
forall k l m (is :: IxList) (js :: IxList) (ks :: IxList) s t u v a
       b.
(JoinKinds k l m, AppendIndices is js ks) =>
Optic k is s t u v -> Optic l js u v a b -> Optic m ks s t a b
% Optic
  An_Iso
  NoIx
  (UnionFind k)
  (UnionFind k)
  (UnionMap k (IntLikeSet k))
  (UnionMap k (IntLikeSet k))
forall k. Iso' (UnionFind k) (UnionMap k (IntLikeSet k))
ufToUmIso)

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
l UnionFindLens s k
-> Optic
     An_Iso
     NoIx
     (UnionFind k)
     (UnionFind k)
     (UnionMap k (IntLikeSet k))
     (UnionMap k (IntLikeSet k))
-> UnionMapLens s k (IntLikeSet k)
forall k l m (is :: IxList) (js :: IxList) (ks :: IxList) s t u v a
       b.
(JoinKinds k l m, AppendIndices is js ks) =>
Optic k is s t u v -> Optic l js u v a b -> Optic m ks s t a b
% Optic
  An_Iso
  NoIx
  (UnionFind k)
  (UnionFind k)
  (UnionMap k (IntLikeSet k))
  (UnionMap k (IntLikeSet k))
forall k. Iso' (UnionFind k) (UnionMap k (IntLikeSet k))
ufToUmIso) 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
l UnionFindLens s k
-> Optic
     An_Iso
     NoIx
     (UnionFind k)
     (UnionFind k)
     (UnionMap k (IntLikeSet k))
     (UnionMap k (IntLikeSet k))
-> UnionMapLens s k (IntLikeSet k)
forall k l m (is :: IxList) (js :: IxList) (ks :: IxList) s t u v a
       b.
(JoinKinds k l m, AppendIndices is js ks) =>
Optic k is s t u v -> Optic l js u v a b -> Optic m ks s t a b
% Optic
  An_Iso
  NoIx
  (UnionFind k)
  (UnionFind k)
  (UnionMap k (IntLikeSet k))
  (UnionMap k (IntLikeSet k))
forall k. Iso' (UnionFind k) (UnionMap k (IntLikeSet k))
ufToUmIso) 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'