{-# LANGUAGE DeriveAnyClass #-}

-- | See 'Assoc'.
module Overeasy.Assoc
  ( Assoc
  , assocFwd
  , assocBwd
  , assocEquiv
  , assocSize
  , assocNew
  , assocSingleton
  , AssocInsertRes (..)
  , assocInsertInc
  , assocInsert
  , assocFromList
  , assocToList
  , assocMember
  , assocLookupByKey
  , assocPartialLookupByKey
  , assocLookupByValue
  , assocPartialLookupByValue
  , assocLookupRoot
  , assocRoots
  , assocLeaves
  , assocMembers
  , assocCanCompact
  , assocCompactInc
  , assocCompact
  , assocRemoveAllInc
  , assocRemoveAll
  , assocUnion
  , assocFootprint
  ) where

import Control.DeepSeq (NFData)
import Control.Monad.State.Strict (MonadState (..), State, modify')
import Data.Coerce (Coercible)
import Data.Foldable (foldl')
import Data.Hashable (Hashable)
import Data.HashMap.Strict (HashMap)
import qualified Data.HashMap.Strict as HashMap
import Data.Maybe (fromJust)
import GHC.Generics (Generic)
import IntLike.Map (IntLikeMap)
import qualified IntLike.Map as ILM
import IntLike.Set (IntLikeSet)
import qualified IntLike.Set as ILS
import Overeasy.EquivFind (EquivAddRes (..), EquivFind, efAddInc, efBwd, efCanCompact, efCompactInc, efEquivs, efLeaves,
                           efLookupRoot, efMember, efMembers, efNew, efRemoveAllInc, efRoots, efSingleton,
                           efUnsafeAddLeafInc, efUnsafeMerge)

-- | Associates keys and values in such a way that inserting
-- duplicate values induces equivalences on their keys.
-- Invariant: fwd and bwd maps contain only root keys.
data Assoc x a = Assoc
  { forall x a. Assoc x a -> IntLikeMap x a
assocFwd :: !(IntLikeMap x a)
  -- ^ Map from id to element
  , forall x a. Assoc x a -> HashMap a x
assocBwd :: !(HashMap a x)
  -- ^ Map from element to id
  , forall x a. Assoc x a -> EquivFind x
assocEquiv :: !(EquivFind x)
  -- ^ Equivalence classes of ids
  } deriving stock (Assoc x a -> Assoc x a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall x a. (Eq a, Eq x) => Assoc x a -> Assoc x a -> Bool
/= :: Assoc x a -> Assoc x a -> Bool
$c/= :: forall x a. (Eq a, Eq x) => Assoc x a -> Assoc x a -> Bool
== :: Assoc x a -> Assoc x a -> Bool
$c== :: forall x a. (Eq a, Eq x) => Assoc x a -> Assoc x a -> Bool
Eq, Int -> Assoc x a -> ShowS
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall x a. (Show a, Show x) => Int -> Assoc x a -> ShowS
forall x a. (Show a, Show x) => [Assoc x a] -> ShowS
forall x a. (Show a, Show x) => Assoc x a -> String
showList :: [Assoc x a] -> ShowS
$cshowList :: forall x a. (Show a, Show x) => [Assoc x a] -> ShowS
show :: Assoc x a -> String
$cshow :: forall x a. (Show a, Show x) => Assoc x a -> String
showsPrec :: Int -> Assoc x a -> ShowS
$cshowsPrec :: forall x a. (Show a, Show x) => Int -> Assoc x a -> ShowS
Show, forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall x a x. Rep (Assoc x a) x -> Assoc x a
forall x a x. Assoc x a -> Rep (Assoc x a) x
$cto :: forall x a x. Rep (Assoc x a) x -> Assoc x a
$cfrom :: forall x a x. Assoc x a -> Rep (Assoc x a) x
Generic)
    deriving anyclass (forall a. (a -> ()) -> NFData a
forall x a. (NFData a, NFData x) => Assoc x a -> ()
rnf :: Assoc x a -> ()
$crnf :: forall x a. (NFData a, NFData x) => Assoc x a -> ()
NFData)

-- | How many values are in the map?
assocSize :: Assoc x a -> Int
assocSize :: forall x a. Assoc x a -> Int
assocSize = forall x a. IntLikeMap x a -> Int
ILM.size forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall x a. Assoc x a -> IntLikeMap x a
assocFwd

-- | Creates an empty assoc
assocNew :: Assoc x a
assocNew :: forall x a. Assoc x a
assocNew = forall x a.
IntLikeMap x a -> HashMap a x -> EquivFind x -> Assoc x a
Assoc forall x a. IntLikeMap x a
ILM.empty forall k v. HashMap k v
HashMap.empty forall x. EquivFind x
efNew

-- | Creates a singleton assoc
assocSingleton :: (Coercible x Int, Hashable a) => x -> a -> Assoc x a
assocSingleton :: forall x a. (Coercible x Int, Hashable a) => x -> a -> Assoc x a
assocSingleton x
x a
a = forall x a.
IntLikeMap x a -> HashMap a x -> EquivFind x -> Assoc x a
Assoc (forall x a. Coercible x Int => x -> a -> IntLikeMap x a
ILM.singleton x
x a
a) (forall k v. Hashable k => k -> v -> HashMap k v
HashMap.singleton a
a x
x) (forall x. Coercible x Int => x -> EquivFind x
efSingleton x
x)

-- | The result of inserting into the assoc, if you're interested.
data AssocInsertRes x =
    AssocInsertResUnchanged
  | AssocInsertResCreated
  | AssocInsertResUpdated
  | AssocInsertResMerged !(IntLikeSet x)
  deriving stock (AssocInsertRes x -> AssocInsertRes x -> Bool
forall x. AssocInsertRes x -> AssocInsertRes x -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: AssocInsertRes x -> AssocInsertRes x -> Bool
$c/= :: forall x. AssocInsertRes x -> AssocInsertRes x -> Bool
== :: AssocInsertRes x -> AssocInsertRes x -> Bool
$c== :: forall x. AssocInsertRes x -> AssocInsertRes x -> Bool
Eq, Int -> AssocInsertRes x -> ShowS
forall x. Int -> AssocInsertRes x -> ShowS
forall x. [AssocInsertRes x] -> ShowS
forall x. AssocInsertRes x -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [AssocInsertRes x] -> ShowS
$cshowList :: forall x. [AssocInsertRes x] -> ShowS
show :: AssocInsertRes x -> String
$cshow :: forall x. AssocInsertRes x -> String
showsPrec :: Int -> AssocInsertRes x -> ShowS
$cshowsPrec :: forall x. Int -> AssocInsertRes x -> ShowS
Show)

-- | Insert into the assoc (raw version)
assocInsertInc :: (Coercible x Int, Ord x, Eq a, Hashable a) => x -> a -> Assoc x a -> ((x, AssocInsertRes x), Assoc x a)
assocInsertInc :: forall x a.
(Coercible x Int, Ord x, Eq a, Hashable a) =>
x -> a -> Assoc x a -> ((x, AssocInsertRes x), Assoc x a)
assocInsertInc x
x a
a1 assoc :: Assoc x a
assoc@(Assoc IntLikeMap x a
fwd HashMap a x
bwd EquivFind x
equiv) = ((x, AssocInsertRes x), Assoc x a)
finalRes where
  finalRes :: ((x, AssocInsertRes x), Assoc x a)
finalRes =
    let (EquivAddRes x
res, EquivFind x
equiv') = forall x.
Coercible x Int =>
x -> EquivFind x -> (EquivAddRes x, EquivFind x)
efAddInc x
x EquivFind x
equiv
    in case EquivAddRes x
res of
      EquivAddRes x
EquivAddResNewRoot -> x -> EquivFind x -> ((x, AssocInsertRes x), Assoc x a)
insertRoot x
x EquivFind x
equiv'
      EquivAddResAlreadyLeafOf x
z -> x -> ((x, AssocInsertRes x), Assoc x a)
updateRoot x
z
      EquivAddRes x
EquivAddResAlreadyRoot -> x -> ((x, AssocInsertRes x), Assoc x a)
updateRoot x
x
  updateRoot :: x -> ((x, AssocInsertRes x), Assoc x a)
updateRoot x
w =
    -- w is existing root and is guaranteed to map to something
    let a0 :: a
a0 = forall x a. Coercible x Int => x -> IntLikeMap x a -> a
ILM.partialLookup x
w IntLikeMap x a
fwd
    in if a
a0 forall a. Eq a => a -> a -> Bool
== a
a1
      -- the value has not changed, don't need to change assoc
      then ((x
w, forall x. AssocInsertRes x
AssocInsertResUnchanged), Assoc x a
assoc)
      else
        -- value has changed, need to check if it's fresh
        case forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HashMap.lookup a
a1 HashMap a x
bwd of
          -- never seen; insert and return
          Maybe x
Nothing ->
            let fwd' :: IntLikeMap x a
fwd' = forall x a.
Coercible x Int =>
x -> a -> IntLikeMap x a -> IntLikeMap x a
ILM.insert x
w a
a1 IntLikeMap x a
fwd
                bwd' :: HashMap a x
bwd' = forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
HashMap.insert a
a1 x
w (forall k v. (Eq k, Hashable k) => k -> HashMap k v -> HashMap k v
HashMap.delete a
a0 HashMap a x
bwd)
            in ((x
w, forall x. AssocInsertRes x
AssocInsertResUpdated), forall x a.
IntLikeMap x a -> HashMap a x -> EquivFind x -> Assoc x a
Assoc IntLikeMap x a
fwd' HashMap a x
bwd' EquivFind x
equiv)
          -- mapped to another set of nodes, merge
          Just x
v ->
            let (x
toKeep, IntLikeSet x
toDelete, EquivFind x
equiv') = forall x.
(Coercible x Int, Ord x) =>
x -> x -> EquivFind x -> (x, IntLikeSet x, EquivFind x)
efUnsafeMerge x
w x
v EquivFind x
equiv
                res :: AssocInsertRes x
res = forall x. IntLikeSet x -> AssocInsertRes x
AssocInsertResMerged IntLikeSet x
toDelete
            in if x
toKeep forall a. Eq a => a -> a -> Bool
== x
w
              -- w wins
              then
                let fwd' :: IntLikeMap x a
fwd' = forall x a.
Coercible x Int =>
x -> a -> IntLikeMap x a -> IntLikeMap x a
ILM.insert x
w a
a1 (forall x a.
Coercible x Int =>
x -> IntLikeMap x a -> IntLikeMap x a
ILM.delete x
v IntLikeMap x a
fwd)
                    bwd' :: HashMap a x
bwd' = forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
HashMap.insert a
a1 x
w (forall k v. (Eq k, Hashable k) => k -> HashMap k v -> HashMap k v
HashMap.delete a
a0 HashMap a x
bwd)
                in ((x
w, AssocInsertRes x
res), forall x a.
IntLikeMap x a -> HashMap a x -> EquivFind x -> Assoc x a
Assoc IntLikeMap x a
fwd' HashMap a x
bwd' EquivFind x
equiv')
              -- v wins
              else
                let fwd' :: IntLikeMap x a
fwd' = forall x a.
Coercible x Int =>
x -> IntLikeMap x a -> IntLikeMap x a
ILM.delete x
w IntLikeMap x a
fwd
                    bwd' :: HashMap a x
bwd' = forall k v. (Eq k, Hashable k) => k -> HashMap k v -> HashMap k v
HashMap.delete a
a0 HashMap a x
bwd
                in ((x
v, AssocInsertRes x
res), forall x a.
IntLikeMap x a -> HashMap a x -> EquivFind x -> Assoc x a
Assoc IntLikeMap x a
fwd' HashMap a x
bwd' EquivFind x
equiv')
  insertRoot :: x -> EquivFind x -> ((x, AssocInsertRes x), Assoc x a)
insertRoot x
w EquivFind x
equiv' =
    -- w is new root that doesn't exist
    case forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HashMap.lookup a
a1 HashMap a x
bwd of
      -- never seen; insert and return
      Maybe x
Nothing ->
        let fwd' :: IntLikeMap x a
fwd' = forall x a.
Coercible x Int =>
x -> a -> IntLikeMap x a -> IntLikeMap x a
ILM.insert x
w a
a1 IntLikeMap x a
fwd
            bwd' :: HashMap a x
bwd' = forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
HashMap.insert a
a1 x
w HashMap a x
bwd
        in ((x
w, forall x. AssocInsertRes x
AssocInsertResCreated), forall x a.
IntLikeMap x a -> HashMap a x -> EquivFind x -> Assoc x a
Assoc IntLikeMap x a
fwd' HashMap a x
bwd' EquivFind x
equiv')
      Just x
v ->
        let (x
toKeep, IntLikeSet x
toDelete, EquivFind x
equiv'') = forall x.
(Coercible x Int, Ord x) =>
x -> x -> EquivFind x -> (x, IntLikeSet x, EquivFind x)
efUnsafeMerge x
w x
v EquivFind x
equiv'
            res :: AssocInsertRes x
res = forall x. IntLikeSet x -> AssocInsertRes x
AssocInsertResMerged IntLikeSet x
toDelete
        in if x
toKeep forall a. Eq a => a -> a -> Bool
== x
w
          -- w wins
          then
            let fwd' :: IntLikeMap x a
fwd' = forall x a.
Coercible x Int =>
x -> a -> IntLikeMap x a -> IntLikeMap x a
ILM.insert x
w a
a1 (forall x a.
Coercible x Int =>
x -> IntLikeMap x a -> IntLikeMap x a
ILM.delete x
v IntLikeMap x a
fwd)
                bwd' :: HashMap a x
bwd' = forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
HashMap.insert a
a1 x
w HashMap a x
bwd
            in ((x
w, AssocInsertRes x
res), forall x a.
IntLikeMap x a -> HashMap a x -> EquivFind x -> Assoc x a
Assoc IntLikeMap x a
fwd' HashMap a x
bwd' EquivFind x
equiv'')
          -- v wins
          else
            let fwd' :: IntLikeMap x a
fwd' = forall x a.
Coercible x Int =>
x -> IntLikeMap x a -> IntLikeMap x a
ILM.delete x
w IntLikeMap x a
fwd
            in ((x
v, AssocInsertRes x
res), forall x a.
IntLikeMap x a -> HashMap a x -> EquivFind x -> Assoc x a
Assoc IntLikeMap x a
fwd' HashMap a x
bwd EquivFind x
equiv'')

-- | Insert into the assoc (the state version)
assocInsert :: (Coercible x Int, Ord x, Eq a, Hashable a) => x -> a -> State (Assoc x a) (x, AssocInsertRes x)
assocInsert :: forall x a.
(Coercible x Int, Ord x, Eq a, Hashable a) =>
x -> a -> State (Assoc x a) (x, AssocInsertRes x)
assocInsert x
x a
a = forall s (m :: * -> *) a. MonadState s m => (s -> (a, s)) -> m a
state (forall x a.
(Coercible x Int, Ord x, Eq a, Hashable a) =>
x -> a -> Assoc x a -> ((x, AssocInsertRes x), Assoc x a)
assocInsertInc x
x a
a)

-- | Build an assoc from a list of pairs
assocFromList :: (Coercible x Int, Ord x, Eq a, Hashable a) => [(x, a)] -> Assoc x a
assocFromList :: forall x a.
(Coercible x Int, Ord x, Eq a, Hashable a) =>
[(x, a)] -> Assoc x a
assocFromList = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\Assoc x a
assoc (x
x, a
a) -> forall a b. (a, b) -> b
snd (forall x a.
(Coercible x Int, Ord x, Eq a, Hashable a) =>
x -> a -> Assoc x a -> ((x, AssocInsertRes x), Assoc x a)
assocInsertInc x
x a
a Assoc x a
assoc)) forall x a. Assoc x a
assocNew

-- | Turn an assoc into a list of pairs (NOTE - emits only root keys!)
assocToList :: Coercible x Int => Assoc x a -> [(x, a)]
assocToList :: forall x a. Coercible x Int => Assoc x a -> [(x, a)]
assocToList = forall x a. Coercible x Int => IntLikeMap x a -> [(x, a)]
ILM.toList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall x a. Assoc x a -> IntLikeMap x a
assocFwd

-- | Is the given key in the assoc?
assocMember :: Coercible x Int => x -> Assoc x a -> Bool
assocMember :: forall x a. Coercible x Int => x -> Assoc x a -> Bool
assocMember x
x (Assoc IntLikeMap x a
_ HashMap a x
_ EquivFind x
equiv) = forall x. Coercible x Int => x -> EquivFind x -> Bool
efMember x
x EquivFind x
equiv

-- | Forward lookup
assocLookupByKey :: Coercible x Int => x -> Assoc x a -> Maybe a
assocLookupByKey :: forall x a. Coercible x Int => x -> Assoc x a -> Maybe a
assocLookupByKey x
x (Assoc IntLikeMap x a
fwd HashMap a x
_ EquivFind x
equiv) = forall x a. Coercible x Int => x -> IntLikeMap x a -> Maybe a
ILM.lookup (forall x. Coercible x Int => x -> EquivFind x -> x
efLookupRoot x
x EquivFind x
equiv) IntLikeMap x a
fwd

-- | PARTIAL forward lookup
assocPartialLookupByKey :: Coercible x Int => x -> Assoc x a -> a
assocPartialLookupByKey :: forall x a. Coercible x Int => x -> Assoc x a -> a
assocPartialLookupByKey x
x = forall a. HasCallStack => Maybe a -> a
fromJust forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall x a. Coercible x Int => x -> Assoc x a -> Maybe a
assocLookupByKey x
x

-- | Backward lookup
assocLookupByValue :: (Eq a, Hashable a) => a -> Assoc x a -> Maybe x
assocLookupByValue :: forall a x. (Eq a, Hashable a) => a -> Assoc x a -> Maybe x
assocLookupByValue a
a = forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HashMap.lookup a
a forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall x a. Assoc x a -> HashMap a x
assocBwd

-- | PARTIAL backward lookup
assocPartialLookupByValue :: (Eq a, Hashable a) => a -> Assoc x a -> x
assocPartialLookupByValue :: forall a x. (Eq a, Hashable a) => a -> Assoc x a -> x
assocPartialLookupByValue a
a = forall a b c. (a -> b -> c) -> b -> a -> c
flip forall k v.
(Eq k, Hashable k, HasCallStack) =>
HashMap k v -> k -> v
(HashMap.!) a
a forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall x a. Assoc x a -> HashMap a x
assocBwd

-- | Finds the root for the given key (id if not found)
assocLookupRoot :: Coercible x Int => x -> Assoc x a -> x
assocLookupRoot :: forall x a. Coercible x Int => x -> Assoc x a -> x
assocLookupRoot x
x = forall x. Coercible x Int => x -> EquivFind x -> x
efLookupRoot x
x forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall x a. Assoc x a -> EquivFind x
assocEquiv

-- | List all root (live, non-compactible) keys
assocRoots :: Coercible x Int => Assoc x a -> [x]
assocRoots :: forall x a. Coercible x Int => Assoc x a -> [x]
assocRoots = forall x. Coercible x Int => EquivFind x -> [x]
efRoots forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall x a. Assoc x a -> EquivFind x
assocEquiv

-- | List all leaf (dead, compactible) keys
assocLeaves :: Coercible x Int => Assoc x a -> [x]
assocLeaves :: forall x a. Coercible x Int => Assoc x a -> [x]
assocLeaves = forall x. Coercible x Int => EquivFind x -> [x]
efLeaves forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall x a. Assoc x a -> EquivFind x
assocEquiv

-- | List all entries (root and leaf)
assocMembers :: Coercible x Int => Assoc x a -> [x]
assocMembers :: forall x a. Coercible x Int => Assoc x a -> [x]
assocMembers = forall x. Coercible x Int => EquivFind x -> [x]
efMembers forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall x a. Assoc x a -> EquivFind x
assocEquiv

-- | Are there dead keys in the equiv from 'assocInsert'?
assocCanCompact :: Assoc x a -> Bool
assocCanCompact :: forall x a. Assoc x a -> Bool
assocCanCompact = forall x. EquivFind x -> Bool
efCanCompact forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall x a. Assoc x a -> EquivFind x
assocEquiv

-- | Removes all dead keys in the equiv (raw version).
assocCompactInc :: Coercible x Int => Assoc x a -> (IntLikeMap x x, Assoc x a)
assocCompactInc :: forall x a.
Coercible x Int =>
Assoc x a -> (IntLikeMap x x, Assoc x a)
assocCompactInc assoc :: Assoc x a
assoc@(Assoc IntLikeMap x a
fwd HashMap a x
bwd EquivFind x
equiv) =
  let replacements :: IntLikeMap x x
replacements = forall x. EquivFind x -> IntLikeMap x x
efBwd EquivFind x
equiv
      assoc' :: Assoc x a
assoc' =
        if forall x a. IntLikeMap x a -> Bool
ILM.null IntLikeMap x x
replacements
          then Assoc x a
assoc
          else let (IntLikeMap x (IntLikeSet x)
_, EquivFind x
equiv') = forall x.
Coercible x Int =>
EquivFind x -> (IntLikeMap x (IntLikeSet x), EquivFind x)
efCompactInc EquivFind x
equiv in forall x a.
IntLikeMap x a -> HashMap a x -> EquivFind x -> Assoc x a
Assoc IntLikeMap x a
fwd HashMap a x
bwd EquivFind x
equiv'
  in (IntLikeMap x x
replacements, Assoc x a
assoc')

-- | Removes all dead keys in the equiv (state version).
-- Returns map of dead leaf node -> live root node
assocCompact :: Coercible x Int => State (Assoc x a) (IntLikeMap x x)
assocCompact :: forall x a. Coercible x Int => State (Assoc x a) (IntLikeMap x x)
assocCompact = forall s (m :: * -> *) a. MonadState s m => (s -> (a, s)) -> m a
state forall x a.
Coercible x Int =>
Assoc x a -> (IntLikeMap x x, Assoc x a)
assocCompactInc

-- | Removes the given keys from the assoc (raw version)
assocRemoveAllInc :: (Coercible x Int, Eq a, Hashable a) => [x] -> Assoc x a -> Assoc x a
assocRemoveAllInc :: forall x a.
(Coercible x Int, Eq a, Hashable a) =>
[x] -> Assoc x a -> Assoc x a
assocRemoveAllInc [x]
xs (Assoc IntLikeMap x a
fwd0 HashMap a x
bwd0 EquivFind x
equiv0) = forall x a.
IntLikeMap x a -> HashMap a x -> EquivFind x -> Assoc x a
Assoc IntLikeMap x a
fwdFinal HashMap a x
bwdFinal EquivFind x
equivFinal where
  (IntLikeMap x x
remap, EquivFind x
equivFinal) = forall x.
Coercible x Int =>
[x] -> EquivFind x -> (IntLikeMap x x, EquivFind x)
efRemoveAllInc [x]
xs EquivFind x
equiv0
  (IntLikeMap x a
fwdFinal, HashMap a x
bwdFinal) = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (IntLikeMap x a, HashMap a x) -> x -> (IntLikeMap x a, HashMap a x)
go (IntLikeMap x a
fwd0, HashMap a x
bwd0) [x]
xs
  go :: (IntLikeMap x a, HashMap a x) -> x -> (IntLikeMap x a, HashMap a x)
go tup :: (IntLikeMap x a, HashMap a x)
tup@(IntLikeMap x a
fwd, HashMap a x
bwd) x
x =
    case forall x a. Coercible x Int => x -> IntLikeMap x a -> Maybe a
ILM.lookup x
x IntLikeMap x a
fwd of
      -- Leaf, ignore
      Maybe a
Nothing -> (IntLikeMap x a, HashMap a x)
tup
      -- Root
      Just a
a ->
        case forall x a. Coercible x Int => x -> IntLikeMap x a -> Maybe a
ILM.lookup x
x IntLikeMap x x
remap of
          -- Singleton root, delete
          Maybe x
Nothing ->
            let fwd' :: IntLikeMap x a
fwd' = forall x a.
Coercible x Int =>
x -> IntLikeMap x a -> IntLikeMap x a
ILM.delete x
x IntLikeMap x a
fwd
                bwd' :: HashMap a x
bwd' = forall k v. (Eq k, Hashable k) => k -> HashMap k v -> HashMap k v
HashMap.delete a
a HashMap a x
bwd
            in (IntLikeMap x a
fwd', HashMap a x
bwd')
          -- Remapped root, rotate
          Just x
y ->
            let fwd' :: IntLikeMap x a
fwd' = forall x a.
Coercible x Int =>
x -> IntLikeMap x a -> IntLikeMap x a
ILM.delete x
x (forall x a.
Coercible x Int =>
x -> a -> IntLikeMap x a -> IntLikeMap x a
ILM.insert x
y a
a IntLikeMap x a
fwd)
                bwd' :: HashMap a x
bwd' = forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
HashMap.insert a
a x
y HashMap a x
bwd
            in (IntLikeMap x a
fwd', HashMap a x
bwd')

-- | Removes the given keys from the assoc (state version).
-- Values will only be removed from the assoc if the key is a singleton root.
-- If a key is not found, it is simply ignored.
assocRemoveAll :: (Coercible x Int, Eq a, Hashable a) => [x] -> State (Assoc x a) ()
assocRemoveAll :: forall x a.
(Coercible x Int, Eq a, Hashable a) =>
[x] -> State (Assoc x a) ()
assocRemoveAll = forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify' forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall x a.
(Coercible x Int, Eq a, Hashable a) =>
[x] -> Assoc x a -> Assoc x a
assocRemoveAllInc

-- | Join two assocs (uses the first as the base)
assocUnion :: (Coercible x Int, Ord x, Eq a, Hashable a) => Assoc x a -> Assoc x a -> Assoc x a
assocUnion :: forall x a.
(Coercible x Int, Ord x, Eq a, Hashable a) =>
Assoc x a -> Assoc x a -> Assoc x a
assocUnion Assoc x a
base (Assoc IntLikeMap x a
fwd HashMap a x
_ EquivFind x
equiv) = forall x a.
IntLikeMap x a -> HashMap a x -> EquivFind x -> Assoc x a
Assoc IntLikeMap x a
fwdFinal HashMap a x
bwdFinal EquivFind x
equivFinal where
  goRoots :: Assoc x a -> (x, a) -> Assoc x a
goRoots Assoc x a
assocGo (x
x, a
a) = forall a b. (a, b) -> b
snd (forall x a.
(Coercible x Int, Ord x, Eq a, Hashable a) =>
x -> a -> Assoc x a -> ((x, AssocInsertRes x), Assoc x a)
assocInsertInc x
x a
a Assoc x a
assocGo)
  goLeaves :: EquivFind x -> (x, x) -> EquivFind x
goLeaves EquivFind x
equivGo (x
leaf, x
oldRoot) = forall x. Coercible x Int => x -> x -> EquivFind x -> EquivFind x
efUnsafeAddLeafInc x
oldRoot x
leaf EquivFind x
equivGo
  Assoc IntLikeMap x a
fwdFinal HashMap a x
bwdFinal EquivFind x
equivMid = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' forall {x} {a}.
(Coercible x Int, Ord x, Hashable a) =>
Assoc x a -> (x, a) -> Assoc x a
goRoots Assoc x a
base (forall x a. Coercible x Int => IntLikeMap x a -> [(x, a)]
ILM.toList IntLikeMap x a
fwd)
  equivFinal :: EquivFind x
equivFinal = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' forall {x}. Coercible x Int => EquivFind x -> (x, x) -> EquivFind x
goLeaves EquivFind x
equivMid (forall x a. Coercible x Int => IntLikeMap x a -> [(x, a)]
ILM.toList (forall x. EquivFind x -> IntLikeMap x x
efBwd EquivFind x
equiv))

-- | Returns the footprint of the given value - all keys that map to it (root and leaf)
assocFootprint :: (Coercible x Int, Eq a, Hashable a) => a -> Assoc x a -> IntLikeSet x
assocFootprint :: forall x a.
(Coercible x Int, Eq a, Hashable a) =>
a -> Assoc x a -> IntLikeSet x
assocFootprint a
a (Assoc IntLikeMap x a
_ HashMap a x
bwd EquivFind x
equiv) =
  case forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HashMap.lookup a
a HashMap a x
bwd of
    Maybe x
Nothing -> forall x. IntLikeSet x
ILS.empty
    Just x
r -> forall x. Coercible x Int => x -> EquivFind x -> IntLikeSet x
efEquivs x
r EquivFind x
equiv