src/Clustering/UnionFind/ST.hs

   1 -- SPDX-License-Identifier: BSD-3-Clause
   2 -- SPDX-FileCopyrightText: 2012 Thomas Schilling <nominolo@googlemail.com>
   3 --
   4 {-# OPTIONS_GHC -funbox-strict-fields #-}
   5
   6 -- | An implementation of Tarjan's UNION-FIND algorithm.
   7 -- (Robert E Tarjan. \"Efficiency of a Good But Not Linear Set Union Algorithm\", JACM 22(2), 1975)
   8 --
   9 -- The algorithm implements three operations efficiently (all amortised @O(1)@):
  10 --
  11 --  1. Check whether two elements are in the same equivalence class.
  12 --
  13 --  2. Create a union of two equivalence classes.
  14 --
  15 --  3. Look up the descriptor of the equivalence class.
  16 --
  17 -- The implementation is based on mutable references.
  18 -- Each equivalence class has exactly one member that serves
  19 -- as its representative element.
  20 -- Every element either is the representative element of its equivalence class
  21 -- or points to another element in the same equivalence class.
  22 -- Equivalence testing thus consists of following the pointers
  23 -- to the representative elements and then comparing these for identity.
  24 --
  25 -- The algorithm performs lazy path compression.
  26 -- That is, whenever we walk along a path greater than length 1
  27 -- we automatically update the pointers along the path to directly point
  28 -- to the representative element.
  29 -- Consequently future lookups will be have a path length of at most 1.
  30 --
  31 -- Adapted from Thomas Schilling's union-find package:
  32 -- https://hackage.haskell.org/package/union-find
  33 module Clustering.UnionFind.ST (
  34   Point,
  35   fresh,
  36   repr,
  37   union,
  38   unionWith,
  39   equivalent,
  40   redundant,
  41   descriptor,
  42   setDescriptor,
  43   modifyDescriptor,
  44 )
  45 where
  46
  47 import Control.Applicative
  48 import Control.Monad (Monad (..), when)
  49 import Control.Monad.ST
  50 import Data.Bool (Bool (..))
  51 import Data.Eq (Eq (..))
  52 import Data.Function (($))
  53 import Data.Int (Int)
  54 import Data.Ord (Ord (..))
  55 import Data.STRef
  56 import Prelude (error, (+))
  57
  58 -- | The abstract type of an element of the sets we work on.  It is
  59 -- parameterised over the type of the descriptor.
  60 newtype Point s a = MkPoint (STRef s (Link s a))
  61   deriving Eq -- Pointer equality on STRef
  62
  63 -- TODO: unpack Info
  64 -- as in https://github.com/haskell/cabal/blob/8815e0a3e76e05cac91b8a88ce7d590afb07ef71/Cabal/src/Distribution/Utils/UnionFind.hs
  65 data Link s a
  66   = -- | This is the descriptive element of the equivalence class.
  67     Info {-# UNPACK #-} !(STRef s (Info a))
  68   | -- | Pointer to some other element of the equivalence class.
  69     Link {-# UNPACK #-} !(Point s a)
  70   deriving (Eq)
  71
  72 unInfo :: Link s a -> STRef s (Info a)
  73 unInfo = \case
  74   Info x -> x
  75   _ -> error "unInfo"
  76
  77 data Info a = MkInfo
  78   { weight :: {-# UNPACK #-} !Int
  79   -- ^ The size of the equivalence class, used by 'union'.
  80   , descr :: a
  81   }
  82   deriving (Eq)
  83
  84 -- | /O(1)/.
  85 -- Create a fresh equivalence class and return it.  A fresh point is in
  86 -- the equivalence class that contains only itself.
  87 fresh :: a -> ST s (Point s a)
  88 fresh desc = do
  89   info <- newSTRef (MkInfo{weight = 1, descr = desc})
  90   l <- newSTRef (Info info)
  91   return (MkPoint l)
  92
  93 -- | /O(1)/. @repr point@
  94 -- returns the representative point of @point@'s equivalence class.
  95 --
  96 -- This method performs the path compresssion.
  97 repr :: Point s a -> ST s (Point s a)
  98 repr point@(MkPoint l) = do
  99   link <- readSTRef l
 100   case link of
 101     Info _ -> return point
 102     Link pt'@(MkPoint l') -> do
 103       pt'' <- repr pt'
 104       when (pt'' /= pt') $ do
 105         -- At this point we know that @pt'@ is not the representative
 106         -- element of @point@'s equivalent class.  Therefore @pt'@'s
 107         -- link must be of the form @Link r@.  We write this same
 108         -- value into @point@'s link reference and thereby perform
 109         -- path compression.
 110         link' <- readSTRef l'
 111         writeSTRef l link'
 112       return pt''
 113
 114 -- | Return the reference to the point's equivalence class's descriptor.
 115 descrRef :: Point s a -> ST s (STRef s (Info a))
 116 descrRef point@(MkPoint link_ref) = do
 117   link <- readSTRef link_ref
 118   case link of
 119     Info info -> return info
 120     Link (MkPoint link'_ref) -> do
 121       -- Unrolling for the length == 1 case.
 122       link' <- readSTRef link'_ref
 123       case link' of
 124         Info info -> return info
 125         _ -> repr point >>= descrRef
 126
 127 -- | /O(1)/. Return the descriptor associated with argument point's
 128 -- equivalence class.
 129 descriptor :: Point s a -> ST s a
 130 descriptor point = descr <$> (descrRef point >>= readSTRef)
 131
 132 -- | /O(1)/. Replace the descriptor of the point's equivalence class
 133 -- with the second argument.
 134 setDescriptor :: Point s a -> a -> ST s ()
 135 setDescriptor point new_descr = do
 136   r <- descrRef point
 137   modifySTRef r $ \i -> i{descr = new_descr}
 138
 139 modifyDescriptor :: Point s a -> (a -> a) -> ST s ()
 140 modifyDescriptor point f = do
 141   r <- descrRef point
 142   modifySTRef r $ \i -> i{descr = f (descr i)}
 143
 144 -- | /O(1)/. Join the equivalence classes of the points (which must be
 145 -- distinct).  The resulting equivalence class will get the descriptor
 146 -- of the second argument.
 147 union :: Point s a -> Point s a -> ST s ()
 148 union p1 p2 = unionWith p1 p2 (\_ d2 -> return d2)
 149
 150 -- | Like 'union', but sets the descriptor returned from the callback.
 151 --
 152 -- The intention is to keep the descriptor of the second argument to
 153 -- the callback, but the callback might adjust the information of the
 154 -- descriptor or perform side effects.
 155 unionWith :: Point s a -> Point s a -> (a -> a -> ST s a) -> ST s ()
 156 unionWith p1 p2 update = do
 157   point1@(MkPoint link_ref1) <- repr p1
 158   point2@(MkPoint link_ref2) <- repr p2
 159   -- The precondition ensures that we don't create cyclic structures.
 160   when (point1 /= point2) $ do
 161     info_ref1 <- unInfo <$> readSTRef link_ref1
 162     info_ref2 <- unInfo <$> readSTRef link_ref2
 163     MkInfo w1 d1 <- readSTRef info_ref1 -- d1 is discarded
 164     MkInfo w2 d2 <- readSTRef info_ref2
 165     d2' <- update d1 d2
 166     -- Make the smaller tree a subtree of the bigger one.
 167     -- The idea is this: We increase the path length of one set by one.
 168     -- Assuming all elements are accessed equally often,
 169     -- this means the penalty is smaller if we do it
 170     -- for the smaller set of the two.
 171     if w1 >= w2
 172       then do
 173         writeSTRef link_ref2 (Link point1)
 174         writeSTRef info_ref1 (MkInfo (w1 + w2) d2')
 175       else do
 176         writeSTRef link_ref1 (Link point2)
 177         writeSTRef info_ref2 (MkInfo (w1 + w2) d2')
 178
 179 -- | /O(1)/. Return @True@ if both points belong to the same
 180 -- | equivalence class.
 181 equivalent :: Point s a -> Point s a -> ST s Bool
 182 equivalent p1 p2 = (==) <$> repr p1 <*> repr p2
 183
 184 -- | /O(1)/. Returns @True@ for all but one element of an equivalence class.
 185 -- That is, if @ps = [p1, .., pn]@ are all in the same
 186 -- equivalence class, then the following assertion holds.
 187 --
 188 -- > do rs <- mapM redundant ps
 189 -- >    assert (length (filter (==False) rs) == 1)
 190 --
 191 -- It is unspecified for which element function returns @False@,
 192 -- so be really careful when using this.
 193 redundant :: Point s a -> ST s Bool
 194 redundant (MkPoint link_r) = do
 195   link <- readSTRef link_r
 196   case link of
 197     Info _ -> return False
 198     Link _ -> return True