packages feed

union-find 0.1 → 0.2

raw patch · 5 files changed

+359/−4 lines, 5 filesdep +containersdep +transformersdep ~base

Dependencies added: containers, transformers

Dependency ranges changed: base

Files

README.md view
@@ -1,2 +1,25 @@-union-find-========+# union-find++A simple Haskell library that implements Tarjan's Union/Find+algorithm.  Useful, for example, to implement unification in a type+inference system.++The Union/Find algorithm implements these operations in+(effectively) constant-time:++ 1. Check whether two elements are in the same equivalence class.++ 2. Create a union of two equivalence classes.++ 3. Look up the descriptor of the equivalence class.+++## Installation++Using `cabal` (which comes with the Haskell Platform):++    $ cabal install union-find++or in the checked-out repository:++    $ cabal install
+ src/Control/Monad/Trans/UnionFind.hs view
@@ -0,0 +1,59 @@+{-# LANGUAGE GeneralizedNewtypeDeriving #-}+module Control.Monad.Trans.UnionFind+  ( UnionFindT, runUnionFind+  , Point, fresh, repr, descriptor, union, equivalent+  ) where++import Control.Applicative (Applicative)+import Control.Monad.Trans.Class (MonadTrans(..))+import Control.Monad.Trans.State (StateT(..), evalStateT)+import Data.UnionFind.IntMap (Point)+import qualified Control.Monad.Trans.State as State+import qualified Data.UnionFind.IntMap as UF++-- | A monad transformer that adds union find operations.+--+-- The @p@ parameter is the type of points.  Uses the+-- "Data.UnionFind.IntMap" as the underlying union-find+-- implementation.+newtype UnionFindT p m a = UnionFindT {+  unUnionFindT :: StateT (UF.PointSupply p) m a+  } deriving (Functor, Applicative, Monad, MonadTrans)++runUnionFind :: Monad m => UnionFindT p m a -> m a+runUnionFind = (`evalStateT` UF.newPointSupply) . unUnionFindT++swap :: (a, b) -> (b, a)+swap (x, y) = (y, x)++-- | Create a new point with the given descriptor.  The returned is+-- only equivalent to itself.+--+-- Note that a 'Point' has its own identity.  That is, if two points+-- are equivalent then their descriptors are equal, but not vice+-- versa.+--+fresh :: Monad m => p -> UnionFindT p m (Point p)+fresh x = UnionFindT . StateT $ return . swap . flip UF.fresh x++-- | /O(1)/. @repr point@ returns the representative point of+-- @point@'s equivalence class.+repr :: Monad m => Point p -> UnionFindT p m (Point p)+repr = UnionFindT . State.gets . flip UF.repr++-- | Return the descriptor of the +descriptor :: Monad m => Point p -> UnionFindT p m p+descriptor = UnionFindT . State.gets . flip UF.descriptor++-- | Join the equivalence classes of the points.  The resulting+-- equivalence class will get the descriptor of the second argument.+union :: Monad m => Point p -> Point p -> UnionFindT p m ()+union p1 p2 = UnionFindT . State.modify $ \x -> UF.union x p1 p2++-- | Test if the two elements are in the same equivalence class.+-- +-- @+-- liftA2 (==) (repr x) (repr y)+-- @+equivalent :: Monad m => Point p -> Point p -> UnionFindT p m Bool+equivalent p1 p2 = UnionFindT . State.gets $ \x -> UF.equivalent x p1 p2
+ src/Data/UnionFind/IntMap.hs view
@@ -0,0 +1,101 @@+{-# LANGUAGE Rank2Types #-}+{-# LANGUAGE BangPatterns #-}+{-# OPTIONS_GHC -funbox-strict-fields #-}+--+-- NOTE: Does not currently ensure that a 'Point' was indeed generated+-- by the specified 'PointSupply'.+--+module Data.UnionFind.IntMap +    ( newPointSupply, fresh, repr, descriptor, union, equivalent,+      PointSupply, Point ) where++import qualified Data.IntMap as IM++data PointSupply a = PointSupply !Int (IM.IntMap (Link a))+  deriving Show++data Link a +    = Info {-# UNPACK #-} !Int a+      -- ^ This is the descriptive element of the equivalence class+      -- and its rank.+    | Link {-# UNPACK #-} !Int+      -- ^ Pointer to some other element of the equivalence class.+     deriving Show++newtype Point a = Point Int++newPointSupply :: PointSupply a+newPointSupply = PointSupply 0 IM.empty++fresh :: PointSupply a -> a -> (PointSupply a, Point a)+fresh (PointSupply next eqs) a =+  (PointSupply (next + 1) (IM.insert next (Info 0 a) eqs), Point next)++-- freshList :: PointSupply a -> [a] -> (PointSupply a, [Point a])+-- freshList ++repr :: PointSupply a -> Point a -> Point a+repr ps p = reprInfo ps p (\n _rank _a -> Point n)++reprInfo :: PointSupply a -> Point a -> (Int -> Int -> a -> r) -> r+reprInfo (PointSupply _next eqs) (Point n) k = go n+  where+    go !i =+      case eqs IM.! i of+        Link i' -> go i'+        Info r a -> k i r a+  +union :: PointSupply a -> Point a -> Point a -> PointSupply a+union ps@(PointSupply next eqs) p1 p2 =+  reprInfo ps p1 $ \i1 r1 _a1 -> +  reprInfo ps p2 $ \i2 r2 a2 ->+  if i1 == i2 then ps else+    case r1 `compare` r2 of+      LT ->+        -- No rank or descriptor update necessary+        let !eqs1 = IM.insert i1 (Link i2) eqs in+        PointSupply next eqs1+      EQ ->+        let !eqs1 = IM.insert i1 (Link i2) eqs+            !eqs2 = IM.insert i2 (Info (r2 + 1) a2) eqs1 in+        PointSupply next eqs2+      GT ->+        let !eqs1 = IM.insert i1 (Info r2 a2) eqs+            !eqs2 = IM.insert i2 (Link i1) eqs1 in+        PointSupply next eqs2++descriptor :: PointSupply a -> Point a -> a+descriptor ps p = reprInfo ps p (\_ _ a -> a)++equivalent :: PointSupply a -> Point a -> Point a -> Bool+equivalent ps p1 p2 =+  reprInfo ps p1 $ \i1 _ _ ->+  reprInfo ps p2 $ \i2 _ _ ->+  i1 == i2++{-+tst1 :: IO ()+tst1 = do+  let ps0 = newPointSupply+      (ps1, p1) = fresh ps0 "hello"+      (ps2, p2) = fresh ps1 "world"+      (ps3, p3) = fresh ps2 "you"+      (ps, p4) = fresh ps3 "there"+  let ps' = union ps p1 p2+  print (descr ps p1, descr ps p2, equivalent ps p1 p2)+  print (descr ps' p1, descr ps' p2, equivalent ps' p1 p2)+  let ps'' = union ps' p3 p1+  print (descr ps'' p1, descr ps'' p3, equivalent ps'' p1 p3)+  print ps''+-}++-- TODO: should fail+{-+tst2 :: IO ()+tst2 = do+  let as0 = newPointSupply+      (as, a1) = fresh as0 "foo"+      bs0 = newPointSupply+      (bs, b1) = fresh bs0 "bar"+  print $ union as a1 b1+  -}    
+ src/Data/UnionFind/ST.hs view
@@ -0,0 +1,169 @@+-- | An implementation of Tarjan's UNION-FIND algorithm.  (Robert E+-- Tarjan. \"Efficiency of a Good But Not Linear Set Union Algorithm\", JACM+-- 22(2), 1975)+--+-- The algorithm implements three operations efficiently (all amortised+-- @O(1)@):+--+--  1. Check whether two elements are in the same equivalence class.+--+--  2. Create a union of two equivalence classes.+--+--  3. Look up the descriptor of the equivalence class.+-- +-- The implementation is based on mutable references.  Each+-- equivalence class has exactly one member that serves as its+-- representative element.  Every element either is the representative+-- element of its equivalence class or points to another element in+-- the same equivalence class.  Equivalence testing thus consists of+-- following the pointers to the representative elements and then+-- comparing these for identity.+--+-- The algorithm performs lazy path compression.  That is, whenever we+-- walk along a path greater than length 1 we automatically update the+-- pointers along the path to directly point to the representative+-- element.  Consequently future lookups will be have a path length of+-- at most 1.+--+{-# OPTIONS_GHC -funbox-strict-fields #-}+module Data.UnionFind.ST+  ( Point, fresh, repr, union, union', equivalent, redundant,+    descriptor, setDescriptor, modifyDescriptor )+where++import Control.Applicative+import Control.Monad ( when )+import Control.Monad.ST+import Data.STRef++-- | The abstract type of an element of the sets we work on.  It is+-- parameterised over the type of the descriptor.+newtype Point s a = Pt (STRef s (Link s a)) deriving Eq++data Link s a+    = Info {-# UNPACK #-} !(STRef s (Info a))+      -- ^ This is the descriptive element of the equivalence class.+    | Link {-# UNPACK #-} !(Point s a)+      -- ^ Pointer to some other element of the equivalence class.+     deriving Eq++data Info a = MkInfo+  { weight :: {-# UNPACK #-} !Int+    -- ^ The size of the equivalence class, used by 'union'.+  , descr  :: a+  } deriving Eq++-- | /O(1)/. Create a fresh point and return it.  A fresh point is in+-- the equivalence class that contains only itself.+fresh :: a -> ST s (Point s a)+fresh desc = do+  info <- newSTRef (MkInfo { weight = 1, descr = desc })+  l <- newSTRef (Info info)+  return (Pt l)++-- | /O(1)/. @repr point@ returns the representative point of+-- @point@'s equivalence class.+--+-- This method performs the path compresssion.+repr :: Point s a -> ST s (Point s a)+repr point@(Pt l) = do+  link <- readSTRef l+  case link of+    Info _ -> return point+    Link pt'@(Pt l') -> do+      pt'' <- repr pt'+      when (pt'' /= pt') $ do+        -- At this point we know that @pt'@ is not the representative+        -- element of @point@'s equivalent class.  Therefore @pt'@'s+        -- link must be of the form @Link r@.  We write this same+        -- value into @point@'s link reference and thereby perform+        -- path compression.+        link' <- readSTRef l'+        writeSTRef l link'+      return pt''++-- | Return the reference to the point's equivalence class's+-- descriptor.+descrRef :: Point s a -> ST s (STRef s (Info a))+descrRef point@(Pt link_ref) = do+  link <- readSTRef link_ref+  case link of+    Info info -> return info+    Link (Pt link'_ref) -> do+      link' <- readSTRef link'_ref+      case link' of+        Info info -> return info+        _ -> descrRef =<< repr point++-- | /O(1)/. Return the descriptor associated with argument point's+-- equivalence class.+descriptor :: Point s a -> ST s a+descriptor point = do+  descr <$> (readSTRef =<< descrRef point)++-- | /O(1)/. Replace the descriptor of the point's equivalence class+-- with the second argument.+setDescriptor :: Point s a -> a -> ST s ()+setDescriptor point new_descr = do+  r <- descrRef point+  modifySTRef r $ \i -> i { descr = new_descr }++modifyDescriptor :: Point s a -> (a -> a) -> ST s ()+modifyDescriptor point f = do+  r <- descrRef point+  modifySTRef r $ \i -> i { descr = f (descr i) }++-- | /O(1)/. Join the equivalence classes of the points (which must be+-- distinct).  The resulting equivalence class will get the descriptor+-- of the second argument.+union :: Point s a -> Point s a -> ST s ()+union p1 p2 = union' p1 p2 (\_ d2 -> return d2)++-- | Like 'union', but sets the descriptor returned from the callback.+-- +-- The intention is to keep the descriptor of the second argument to+-- the callback, but the callback might adjust the information of the+-- descriptor or perform side effects.+union' :: Point s a -> Point s a -> (a -> a -> ST s a) -> ST s ()+union' p1 p2 update = do+  point1@(Pt link_ref1) <- repr p1+  point2@(Pt link_ref2) <- repr p2+  -- The precondition ensures that we don't create cyclic structures.+  when (point1 /= point2) $ do+    Info info_ref1 <- readSTRef link_ref1+    Info info_ref2 <- readSTRef link_ref2+    MkInfo w1 d1 <- readSTRef info_ref1 -- d1 is discarded+    MkInfo w2 d2 <- readSTRef info_ref2+    d2' <- update d1 d2+    -- Make the smaller tree a a subtree of the bigger one.  The idea+    -- is this: We increase the path length of one set by one.+    -- Assuming all elements are accessed equally often, this means+    -- the penalty is smaller if we do it for the smaller set of the+    -- two.+    if w1 >= w2 then do+      writeSTRef link_ref2 (Link point1)+      writeSTRef info_ref1 (MkInfo (w1 + w2) d2')+     else do+      writeSTRef link_ref1 (Link point2)+      writeSTRef info_ref2 (MkInfo (w1 + w2) d2')++-- | /O(1)/. Return @True@ if both points belong to the same+-- | equivalence class.+equivalent :: Point s a -> Point s a -> ST s Bool+equivalent p1 p2 = (==) <$> repr p1 <*> repr p2++-- | /O(1)/. Returns @True@ for all but one element of an equivalence+-- class.  That is, if @ps = [p1, .., pn]@ are all in the same+-- equivalence class, then the following assertion holds.+-- +-- > do rs <- mapM redundant ps+-- >    assert (length (filter (==False) rs) == 1)+-- +-- It is unspecified for which element function returns @False@, so be+-- really careful when using this.+redundant :: Point s a -> ST s Bool+redundant (Pt link_r) = do+  link <- readSTRef link_r+  case link of+    Info _ -> return False+    Link _ -> return True
union-find.cabal view
@@ -1,5 +1,5 @@ Name:            union-find-Version:         0.1+Version:         0.2 License:         BSD3 License-File:    LICENSE Author:          Thomas Schilling <nominolo@googlemail.com>@@ -28,10 +28,13 @@  Library   Build-Depends:-    base >= 4 && < 5+    base >= 4 && < 5, containers >= 0.3, transformers >= 0.2   GHC-Options:     -Wall   Exposed-Modules:+    Control.Monad.Trans.UnionFind     Data.UnionFind.IO+    Data.UnionFind.ST+    Data.UnionFind.IntMap   Hs-Source-Dirs: src