git.iko.soy/elk
on ea77dcf help sign in
elk/src/Type/UnionFind.hs
173 lines 4.0 kB history raw
1{-# OPTIONS_GHC -funbox-strict-fields #-}
2{-# LANGUAGE BangPatterns #-}
3module Type.UnionFind
4 ( Point
5 , fresh
6 , repr
7 , union
8 , equivalent
9 , redundant
10 , descriptor
11 , setDescriptor
12 , modifyDescriptor
13 )
14 where
17{- This is based on the following implementations:
19 - https://hackage.haskell.org/package/union-find-0.2/docs/src/Data-UnionFind-IO.html
20 - http://yann.regis-gianas.org/public/mini/code_UnionFind.html
22It seems like the OCaml one came first, but I am not sure.
24Compared to the Haskell implementation, the major changes here include:
26 1. No more reallocating PointInfo when changing the weight
27 2. Using the strict modifyIORef
29-}
32import Control.Monad ( when )
33import Data.IORef (IORef, modifyIORef', newIORef, readIORef, writeIORef)
34import Data.Word (Word32)
38-- POINT
41newtype Point a =
42 Pt (IORef (PointInfo a))
43 deriving Eq
46data PointInfo a
47 = Info {-# UNPACK #-} !(IORef Word32) {-# UNPACK #-} !(IORef a)
48 | Link {-# UNPACK #-} !(Point a)
52-- HELPERS
55fresh :: a -> IO (Point a)
56fresh value =
57 do weight <- newIORef 1
58 desc <- newIORef value
59 link <- newIORef (Info weight desc)
60 return (Pt link)
63repr :: Point a -> IO (Point a)
64repr point@(Pt ref) =
65 do pInfo <- readIORef ref
66 case pInfo of
67 Info _ _ ->
68 return point
70 Link point1@(Pt ref1) ->
71 do point2 <- repr point1
72 when (point2 /= point1) $
73 do pInfo1 <- readIORef ref1
74 writeIORef ref pInfo1
75 return point2
78descriptor :: Point a -> IO a
79descriptor point@(Pt ref) =
80 do pInfo <- readIORef ref
81 case pInfo of
82 Info _ descRef ->
83 readIORef descRef
85 Link (Pt ref1) ->
86 do link' <- readIORef ref1
87 case link' of
88 Info _ descRef ->
89 readIORef descRef
91 Link _ ->
92 descriptor =<< repr point
95setDescriptor :: Point a -> a -> IO ()
96setDescriptor point@(Pt ref) newDesc =
97 do pInfo <- readIORef ref
98 case pInfo of
99 Info _ descRef ->
100 writeIORef descRef newDesc
102 Link (Pt ref1) ->
103 do link' <- readIORef ref1
104 case link' of
105 Info _ descRef ->
106 writeIORef descRef newDesc
108 Link _ ->
109 do newPoint <- repr point
110 setDescriptor newPoint newDesc
113modifyDescriptor :: Point a -> (a -> a) -> IO ()
114modifyDescriptor point@(Pt ref) func =
115 do pInfo <- readIORef ref
116 case pInfo of
117 Info _ descRef ->
118 modifyIORef' descRef func
120 Link (Pt ref1) ->
121 do link' <- readIORef ref1
122 case link' of
123 Info _ descRef ->
124 modifyIORef' descRef func
126 Link _ ->
127 do newPoint <- repr point
128 modifyDescriptor newPoint func
131union :: Point a -> Point a -> a -> IO ()
132union p1 p2 newDesc =
133 do point1@(Pt ref1) <- repr p1
134 point2@(Pt ref2) <- repr p2
136 Info w1 d1 <- readIORef ref1
137 Info w2 d2 <- readIORef ref2
139 if point1 == point2
140 then writeIORef d1 newDesc
141 else do
142 weight1 <- readIORef w1
143 weight2 <- readIORef w2
145 let !newWeight = weight1 + weight2
147 if weight1 >= weight2
148 then
149 do writeIORef ref2 (Link point1)
150 writeIORef w1 newWeight
151 writeIORef d1 newDesc
152 else
153 do writeIORef ref1 (Link point2)
154 writeIORef w2 newWeight
155 writeIORef d2 newDesc
158equivalent :: Point a -> Point a -> IO Bool
159equivalent p1 p2 =
160 do v1 <- repr p1
161 v2 <- repr p2
162 return (v1 == v2)
165redundant :: Point a -> IO Bool
166redundant (Pt ref) =
167 do pInfo <- readIORef ref
168 case pInfo of
169 Info _ _ ->
170 return False
172 Link _ ->
173 return True