{-# LANGUAGE RecursiveDo #-}

module Unison.PatternMatchCoverage.UFMap
  ( UFMap,
    UFValue (..),
    empty,
    lookupCanon,
    insert,
    union,
    alterF,
    alter,
    keys,
    toClasses,
  )
where

import Control.Monad.Fix (MonadFix)
import Control.Monad.Trans.Class
import Control.Monad.Trans.Except (ExceptT (..))
import Data.Foldable (foldl')
import Data.Functor ((<&>))
import Data.Functor.Compose (Compose (..))
import Data.Functor.Identity (Identity (Identity, runIdentity))
import Data.Functor.Sum (Sum (..))
import Data.Map (Map)
import Data.Map.Lazy qualified as LazyMap
import Data.Map.Strict qualified as Map
import Data.Maybe (fromJust)
import Data.Set (Set)
import Data.Set qualified as Set

-- | A union-find structure. Used by
-- 'Unison.PatternMatchCoverage.NormalizedConstraints.NormalizedConstraints'
-- to provide efficient unification.
newtype UFMap k v = UFMap (Map k (UFValue k v))
  deriving stock (UFMap k v -> UFMap k v -> Bool
(UFMap k v -> UFMap k v -> Bool)
-> (UFMap k v -> UFMap k v -> Bool) -> Eq (UFMap k v)
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall k v. (Eq k, Eq v) => UFMap k v -> UFMap k v -> Bool
$c== :: forall k v. (Eq k, Eq v) => UFMap k v -> UFMap k v -> Bool
== :: UFMap k v -> UFMap k v -> Bool
$c/= :: forall k v. (Eq k, Eq v) => UFMap k v -> UFMap k v -> Bool
/= :: UFMap k v -> UFMap k v -> Bool
Eq, Eq (UFMap k v)
Eq (UFMap k v) =>
(UFMap k v -> UFMap k v -> Ordering)
-> (UFMap k v -> UFMap k v -> Bool)
-> (UFMap k v -> UFMap k v -> Bool)
-> (UFMap k v -> UFMap k v -> Bool)
-> (UFMap k v -> UFMap k v -> Bool)
-> (UFMap k v -> UFMap k v -> UFMap k v)
-> (UFMap k v -> UFMap k v -> UFMap k v)
-> Ord (UFMap k v)
UFMap k v -> UFMap k v -> Bool
UFMap k v -> UFMap k v -> Ordering
UFMap k v -> UFMap k v -> UFMap k v
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
forall k v. (Ord k, Ord v) => Eq (UFMap k v)
forall k v. (Ord k, Ord v) => UFMap k v -> UFMap k v -> Bool
forall k v. (Ord k, Ord v) => UFMap k v -> UFMap k v -> Ordering
forall k v. (Ord k, Ord v) => UFMap k v -> UFMap k v -> UFMap k v
$ccompare :: forall k v. (Ord k, Ord v) => UFMap k v -> UFMap k v -> Ordering
compare :: UFMap k v -> UFMap k v -> Ordering
$c< :: forall k v. (Ord k, Ord v) => UFMap k v -> UFMap k v -> Bool
< :: UFMap k v -> UFMap k v -> Bool
$c<= :: forall k v. (Ord k, Ord v) => UFMap k v -> UFMap k v -> Bool
<= :: UFMap k v -> UFMap k v -> Bool
$c> :: forall k v. (Ord k, Ord v) => UFMap k v -> UFMap k v -> Bool
> :: UFMap k v -> UFMap k v -> Bool
$c>= :: forall k v. (Ord k, Ord v) => UFMap k v -> UFMap k v -> Bool
>= :: UFMap k v -> UFMap k v -> Bool
$cmax :: forall k v. (Ord k, Ord v) => UFMap k v -> UFMap k v -> UFMap k v
max :: UFMap k v -> UFMap k v -> UFMap k v
$cmin :: forall k v. (Ord k, Ord v) => UFMap k v -> UFMap k v -> UFMap k v
min :: UFMap k v -> UFMap k v -> UFMap k v
Ord, Int -> UFMap k v -> ShowS
[UFMap k v] -> ShowS
UFMap k v -> String
(Int -> UFMap k v -> ShowS)
-> (UFMap k v -> String)
-> ([UFMap k v] -> ShowS)
-> Show (UFMap k v)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall k v. (Show k, Show v) => Int -> UFMap k v -> ShowS
forall k v. (Show k, Show v) => [UFMap k v] -> ShowS
forall k v. (Show k, Show v) => UFMap k v -> String
$cshowsPrec :: forall k v. (Show k, Show v) => Int -> UFMap k v -> ShowS
showsPrec :: Int -> UFMap k v -> ShowS
$cshow :: forall k v. (Show k, Show v) => UFMap k v -> String
show :: UFMap k v -> String
$cshowList :: forall k v. (Show k, Show v) => [UFMap k v] -> ShowS
showList :: [UFMap k v] -> ShowS
Show)

data UFValue k v
  = -- | This is not the canonical value, lookup k in the map to try again
    Indirection !k
  | -- | The number of elements in the equivalence class
    Canonical !Int !v
  deriving stock (UFValue k v -> UFValue k v -> Bool
(UFValue k v -> UFValue k v -> Bool)
-> (UFValue k v -> UFValue k v -> Bool) -> Eq (UFValue k v)
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall k v. (Eq k, Eq v) => UFValue k v -> UFValue k v -> Bool
$c== :: forall k v. (Eq k, Eq v) => UFValue k v -> UFValue k v -> Bool
== :: UFValue k v -> UFValue k v -> Bool
$c/= :: forall k v. (Eq k, Eq v) => UFValue k v -> UFValue k v -> Bool
/= :: UFValue k v -> UFValue k v -> Bool
Eq, Eq (UFValue k v)
Eq (UFValue k v) =>
(UFValue k v -> UFValue k v -> Ordering)
-> (UFValue k v -> UFValue k v -> Bool)
-> (UFValue k v -> UFValue k v -> Bool)
-> (UFValue k v -> UFValue k v -> Bool)
-> (UFValue k v -> UFValue k v -> Bool)
-> (UFValue k v -> UFValue k v -> UFValue k v)
-> (UFValue k v -> UFValue k v -> UFValue k v)
-> Ord (UFValue k v)
UFValue k v -> UFValue k v -> Bool
UFValue k v -> UFValue k v -> Ordering
UFValue k v -> UFValue k v -> UFValue k v
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
forall k v. (Ord k, Ord v) => Eq (UFValue k v)
forall k v. (Ord k, Ord v) => UFValue k v -> UFValue k v -> Bool
forall k v.
(Ord k, Ord v) =>
UFValue k v -> UFValue k v -> Ordering
forall k v.
(Ord k, Ord v) =>
UFValue k v -> UFValue k v -> UFValue k v
$ccompare :: forall k v.
(Ord k, Ord v) =>
UFValue k v -> UFValue k v -> Ordering
compare :: UFValue k v -> UFValue k v -> Ordering
$c< :: forall k v. (Ord k, Ord v) => UFValue k v -> UFValue k v -> Bool
< :: UFValue k v -> UFValue k v -> Bool
$c<= :: forall k v. (Ord k, Ord v) => UFValue k v -> UFValue k v -> Bool
<= :: UFValue k v -> UFValue k v -> Bool
$c> :: forall k v. (Ord k, Ord v) => UFValue k v -> UFValue k v -> Bool
> :: UFValue k v -> UFValue k v -> Bool
$c>= :: forall k v. (Ord k, Ord v) => UFValue k v -> UFValue k v -> Bool
>= :: UFValue k v -> UFValue k v -> Bool
$cmax :: forall k v.
(Ord k, Ord v) =>
UFValue k v -> UFValue k v -> UFValue k v
max :: UFValue k v -> UFValue k v -> UFValue k v
$cmin :: forall k v.
(Ord k, Ord v) =>
UFValue k v -> UFValue k v -> UFValue k v
min :: UFValue k v -> UFValue k v -> UFValue k v
Ord, Int -> UFValue k v -> ShowS
[UFValue k v] -> ShowS
UFValue k v -> String
(Int -> UFValue k v -> ShowS)
-> (UFValue k v -> String)
-> ([UFValue k v] -> ShowS)
-> Show (UFValue k v)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall k v. (Show k, Show v) => Int -> UFValue k v -> ShowS
forall k v. (Show k, Show v) => [UFValue k v] -> ShowS
forall k v. (Show k, Show v) => UFValue k v -> String
$cshowsPrec :: forall k v. (Show k, Show v) => Int -> UFValue k v -> ShowS
showsPrec :: Int -> UFValue k v -> ShowS
$cshow :: forall k v. (Show k, Show v) => UFValue k v -> String
show :: UFValue k v -> String
$cshowList :: forall k v. (Show k, Show v) => [UFValue k v] -> ShowS
showList :: [UFValue k v] -> ShowS
Show)

empty :: UFMap k v
empty :: forall k v. UFMap k v
empty = Map k (UFValue k v) -> UFMap k v
forall k v. Map k (UFValue k v) -> UFMap k v
UFMap Map k (UFValue k v)
forall k a. Map k a
Map.empty

insert :: (Ord k) => k -> v -> UFMap k v -> UFMap k v
insert :: forall k v. Ord k => k -> v -> UFMap k v -> UFMap k v
insert k
k !v
v UFMap k v
m =
  k
-> Maybe v
-> (k -> Int -> v -> UFValue k v)
-> UFMap k v
-> UFMap k v
forall k v.
Ord k =>
k
-> Maybe v
-> (k -> Int -> v -> UFValue k v)
-> UFMap k v
-> UFMap k v
alter k
k (v -> Maybe v
forall a. a -> Maybe a
Just v
v) (\k
_ Int
s v
_ -> Int -> v -> UFValue k v
forall k v. Int -> v -> UFValue k v
Canonical Int
s v
v) UFMap k v
m

alterF' ::
  forall f k v.
  (Functor f, Ord k) =>
  -- | The key to lookup
  k ->
  -- | The canonical key (use laziness to supply if unknown)
  k ->
  -- | Return Just to short-circuit the indirection lookup loop
  (k -> UFMap k v -> Maybe (f (UFMap k v))) ->
  -- | Nothing case
  f (Maybe v) ->
  -- | Just case
  --
  -- @canonicalKey -> size -> value -> new value@
  --
  -- /N.B./ deleting a value is not supported
  (k -> Int -> v -> f (UFValue k v)) ->
  UFMap k v ->
  -- | Returns the canonical k, the size, the value, and the path
  -- compressed UFMap
  f (UFMap k v)
alterF' :: forall (f :: * -> *) k v.
(Functor f, Ord k) =>
k
-> k
-> (k -> UFMap k v -> Maybe (f (UFMap k v)))
-> f (Maybe v)
-> (k -> Int -> v -> f (UFValue k v))
-> UFMap k v
-> f (UFMap k v)
alterF' k
k0 k
kcanon k -> UFMap k v -> Maybe (f (UFMap k v))
loopGuard f (Maybe v)
handleNothing k -> Int -> v -> f (UFValue k v)
handleJust UFMap k v
map0 =
  let phi :: k -> Maybe (UFValue k v) -> Sum ((,) k) f (Maybe (UFValue k v))
      phi :: k -> Maybe (UFValue k v) -> Sum ((,) k) f (Maybe (UFValue k v))
phi k
k =
        \case
          Maybe (UFValue k v)
Nothing -> f (Maybe (UFValue k v)) -> Sum ((,) k) f (Maybe (UFValue k v))
forall {k} (f :: k -> *) (g :: k -> *) (a :: k). g a -> Sum f g a
InR ((v -> UFValue k v) -> Maybe v -> Maybe (UFValue k v)
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Int -> v -> UFValue k v
forall k v. Int -> v -> UFValue k v
Canonical Int
1) (Maybe v -> Maybe (UFValue k v))
-> f (Maybe v) -> f (Maybe (UFValue k v))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f (Maybe v)
handleNothing)
          Just UFValue k v
alpha -> case UFValue k v
alpha of
            Indirection k
k -> (k, Maybe (UFValue k v)) -> Sum ((,) k) f (Maybe (UFValue k v))
forall {k} (f :: k -> *) (g :: k -> *) (a :: k). f a -> Sum f g a
InL (k
k, UFValue k v -> Maybe (UFValue k v)
forall a. a -> Maybe a
Just (k -> UFValue k v
forall k v. k -> UFValue k v
Indirection k
kcanon))
            Canonical Int
sizeOrig v
v -> f (Maybe (UFValue k v)) -> Sum ((,) k) f (Maybe (UFValue k v))
forall {k} (f :: k -> *) (g :: k -> *) (a :: k). g a -> Sum f g a
InR (UFValue k v -> Maybe (UFValue k v)
forall a. a -> Maybe a
Just (UFValue k v -> Maybe (UFValue k v))
-> f (UFValue k v) -> f (Maybe (UFValue k v))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> k -> Int -> v -> f (UFValue k v)
handleJust k
k Int
sizeOrig v
v)
      go :: k -> UFMap k v -> f (UFMap k v)
      go :: k -> UFMap k v -> f (UFMap k v)
go k
k ufm :: UFMap k v
ufm@(UFMap Map k (UFValue k v)
m) = case k -> UFMap k v -> Maybe (f (UFMap k v))
loopGuard k
k UFMap k v
ufm of
        Just f (UFMap k v)
short -> f (UFMap k v)
short
        Maybe (f (UFMap k v))
Nothing -> case (Maybe (UFValue k v) -> Sum ((,) k) f (Maybe (UFValue k v)))
-> k -> Map k (UFValue k v) -> Sum ((,) k) f (Map k (UFValue k v))
forall (f :: * -> *) k a.
(Functor f, Ord k) =>
(Maybe a -> f (Maybe a)) -> k -> Map k a -> f (Map k a)
LazyMap.alterF (k -> Maybe (UFValue k v) -> Sum ((,) k) f (Maybe (UFValue k v))
phi k
k) k
k Map k (UFValue k v)
m of
          InL (k
k, Map k (UFValue k v)
m') -> k -> UFMap k v -> f (UFMap k v)
go k
k (Map k (UFValue k v) -> UFMap k v
forall k v. Map k (UFValue k v) -> UFMap k v
UFMap Map k (UFValue k v)
m')
          InR f (Map k (UFValue k v))
res -> Map k (UFValue k v) -> UFMap k v
forall k v. Map k (UFValue k v) -> UFMap k v
UFMap (Map k (UFValue k v) -> UFMap k v)
-> f (Map k (UFValue k v)) -> f (UFMap k v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f (Map k (UFValue k v))
res
   in k -> UFMap k v -> f (UFMap k v)
go k
k0 UFMap k v
map0
{-# INLINE alterF' #-}

alterFWithHalt ::
  forall f k v.
  (Functor f, Ord k) =>
  k ->
  (k -> UFMap k v -> Maybe (f (UFMap k v))) ->
  f (Maybe v) ->
  (k -> Int -> v -> f (UFValue k v)) ->
  UFMap k v ->
  f (UFMap k v)
alterFWithHalt :: forall (f :: * -> *) k v.
(Functor f, Ord k) =>
k
-> (k -> UFMap k v -> Maybe (f (UFMap k v)))
-> f (Maybe v)
-> (k -> Int -> v -> f (UFValue k v))
-> UFMap k v
-> f (UFMap k v)
alterFWithHalt k
k0 k -> UFMap k v -> Maybe (f (UFMap k v))
isCanonical f (Maybe v)
handleNothing k -> Int -> v -> f (UFValue k v)
handleJust UFMap k v
map0 =
  -- tie the canonicalK knot
  let (k
canonicalK, f (UFMap k v)
res) = Compose ((,) k) f (UFMap k v) -> (k, f (UFMap k v))
forall {k1} {k2} (f :: k1 -> *) (g :: k2 -> k1) (a :: k2).
Compose f g a -> f (g a)
getCompose (k
-> k
-> (k -> UFMap k v -> Maybe (Compose ((,) k) f (UFMap k v)))
-> Compose ((,) k) f (Maybe v)
-> (k -> Int -> v -> Compose ((,) k) f (UFValue k v))
-> UFMap k v
-> Compose ((,) k) f (UFMap k v)
forall (f :: * -> *) k v.
(Functor f, Ord k) =>
k
-> k
-> (k -> UFMap k v -> Maybe (f (UFMap k v)))
-> f (Maybe v)
-> (k -> Int -> v -> f (UFValue k v))
-> UFMap k v
-> f (UFMap k v)
alterF' k
k0 k
canonicalK k -> UFMap k v -> Maybe (Compose ((,) k) f (UFMap k v))
loopGuard Compose ((,) k) f (Maybe v)
handleNothing' k -> Int -> v -> Compose ((,) k) f (UFValue k v)
handleJust' UFMap k v
map0)
      handleNothing' :: Compose ((,) k) f (Maybe v)
handleNothing' = (k, f (Maybe v)) -> Compose ((,) k) f (Maybe v)
forall {k} {k1} (f :: k -> *) (g :: k1 -> k) (a :: k1).
f (g a) -> Compose f g a
Compose (k
k0, f (Maybe v)
handleNothing)
      handleJust' :: k -> Int -> v -> Compose ((,) k) f (UFValue k v)
handleJust' k
k Int
s v
v = (k, f (UFValue k v)) -> Compose ((,) k) f (UFValue k v)
forall {k} {k1} (f :: k -> *) (g :: k1 -> k) (a :: k1).
f (g a) -> Compose f g a
Compose (k
k, k -> Int -> v -> f (UFValue k v)
handleJust k
k Int
s v
v)
      -- if the key is canonical then we halt and return it as the
      -- left element of the tuple
      loopGuard :: k -> UFMap k v -> Maybe (Compose ((,) k) f (UFMap k v))
loopGuard k
k UFMap k v
m = (k, f (UFMap k v)) -> Compose ((,) k) f (UFMap k v)
forall {k} {k1} (f :: k -> *) (g :: k1 -> k) (a :: k1).
f (g a) -> Compose f g a
Compose ((k, f (UFMap k v)) -> Compose ((,) k) f (UFMap k v))
-> (f (UFMap k v) -> (k, f (UFMap k v)))
-> f (UFMap k v)
-> Compose ((,) k) f (UFMap k v)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (k
k,) (f (UFMap k v) -> Compose ((,) k) f (UFMap k v))
-> Maybe (f (UFMap k v)) -> Maybe (Compose ((,) k) f (UFMap k v))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> k -> UFMap k v -> Maybe (f (UFMap k v))
isCanonical k
k UFMap k v
m
   in f (UFMap k v)
res
{-# INLINE alterFWithHalt #-}

alterF ::
  forall f k v.
  (Functor f, Ord k) =>
  k ->
  f (Maybe v) ->
  (k -> Int -> v -> f (UFValue k v)) ->
  UFMap k v ->
  f (UFMap k v)
alterF :: forall (f :: * -> *) k v.
(Functor f, Ord k) =>
k
-> f (Maybe v)
-> (k -> Int -> v -> f (UFValue k v))
-> UFMap k v
-> f (UFMap k v)
alterF k
k = k
-> (k -> UFMap k v -> Maybe (f (UFMap k v)))
-> f (Maybe v)
-> (k -> Int -> v -> f (UFValue k v))
-> UFMap k v
-> f (UFMap k v)
forall (f :: * -> *) k v.
(Functor f, Ord k) =>
k
-> (k -> UFMap k v -> Maybe (f (UFMap k v)))
-> f (Maybe v)
-> (k -> Int -> v -> f (UFValue k v))
-> UFMap k v
-> f (UFMap k v)
alterFWithHalt k
k (\k
_ UFMap k v
_ -> Maybe (f (UFMap k v))
forall a. Maybe a
Nothing)
{-# INLINE alterF #-}

alter ::
  forall k v.
  (Ord k) =>
  k ->
  Maybe v ->
  (k -> Int -> v -> UFValue k v) ->
  UFMap k v ->
  UFMap k v
alter :: forall k v.
Ord k =>
k
-> Maybe v
-> (k -> Int -> v -> UFValue k v)
-> UFMap k v
-> UFMap k v
alter k
k Maybe v
handleNothing k -> Int -> v -> UFValue k v
handleJust UFMap k v
map0 =
  Identity (UFMap k v) -> UFMap k v
forall a. Identity a -> a
runIdentity (k
-> Identity (Maybe v)
-> (k -> Int -> v -> Identity (UFValue k v))
-> UFMap k v
-> Identity (UFMap k v)
forall (f :: * -> *) k v.
(Functor f, Ord k) =>
k
-> f (Maybe v)
-> (k -> Int -> v -> f (UFValue k v))
-> UFMap k v
-> f (UFMap k v)
alterF k
k (Maybe v -> Identity (Maybe v)
forall a. a -> Identity a
Identity Maybe v
handleNothing) (\k
k Int
s v
v -> UFValue k v -> Identity (UFValue k v)
forall a. a -> Identity a
Identity (k -> Int -> v -> UFValue k v
handleJust k
k Int
s v
v)) UFMap k v
map0)

-- | Lookup the canonical value
lookupCanon ::
  (Ord k) =>
  k ->
  UFMap k v ->
  -- | returns:
  --
  -- * the canonical member of the equivalence set
  -- * the size of the equivalence set
  -- * the associated value
  -- * the @UFMap@ after path compression
  Maybe (k, Int, v, UFMap k v)
lookupCanon :: forall k v. Ord k => k -> UFMap k v -> Maybe (k, Int, v, UFMap k v)
lookupCanon k
k UFMap k v
m =
  Compose Maybe ((,,,) k Int v) (UFMap k v)
-> Maybe (k, Int, v, UFMap k v)
forall {k1} {k2} (f :: k1 -> *) (g :: k2 -> k1) (a :: k2).
Compose f g a -> f (g a)
getCompose (k
-> Compose Maybe ((,,,) k Int v) (Maybe v)
-> (k -> Int -> v -> Compose Maybe ((,,,) k Int v) (UFValue k v))
-> UFMap k v
-> Compose Maybe ((,,,) k Int v) (UFMap k v)
forall (f :: * -> *) k v.
(Functor f, Ord k) =>
k
-> f (Maybe v)
-> (k -> Int -> v -> f (UFValue k v))
-> UFMap k v
-> f (UFMap k v)
alterF k
k Compose Maybe ((,,,) k Int v) (Maybe v)
forall {g :: * -> *} {a}. Compose Maybe g a
nothing k -> Int -> v -> Compose Maybe ((,,,) k Int v) (UFValue k v)
forall {a} {v} {k}.
a -> Int -> v -> Compose Maybe ((,,,) a Int v) (UFValue k v)
just UFMap k v
m)
  where
    nothing :: Compose Maybe g a
nothing = Maybe (g a) -> Compose Maybe g a
forall {k} {k1} (f :: k -> *) (g :: k1 -> k) (a :: k1).
f (g a) -> Compose f g a
Compose Maybe (g a)
forall a. Maybe a
Nothing
    just :: a -> Int -> v -> Compose Maybe ((,,,) a Int v) (UFValue k v)
just a
k Int
s v
v = Maybe (a, Int, v, UFValue k v)
-> Compose Maybe ((,,,) a Int v) (UFValue k v)
forall {k} {k1} (f :: k -> *) (g :: k1 -> k) (a :: k1).
f (g a) -> Compose f g a
Compose ((a, Int, v, UFValue k v) -> Maybe (a, Int, v, UFValue k v)
forall a. a -> Maybe a
Just (a
k, Int
s, v
v, Int -> v -> UFValue k v
forall k v. Int -> v -> UFValue k v
Canonical Int
s v
v))

data UnionHaltReason k v
  = KeyNotFound k
  | MergeFailed v v

data UnionValue k v a
  = UnionValue k Int v (UFValue k v) a
  deriving stock ((forall a b. (a -> b) -> UnionValue k v a -> UnionValue k v b)
-> (forall a b. a -> UnionValue k v b -> UnionValue k v a)
-> Functor (UnionValue k v)
forall a b. a -> UnionValue k v b -> UnionValue k v a
forall a b. (a -> b) -> UnionValue k v a -> UnionValue k v b
forall k v a b. a -> UnionValue k v b -> UnionValue k v a
forall k v a b. (a -> b) -> UnionValue k v a -> UnionValue k v 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 v a b. (a -> b) -> UnionValue k v a -> UnionValue k v b
fmap :: forall a b. (a -> b) -> UnionValue k v a -> UnionValue k v b
$c<$ :: forall k v a b. a -> UnionValue k v b -> UnionValue k v a
<$ :: forall a b. a -> UnionValue k v b -> UnionValue k v a
Functor)

union ::
  forall m k v r.
  (MonadFix m, Ord k) =>
  k ->
  k ->
  UFMap k v ->
  (UFMap k v -> m r) ->
  (k -> v -> UFMap k v -> m (Maybe r)) ->
  m (Maybe r)
union :: forall (m :: * -> *) k v r.
(MonadFix m, Ord k) =>
k
-> k
-> UFMap k v
-> (UFMap k v -> m r)
-> (k -> v -> UFMap k v -> m (Maybe r))
-> m (Maybe r)
union k
k0 k
k1 UFMap k v
mapinit UFMap k v -> m r
alreadyMerged k -> v -> UFMap k v -> m (Maybe r)
mergeValues = ExceptT (UnionHaltReason k v) m r -> m (Maybe r)
toMaybe do
  rec let lu ::
            k ->
            UFMap k v ->
            (k -> UFMap k v -> Maybe (Compose (Either (UnionHaltReason k v)) (UnionValue k v) (UFMap k v))) ->
            Either (UnionHaltReason k v) (UnionValue k v (UFMap k v))
          lu :: k
-> UFMap k v
-> (k
    -> UFMap k v
    -> Maybe
         (Compose
            (Either (UnionHaltReason k v)) (UnionValue k v) (UFMap k v)))
-> Either (UnionHaltReason k v) (UnionValue k v (UFMap k v))
lu k
k UFMap k v
m k
-> UFMap k v
-> Maybe
     (Compose
        (Either (UnionHaltReason k v)) (UnionValue k v) (UFMap k v))
loopGuard = Compose (Either (UnionHaltReason k v)) (UnionValue k v) (UFMap k v)
-> Either (UnionHaltReason k v) (UnionValue k v (UFMap k v))
forall {k1} {k2} (f :: k1 -> *) (g :: k2 -> k1) (a :: k2).
Compose f g a -> f (g a)
getCompose (k
-> (k
    -> UFMap k v
    -> Maybe
         (Compose
            (Either (UnionHaltReason k v)) (UnionValue k v) (UFMap k v)))
-> Compose
     (Either (UnionHaltReason k v)) (UnionValue k v) (Maybe v)
-> (k
    -> Int
    -> v
    -> Compose
         (Either (UnionHaltReason k v)) (UnionValue k v) (UFValue k v))
-> UFMap k v
-> Compose
     (Either (UnionHaltReason k v)) (UnionValue k v) (UFMap k v)
forall (f :: * -> *) k v.
(Functor f, Ord k) =>
k
-> (k -> UFMap k v -> Maybe (f (UFMap k v)))
-> f (Maybe v)
-> (k -> Int -> v -> f (UFValue k v))
-> UFMap k v
-> f (UFMap k v)
alterFWithHalt k
k k
-> UFMap k v
-> Maybe
     (Compose
        (Either (UnionHaltReason k v)) (UnionValue k v) (UFMap k v))
loopGuard Compose (Either (UnionHaltReason k v)) (UnionValue k v) (Maybe v)
luNothing k
-> Int
-> v
-> Compose
     (Either (UnionHaltReason k v)) (UnionValue k v) (UFValue k v)
luJust UFMap k v
m)
            where
              luNothing :: Compose (Either (UnionHaltReason k v)) (UnionValue k v) (Maybe v)
luNothing = Either (UnionHaltReason k v) (UnionValue k v (Maybe v))
-> Compose
     (Either (UnionHaltReason k v)) (UnionValue k v) (Maybe v)
forall {k} {k1} (f :: k -> *) (g :: k1 -> k) (a :: k1).
f (g a) -> Compose f g a
Compose (UnionHaltReason k v
-> Either (UnionHaltReason k v) (UnionValue k v (Maybe v))
forall a b. a -> Either a b
Left (k -> UnionHaltReason k v
forall k v. k -> UnionHaltReason k v
KeyNotFound k
k))
          luJust :: k
-> Int
-> v
-> Compose
     (Either (UnionHaltReason k v)) (UnionValue k v) (UFValue k v)
luJust k
k Int
s v
v =
            -- a final value thunk is inserted before it is resolved,
            -- as the final result cannot be known before we have
            -- looked up both values and merged them
            let newValue :: UFValue k v
newValue =
                  let newSize :: Int
newSize = case k
kcanon0 k -> k -> Bool
forall a. Eq a => a -> a -> Bool
== k
kcanon1 of
                        Bool
True -> Int
size0
                        Bool
False -> Int
size0 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
size1
                   in case k
chosenCanon k -> k -> Bool
forall a. Eq a => a -> a -> Bool
== k
k of
                        Bool
True -> Int -> v -> UFValue k v
forall k v. Int -> v -> UFValue k v
Canonical Int
newSize v
canonValue
                        Bool
False -> k -> UFValue k v
forall k v. k -> UFValue k v
Indirection k
chosenCanon
             in Either (UnionHaltReason k v) (UnionValue k v (UFValue k v))
-> Compose
     (Either (UnionHaltReason k v)) (UnionValue k v) (UFValue k v)
forall {k} {k1} (f :: k -> *) (g :: k1 -> k) (a :: k1).
f (g a) -> Compose f g a
Compose (UnionValue k v (UFValue k v)
-> Either (UnionHaltReason k v) (UnionValue k v (UFValue k v))
forall a b. b -> Either a b
Right (k
-> Int
-> v
-> UFValue k v
-> UFValue k v
-> UnionValue k v (UFValue k v)
forall k v a. k -> Int -> v -> UFValue k v -> a -> UnionValue k v a
UnionValue k
k Int
s v
v UFValue k v
newValue UFValue k v
newValue))
      UnionValue k
kcanon0 Int
size0 v
v0 UFValue k v
vfinal0 UFMap k v
map0 <- m (Either (UnionHaltReason k v) (UnionValue k v (UFMap k v)))
-> ExceptT (UnionHaltReason k v) m (UnionValue k v (UFMap k v))
forall e (m :: * -> *) a. m (Either e a) -> ExceptT e m a
ExceptT (m (Either (UnionHaltReason k v) (UnionValue k v (UFMap k v)))
 -> ExceptT (UnionHaltReason k v) m (UnionValue k v (UFMap k v)))
-> m (Either (UnionHaltReason k v) (UnionValue k v (UFMap k v)))
-> ExceptT (UnionHaltReason k v) m (UnionValue k v (UFMap k v))
forall a b. (a -> b) -> a -> b
$ Either (UnionHaltReason k v) (UnionValue k v (UFMap k v))
-> m (Either (UnionHaltReason k v) (UnionValue k v (UFMap k v)))
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either (UnionHaltReason k v) (UnionValue k v (UFMap k v))
 -> m (Either (UnionHaltReason k v) (UnionValue k v (UFMap k v))))
-> Either (UnionHaltReason k v) (UnionValue k v (UFMap k v))
-> m (Either (UnionHaltReason k v) (UnionValue k v (UFMap k v)))
forall a b. (a -> b) -> a -> b
$ k
-> UFMap k v
-> (k
    -> UFMap k v
    -> Maybe
         (Compose
            (Either (UnionHaltReason k v)) (UnionValue k v) (UFMap k v)))
-> Either (UnionHaltReason k v) (UnionValue k v (UFMap k v))
lu k
k0 UFMap k v
mapinit \k
_ UFMap k v
_ -> Maybe
  (Compose
     (Either (UnionHaltReason k v)) (UnionValue k v) (UFMap k v))
forall a. Maybe a
Nothing
      UnionValue k
kcanon1 Int
size1 v
v1 UFValue k v
vfinal1 UFMap k v
map1 <- m (Either (UnionHaltReason k v) (UnionValue k v (UFMap k v)))
-> ExceptT (UnionHaltReason k v) m (UnionValue k v (UFMap k v))
forall e (m :: * -> *) a. m (Either e a) -> ExceptT e m a
ExceptT (m (Either (UnionHaltReason k v) (UnionValue k v (UFMap k v)))
 -> ExceptT (UnionHaltReason k v) m (UnionValue k v (UFMap k v)))
-> m (Either (UnionHaltReason k v) (UnionValue k v (UFMap k v)))
-> ExceptT (UnionHaltReason k v) m (UnionValue k v (UFMap k v))
forall a b. (a -> b) -> a -> b
$
        Either (UnionHaltReason k v) (UnionValue k v (UFMap k v))
-> m (Either (UnionHaltReason k v) (UnionValue k v (UFMap k v)))
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either (UnionHaltReason k v) (UnionValue k v (UFMap k v))
 -> m (Either (UnionHaltReason k v) (UnionValue k v (UFMap k v))))
-> Either (UnionHaltReason k v) (UnionValue k v (UFMap k v))
-> m (Either (UnionHaltReason k v) (UnionValue k v (UFMap k v)))
forall a b. (a -> b) -> a -> b
$ k
-> UFMap k v
-> (k
    -> UFMap k v
    -> Maybe
         (Compose
            (Either (UnionHaltReason k v)) (UnionValue k v) (UFMap k v)))
-> Either (UnionHaltReason k v) (UnionValue k v (UFMap k v))
lu k
k1 UFMap k v
map0 \k
k UFMap k v
m -> case k
k k -> k -> Bool
forall a. Eq a => a -> a -> Bool
== k
kcanon0 of
          Bool
False -> Maybe
  (Compose
     (Either (UnionHaltReason k v)) (UnionValue k v) (UFMap k v))
forall a. Maybe a
Nothing
          Bool
True -> Compose (Either (UnionHaltReason k v)) (UnionValue k v) (UFMap k v)
-> Maybe
     (Compose
        (Either (UnionHaltReason k v)) (UnionValue k v) (UFMap k v))
forall a. a -> Maybe a
Just (Either (UnionHaltReason k v) (UnionValue k v (UFMap k v))
-> Compose
     (Either (UnionHaltReason k v)) (UnionValue k v) (UFMap k v)
forall {k} {k1} (f :: k -> *) (g :: k1 -> k) (a :: k1).
f (g a) -> Compose f g a
Compose (UnionValue k v (UFMap k v)
-> Either (UnionHaltReason k v) (UnionValue k v (UFMap k v))
forall a b. b -> Either a b
Right (k
-> Int
-> v
-> UFValue k v
-> UFMap k v
-> UnionValue k v (UFMap k v)
forall k v a. k -> Int -> v -> UFValue k v -> a -> UnionValue k v a
UnionValue k
k Int
size0 v
v0 UFValue k v
vfinal0 UFMap k v
m)))
      -- Join the smaller equivalence class to the larger to bound
      -- worst case number of lookups to log(n). This is the same
      -- strategy as the weighted fast-union algorithm.
      let (k
chosenCanon, v
canonValue, v
nonCanonValue) = case Int
size0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
size1 of
            Bool
True -> (k
kcanon0, v
v0, v
v1)
            Bool
False -> (k
kcanon1, v
v1, v
v0)
  case k
kcanon0 k -> k -> Bool
forall a. Eq a => a -> a -> Bool
== k
kcanon1 of
    Bool
True -> do
      r
res <- m r -> ExceptT (UnionHaltReason k v) m r
forall (m :: * -> *) a.
Monad m =>
m a -> ExceptT (UnionHaltReason k v) m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (UFMap k v -> m r
alreadyMerged UFMap k v
map1)
      pure (UFValue k v
vfinal0 UFValue k v -> r -> r
forall a b. a -> b -> b
`seq` r
res)
    Bool
False -> do
      r
map2 <-
        let res :: ExceptT (UnionHaltReason k v) m r
res =
              m (Either (UnionHaltReason k v) r)
-> ExceptT (UnionHaltReason k v) m r
forall e (m :: * -> *) a. m (Either e a) -> ExceptT e m a
ExceptT (m (Either (UnionHaltReason k v) r)
 -> ExceptT (UnionHaltReason k v) m r)
-> m (Either (UnionHaltReason k v) r)
-> ExceptT (UnionHaltReason k v) m r
forall a b. (a -> b) -> a -> b
$
                k -> v -> UFMap k v -> m (Maybe r)
mergeValues k
chosenCanon v
nonCanonValue UFMap k v
map1 m (Maybe r)
-> (Maybe r -> Either (UnionHaltReason k v) r)
-> m (Either (UnionHaltReason k v) r)
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \case
                  Maybe r
Nothing -> UnionHaltReason k v -> Either (UnionHaltReason k v) r
forall a b. a -> Either a b
Left (v -> v -> UnionHaltReason k v
forall k v. v -> v -> UnionHaltReason k v
MergeFailed v
v0 v
v1)
                  Just r
x -> r -> Either (UnionHaltReason k v) r
forall a b. b -> Either a b
Right r
x
         in -- Now that both lookups have completed we can safely force the
            -- final values
            UFValue k v
vfinal0 UFValue k v
-> ExceptT (UnionHaltReason k v) m r
-> ExceptT (UnionHaltReason k v) m r
forall a b. a -> b -> b
`seq` UFValue k v
vfinal1 UFValue k v
-> ExceptT (UnionHaltReason k v) m r
-> ExceptT (UnionHaltReason k v) m r
forall a b. a -> b -> b
`seq` ExceptT (UnionHaltReason k v) m r
res
      pure r
map2
  where
    toMaybe :: ExceptT (UnionHaltReason k v) m r -> m (Maybe r)
    toMaybe :: ExceptT (UnionHaltReason k v) m r -> m (Maybe r)
toMaybe (ExceptT m (Either (UnionHaltReason k v) r)
action) =
      m (Either (UnionHaltReason k v) r)
action m (Either (UnionHaltReason k v) r)
-> (Either (UnionHaltReason k v) r -> Maybe r) -> m (Maybe r)
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \case
        Right r
m -> r -> Maybe r
forall a. a -> Maybe a
Just r
m
        Left UnionHaltReason k v
r -> case UnionHaltReason k v
r of
          KeyNotFound k
_k -> Maybe r
forall a. Maybe a
Nothing
          MergeFailed v
_v0 v
_v1 -> Maybe r
forall a. Maybe a
Nothing

-- | Dump the @UFmap@ to a list grouped by equivalence class
toClasses ::
  forall k v.
  (Ord k) =>
  UFMap k v ->
  -- | [(canonical key, equivalence class, value)]
  [(k, Set k, v)]
toClasses :: forall k v. Ord k => UFMap k v -> [(k, Set k, v)]
toClasses UFMap k v
m0 =
  let cmFinal :: Map k (k, Set k, v)
      (UFMap k v
_mfinal, Map k (k, Set k, v)
cmFinal) =
        -- we fold over the UFMap's keys and build up a Map that
        -- groups the keys by equivalence class.
        ((UFMap k v, Map k (k, Set k, v))
 -> k -> (UFMap k v, Map k (k, Set k, v)))
-> (UFMap k v, Map k (k, Set k, v))
-> [k]
-> (UFMap k v, Map k (k, Set k, v))
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (UFMap k v, Map k (k, Set k, v))
-> k -> (UFMap k v, Map k (k, Set k, v))
forall {a} {c}.
Ord a =>
(UFMap a c, Map a (a, Set a, c))
-> a -> (UFMap a c, Map a (a, Set a, c))
buildCmFinal (UFMap k v
m0, Map k (k, Set k, v)
forall k a. Map k a
Map.empty) [k]
keys
      keys :: [k]
keys = case UFMap k v
m0 of
        UFMap Map k (UFValue k v)
m -> Map k (UFValue k v) -> [k]
forall k a. Map k a -> [k]
Map.keys Map k (UFValue k v)
m
      buildCmFinal :: (UFMap a c, Map a (a, Set a, c))
-> a -> (UFMap a c, Map a (a, Set a, c))
buildCmFinal (UFMap a c
m, Map a (a, Set a, c)
cm) a
k =
        let (a
kcanon, Int
_, c
v, UFMap a c
m') = Maybe (a, Int, c, UFMap a c) -> (a, Int, c, UFMap a c)
forall a. HasCallStack => Maybe a -> a
fromJust (a -> UFMap a c -> Maybe (a, Int, c, UFMap a c)
forall k v. Ord k => k -> UFMap k v -> Maybe (k, Int, v, UFMap k v)
lookupCanon a
k UFMap a c
m)
            cm' :: Map a (a, Set a, c)
cm' =
              ((a, Set a, c) -> (a, Set a, c) -> (a, Set a, c))
-> a -> (a, Set a, c) -> Map a (a, Set a, c) -> Map a (a, Set a, c)
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
Map.insertWith
                (\(a
k0, Set a
s0, c
v0) (a
_k1, Set a
s1, c
_v1) -> (a
k0, Set a
s0 Set a -> Set a -> Set a
forall a. Semigroup a => a -> a -> a
<> Set a
s1, c
v0))
                a
kcanon
                (a
k, a -> Set a
forall a. a -> Set a
Set.singleton a
k, c
v)
                Map a (a, Set a, c)
cm
         in (UFMap a c
m', Map a (a, Set a, c)
cm')
   in Map k (k, Set k, v) -> [(k, Set k, v)]
forall k a. Map k a -> [a]
Map.elems Map k (k, Set k, v)
cmFinal

keys :: UFMap k v -> [k]
keys :: forall k v. UFMap k v -> [k]
keys (UFMap Map k (UFValue k v)
m) = Map k (UFValue k v) -> [k]
forall k a. Map k a -> [k]
Map.keys Map k (UFValue k v)
m