git.iko.soy/elk
on ca2be8a help sign in
elk/src/Type/UnionFind.hs
172 lines 3.8 kB history raw
1{-# OPTIONS_GHC -funbox-strict-fields #-}
2{-# LANGUAGE BangPatterns #-}
3module Type.UnionFind
4 ( Point
5 , fresh
6 , union
7 , equivalent
8 , redundant
9 , get
10 , set
11 , modify
12 )
13 where
16{- This is based on the following implementations:
18 - https://hackage.haskell.org/package/union-find-0.2/docs/src/Data-UnionFind-IO.html
19 - http://yann.regis-gianas.org/public/mini/code_UnionFind.html
21It seems like the OCaml one came first, but I am not sure.
23Compared to the Haskell implementation, the major changes here include:
25 1. No more reallocating PointInfo when changing the weight
26 2. Using the strict modifyIORef
28-}
31import Control.Monad ( when )
32import Data.IORef (IORef, modifyIORef', newIORef, readIORef, writeIORef)
33import Data.Word (Word32)
37-- POINT
40newtype Point a =
41 Pt (IORef (PointInfo a))
42 deriving Eq
45data PointInfo a
46 = Info {-# UNPACK #-} !(IORef Word32) {-# UNPACK #-} !(IORef a)
47 | Link {-# UNPACK #-} !(Point a)
51-- HELPERS
54fresh :: a -> IO (Point a)
55fresh value =
56 do weight <- newIORef 1
57 desc <- newIORef value
58 link <- newIORef (Info weight desc)
59 return (Pt link)
62repr :: Point a -> IO (Point a)
63repr point@(Pt ref) =
64 do pInfo <- readIORef ref
65 case pInfo of
66 Info _ _ ->
67 return point
69 Link point1@(Pt ref1) ->
70 do point2 <- repr point1
71 when (point2 /= point1) $
72 do pInfo1 <- readIORef ref1
73 writeIORef ref pInfo1
74 return point2
77get :: Point a -> IO a
78get point@(Pt ref) =
79 do pInfo <- readIORef ref
80 case pInfo of
81 Info _ descRef ->
82 readIORef descRef
84 Link (Pt ref1) ->
85 do link' <- readIORef ref1
86 case link' of
87 Info _ descRef ->
88 readIORef descRef
90 Link _ ->
91 get =<< repr point
94set :: Point a -> a -> IO ()
95set point@(Pt ref) newDesc =
96 do pInfo <- readIORef ref
97 case pInfo of
98 Info _ descRef ->
99 writeIORef descRef newDesc
101 Link (Pt ref1) ->
102 do link' <- readIORef ref1
103 case link' of
104 Info _ descRef ->
105 writeIORef descRef newDesc
107 Link _ ->
108 do newPoint <- repr point
109 set newPoint newDesc
112modify :: Point a -> (a -> a) -> IO ()
113modify point@(Pt ref) func =
114 do pInfo <- readIORef ref
115 case pInfo of
116 Info _ descRef ->
117 modifyIORef' descRef func
119 Link (Pt ref1) ->
120 do link' <- readIORef ref1
121 case link' of
122 Info _ descRef ->
123 modifyIORef' descRef func
125 Link _ ->
126 do newPoint <- repr point
127 modify newPoint func
130union :: Point a -> Point a -> a -> IO ()
131union p1 p2 newDesc =
132 do point1@(Pt ref1) <- repr p1
133 point2@(Pt ref2) <- repr p2
135 Info w1 d1 <- readIORef ref1
136 Info w2 d2 <- readIORef ref2
138 if point1 == point2
139 then writeIORef d1 newDesc
140 else do
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