Union find with groups! (completely untested)

Group.hs

module Data.UnionFind.Group where

import Control.Monad
import Control.Monad.Primitive
import Data.Group
import Data.Primitive

data UFG' m g = UFG'
  { operator :: g
  , root :: UFG m g
  } deriving Eq

type UFG m g = MutVar (PrimState m) (Either Word (UFG' m g))

newUFG :: PrimMonad m => m (UFG m g)
newUFG = newMutVar $ Left 0

find :: (PrimMonad m, Monoid g) => UFG m g -> m (g, UFG m g)
find ref = do
  deref <- readMutVar ref
  case deref of
    Left _ -> return (mempty, ref)
    Right UFG' {root = next, operator = operator} -> do
      (operator', root) <- find next
      let operator'' = operator <> operator'
      writeMutVar ref $ Right UFG' { operator = operator'', root = root}
      return (operator'', root)

-- 'MonadPlus' constraint because we check that two elements that are 'union'ed
-- which already have the same route also have the same operator.
--
-- Group argument should go from l to r
union :: (PrimMonad m, MonadPlus m, Eq g, Group g) => UFG m g -> UFG m g -> g -> m ()
union l r o = do
  (lo, lr) <- find l
  (ro, rr) <- find r
  if (lr == rr) 
  -- they are already united but we check for consistency
  then guard (lo == ro)
  -- they are not already united, so we don't need to check for consistency
  else do
    -- these we know should be Left unless there's some ghosts or concurrency
    lrank <- readMutVar lr >>= either return (const mzero)
    rrank <- readMutVar rr >>= either return (const mzero)
    case lrank `compare` rrank of
      LT -> writeMutVar lr $ Right UFG' { operator = invert o, root = rr }
      GT -> writeMutVar rr $ Right UFG' { operator =        o, root = lr }
      EQ -> do
        -- arbitrarily set the left as root, mostly so we don't need to compute
        -- inverse, in case that's expensive
        writeMutVar lr $ Left (lrank + 1)
        writeMutVar rr $ Right UFG' { operator = o, root = lr}