Transaction.hs

-- | A sealed mutable execution mode for the canonical persistent union-find.
-- Dense non-negative prefixes use mutable parent/rank arrays with in-place path
-- compression; sparse, negative, and giant keys stay in an 'IntMap' overlay.
-- Successful transactions freeze back to the immutable 'UnionFind' owner.
--
-- This module is the public face: the two transaction runners (the only
-- @runST@ seals) and the query/mutation verbs. The mutable arena lives in
-- sealed @Transaction.Internal.*@ owners; 'UnionFindEditor' is re-exported
-- abstractly so no external caller can forge or dissect editor state.
module Moonlight.Core.UnionFind.Transaction
  ( UnionFindEditor,
    UnionOutcome (..),
    UnionFindAllocationError (..),
    runUnionFindTransaction,
    runUnionFindTransactionEither,
    transactionMember,
    transactionFind,
    transactionFindExisting,
    transactionCanonicalClass,
    transactionInsertClassId,
    transactionMakeSet,
    transactionUnion,
    transactionEquivalent,
    transactionCanonicalMapAndCompress,
  )
where

import Control.Monad.ST (ST, runST)
import Data.IntMap.Strict (IntMap)
import Data.IntMap.Strict qualified as IntMap
import Data.Maybe (isJust)
import Data.STRef (modifySTRef', readSTRef, writeSTRef)
import Moonlight.Core.Identifier.EGraph (ClassId (..), classIdKey)
import Moonlight.Core.UnionFind.Internal.Semantics (chooseLink, rootKeyBy)
import Moonlight.Core.UnionFind.Internal.Types
  ( UnionFind,
    UnionFindAllocationError (..),
    advanceNextFreshForClassIdKey,
    allocateNextClassId,
  )
import Moonlight.Core.UnionFind.Transaction.Internal.Algorithm
  ( applyLinkDecision,
    compressRootKey,
  )
import Moonlight.Core.UnionFind.Transaction.Internal.Lifecycle
  ( freezeUnionFind,
    thawUnionFind,
  )
import Moonlight.Core.UnionFind.Transaction.Internal.Snapshot
  ( parentMap,
  )
import Moonlight.Core.UnionFind.Transaction.Internal.Storage
  ( readParentKey,
    readRankValue,
    writeFreshClassIdentity,
  )
import Moonlight.Core.UnionFind.Transaction.Internal.Types
  ( UnionFindEditor (nextFresh),
    UnionOutcome (..),
  )
import Prelude

runUnionFindTransaction ::
  UnionFind ->
  (forall state. UnionFindEditor state -> ST state value) ->
  (value, UnionFind)
runUnionFindTransaction base action =
  runST $ do
    editor <- thawUnionFind base
    value <- action editor
    committed <- freezeUnionFind editor
    pure (value, committed)

runUnionFindTransactionEither ::
  UnionFind ->
  (forall state. UnionFindEditor state -> ST state (Either err value)) ->
  Either err (value, UnionFind)
runUnionFindTransactionEither base action =
  runST $ do
    editor <- thawUnionFind base
    outcome <- action editor
    case outcome of
      Left err ->
        pure (Left err)
      Right value -> do
        committed <- freezeUnionFind editor
        pure (Right (value, committed))

transactionMember :: UnionFindEditor state -> ClassId -> ST state Bool
transactionMember editor classId =
  fmap isJust (readParentKey editor (classIdKey classId))
{-# INLINE transactionMember #-}

transactionFind :: UnionFindEditor state -> ClassId -> ST state ClassId
transactionFind editor classId =
  fmap ClassId (compressRootKey editor (classIdKey classId))
{-# INLINE transactionFind #-}

transactionFindExisting ::
  UnionFindEditor state ->
  ClassId ->
  ST state (Maybe ClassId)
transactionFindExisting editor classId = do
  exists <- transactionMember editor classId
  if exists
    then fmap Just (transactionFind editor classId)
    else pure Nothing
{-# INLINE transactionFindExisting #-}

transactionCanonicalClass ::
  UnionFindEditor state ->
  ClassId ->
  ST state (Maybe ClassId)
transactionCanonicalClass editor classId = do
  exists <- transactionMember editor classId
  if exists
    then fmap (Just . ClassId) (rootKeyBy (readParentKey editor) (classIdKey classId))
    else pure Nothing
{-# INLINE transactionCanonicalClass #-}

transactionInsertClassId :: UnionFindEditor state -> ClassId -> ST state ()
transactionInsertClassId editor classId = do
  let key = classIdKey classId
  modifySTRef' (nextFresh editor) (advanceNextFreshForClassIdKey key)
  exists <- fmap isJust (readParentKey editor key)
  if exists
    then pure ()
    else writeFreshClassIdentity editor key
{-# INLINE transactionInsertClassId #-}

transactionMakeSet :: UnionFindEditor state -> ST state (Either UnionFindAllocationError ClassId)
transactionMakeSet editor = do
  currentNextFresh <- readSTRef (nextFresh editor)
  case allocateNextClassId currentNextFresh of
    Left allocationError ->
      pure (Left allocationError)
    Right (classId, updatedNextFresh) -> do
      writeFreshClassIdentity editor (classIdKey classId)
      writeSTRef (nextFresh editor) updatedNextFresh
      pure (Right classId)
{-# INLINE transactionMakeSet #-}

transactionUnion ::
  UnionFindEditor state ->
  ClassId ->
  ClassId ->
  ST state UnionOutcome
transactionUnion editor leftClassId rightClassId = do
  transactionInsertClassId editor leftClassId
  transactionInsertClassId editor rightClassId
  leftRoot <- transactionFind editor leftClassId
  rightRoot <- transactionFind editor rightClassId
  if leftRoot == rightRoot
    then pure (AlreadyEquivalent leftRoot)
    else do
      leftRank <- readRankValue editor (classIdKey leftRoot)
      rightRank <- readRankValue editor (classIdKey rightRoot)
      applyLinkDecision editor (chooseLink leftRoot leftRank rightRoot rightRank)
{-# INLINE transactionUnion #-}

transactionEquivalent ::
  UnionFindEditor state ->
  ClassId ->
  ClassId ->
  ST state Bool
transactionEquivalent editor leftClassId rightClassId = do
  leftRoot <- rootKeyBy (readParentKey editor) (classIdKey leftClassId)
  rightRoot <- rootKeyBy (readParentKey editor) (classIdKey rightClassId)
  pure (leftRoot == rightRoot)
{-# INLINE transactionEquivalent #-}

transactionCanonicalMapAndCompress ::
  UnionFindEditor state ->
  ST state (IntMap ClassId)
transactionCanonicalMapAndCompress editor = do
  currentParents <- parentMap editor
  IntMap.traverseWithKey
    (\key _ -> transactionFind editor (ClassId key))
    currentParents