git.iko.soy/elk
on e00a3f3 help sign in
elk/src/Type/UnionFind.hs
172 lines 3.9 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 -- assert (point1 /= point2)
138 Info w1 d1 <- readIORef ref1
139 Info w2 d2 <- readIORef ref2
141 weight1 <- readIORef w1
142 weight2 <- readIORef w2
144 let !newWeight = weight1 + weight2
146 if weight1 >= weight2
147 then
148 do writeIORef ref2 (Link point1)
149 writeIORef w1 newWeight
150 writeIORef d1 newDesc
151 else
152 do writeIORef ref1 (Link point2)
153 writeIORef w2 newWeight
154 writeIORef d2 newDesc
157equivalent :: Point a -> Point a -> IO Bool
158equivalent p1 p2 =
159 do v1 <- repr p1
160 v2 <- repr p2
161 return (v1 == v2)
164redundant :: Point a -> IO Bool
165redundant (Pt ref) =
166 do pInfo <- readIORef ref
167 case pInfo of
168 Info _ _ ->
169 return False
171 Link _ ->
172 return True