Find.hs

-- | (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 {unUnionFind :: UnionMap k (IntLikeSet k)}
  deriving stock (Eq, Show)

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

ufToUmIso :: Iso' (UnionFind k) (UnionMap k (IntLikeSet k))
ufToUmIso = coerced

empty :: UnionFind k
empty = UnionFind UM.empty

singleton :: (Coercible k Int) => k -> UnionFind k
singleton k = UnionFind (UM.singleton k (ILS.singleton k))

size :: UnionFind k -> Int
size = UM.size . unUnionFind

member :: (Coercible k Int) => k -> UnionFind k -> Bool
member k = UM.member k . unUnionFind

data InsertRes k
  = InsertResAdded !(UnionFind k)
  | InsertResDuplicate
  deriving stock (Eq, Show)

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

data InsertVal
  = InsertValInserted
  | InsertValDuplicate
  deriving stock (Eq, Show)

insertLM :: (Coercible k Int, MonadState s m) => UnionFindLens s k -> k -> m InsertVal
insertLM l k =
  UM.addLM (l % ufToUmIso) k (ILS.singleton k) <&> \case
    UM.AddValAdded -> InsertValInserted
    UM.AddValDuplicate -> InsertValDuplicate

insertM :: (Coercible k Int, MonadState (UnionFind k) m) => k -> m InsertVal
insertM = insertLM equality'

data LookupRes k
  = LookupResMissing !k
  | LookupResFound !k !(IntLikeSet k) !(Maybe (UnionFind k))
  deriving stock (Eq, Show)

lookup :: (Coercible k Int) => k -> UnionFind k -> LookupRes k
lookup k (UnionFind um) =
  case UM.lookup k um of
    UM.LookupResMissing k' -> LookupResMissing k'
    UM.LookupResFound k' x y -> LookupResFound k' x (fmap UnionFind y)

data LookupVal k
  = LookupValMissing !k
  | LookupValOk !k !(IntLikeSet k) !Changed
  deriving stock (Eq, Show)

lookupLM :: (Coercible k Int, MonadState s m) => UnionFindLens s k -> k -> m (LookupVal k)
lookupLM l k =
  UM.lookupLM (l % ufToUmIso) k <&> \case
    UM.LookupValMissing k' -> LookupValMissing k'
    UM.LookupValOk x y c -> LookupValOk x y c

lookupM :: (Coercible k Int, MonadState (UnionFind k) m) => k -> m (LookupVal k)
lookupM = lookupLM equality'

equiv :: (Coercible k Int) => UnionFind k -> (Equiv k, Maybe (UnionFind k))
equiv = second (fmap UnionFind) . UM.equiv . unUnionFind

equivLM :: (Coercible k Int, MonadState s m) => UnionFindLens s k -> m (Equiv k)
equivLM l = UM.equivLM (l % ufToUmIso)

equivM :: (Coercible k Int, MonadState (UnionFind k) m) => m (Equiv k)
equivM = equivLM equality'

compact :: (Coercible k Int) => UnionFind k -> (IntLikeMap k k, UnionFind k)
compact = second UnionFind . UM.compact . unUnionFind

compactLM :: (Coercible k Int, MonadState s m) => UnionFindLens s k -> m (IntLikeMap k k)
compactLM l = UM.compactLM (l % ufToUmIso)

compactM :: (Coercible k Int, MonadState (UnionFind k) m) => m (IntLikeMap k k)
compactM = compactLM equality'

data MergeRes k
  = MergeResMissing !k
  | MergeResMerged !k !(IntLikeSet k) !(UnionFind k)
  deriving stock (Eq, Show)

data MergeVal k
  = MergeValMissing !k
  | MergeValMerged !k !(IntLikeSet k)
  deriving stock (Eq, 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 = UM.foldMergeOne (\x y -> Right ((), x <> y))

manyFn :: (Foldable f) => UM.MergeMany f Void (IntLikeSet k) ()
manyFn = UM.foldMergeMany (Right mempty) (\x y -> Right ((), x <> y))

mergeOne :: (Coercible k Int, Eq k) => k -> k -> UnionFind k -> MergeRes k
mergeOne k j (UnionFind um) =
  case UM.mergeOne oneFn k j um of
    UM.MergeResMissing k' -> MergeResMissing k'
    UM.MergeResMerged k' x _ y -> MergeResMerged k' x (UnionFind y)

mergeOneLM
  :: (Coercible k Int, Eq k, MonadState s m) => UnionFindLens s k -> k -> k -> m (MergeVal k)
mergeOneLM l k j =
  UM.mergeOneLM (l % ufToUmIso) oneFn k j <&> \case
    UM.MergeValMissing k' -> MergeValMissing k'
    UM.MergeValMerged k' x _ -> MergeValMerged k' x

mergeOneM :: (Coercible k Int, Eq k, MonadState (UnionFind k) m) => k -> k -> m (MergeVal k)
mergeOneM = mergeOneLM equality'

mergeMany :: (Traversable f, Coercible k Int, Eq k) => k -> f k -> UnionFind k -> MergeRes k
mergeMany k js (UnionFind um) =
  case UM.mergeMany manyFn k js um of
    UM.MergeResMissing k' -> MergeResMissing k'
    UM.MergeResMerged k' x _ y -> MergeResMerged k' x (UnionFind y)

mergeManyLM
  :: (Traversable f, Coercible k Int, Eq k, MonadState s m)
  => UnionFindLens s k
  -> k
  -> f k
  -> m (MergeVal k)
mergeManyLM l k js =
  UM.mergeManyLM (l % ufToUmIso) manyFn k js <&> \case
    UM.MergeValMissing k' -> MergeValMissing k'
    UM.MergeValMerged k' x _ -> MergeValMerged k' x

mergeManyM
  :: (Traversable f, Coercible k Int, Eq k, MonadState (UnionFind k) m)
  => k
  -> f k
  -> m (MergeVal k)
mergeManyM = mergeManyLM equality'