never executed always true always false
1 {- Union-find data structure compiled from Distribution.Utils.UnionFind -}
2 module GHC.Data.UnionFind where
3
4 import GHC.Prelude
5 import Data.STRef
6 import Control.Monad.ST
7 import Control.Monad
8
9 -- | A variable which can be unified; alternately, this can be thought
10 -- of as an equivalence class with a distinguished representative.
11 newtype Point s a = Point (STRef s (Link s a))
12 deriving (Eq)
13
14 -- | Mutable write to a 'Point'
15 writePoint :: Point s a -> Link s a -> ST s ()
16 writePoint (Point v) = writeSTRef v
17
18 -- | Read the current value of 'Point'.
19 readPoint :: Point s a -> ST s (Link s a)
20 readPoint (Point v) = readSTRef v
21
22 -- | The internal data structure for a 'Point', which either records
23 -- the representative element of an equivalence class, or a link to
24 -- the 'Point' that actually stores the representative type.
25 data Link s a
26 -- NB: it is too bad we can't say STRef Int#; the weights remain boxed
27 = Info {-# UNPACK #-} !(STRef s Int) {-# UNPACK #-} !(STRef s a)
28 | Link {-# UNPACK #-} !(Point s a)
29
30 -- | Create a fresh equivalence class with one element.
31 fresh :: a -> ST s (Point s a)
32 fresh desc = do
33 weight <- newSTRef 1
34 descriptor <- newSTRef desc
35 Point `fmap` newSTRef (Info weight descriptor)
36
37 -- | Flatten any chains of links, returning a 'Point'
38 -- which points directly to the canonical representation.
39 repr :: Point s a -> ST s (Point s a)
40 repr point = readPoint point >>= \r ->
41 case r of
42 Link point' -> do
43 point'' <- repr point'
44 when (point'' /= point') $ do
45 writePoint point =<< readPoint point'
46 return point''
47 Info _ _ -> return point
48
49 -- | Return the canonical element of an equivalence
50 -- class 'Point'.
51 find :: Point s a -> ST s a
52 find point =
53 -- Optimize length 0 and 1 case at expense of
54 -- general case
55 readPoint point >>= \r ->
56 case r of
57 Info _ d_ref -> readSTRef d_ref
58 Link point' -> readPoint point' >>= \r' ->
59 case r' of
60 Info _ d_ref -> readSTRef d_ref
61 Link _ -> repr point >>= find
62
63 -- | Unify two equivalence classes, so that they share
64 -- a canonical element. Keeps the descriptor of point2.
65 union :: Point s a -> Point s a -> ST s ()
66 union refpoint1 refpoint2 = do
67 point1 <- repr refpoint1
68 point2 <- repr refpoint2
69 when (point1 /= point2) $ do
70 l1 <- readPoint point1
71 l2 <- readPoint point2
72 case (l1, l2) of
73 (Info wref1 dref1, Info wref2 dref2) -> do
74 weight1 <- readSTRef wref1
75 weight2 <- readSTRef wref2
76 -- Should be able to optimize the == case separately
77 if weight1 >= weight2
78 then do
79 writePoint point2 (Link point1)
80 -- The weight calculation here seems a bit dodgy
81 writeSTRef wref1 (weight1 + weight2)
82 writeSTRef dref1 =<< readSTRef dref2
83 else do
84 writePoint point1 (Link point2)
85 writeSTRef wref2 (weight1 + weight2)
86 _ -> error "UnionFind.union: repr invariant broken"
87
88 -- | Test if two points are in the same equivalence class.
89 equivalent :: Point s a -> Point s a -> ST s Bool
90 equivalent point1 point2 = liftM2 (==) (repr point1) (repr point2)
91