@@ -0,0 +1,172 @@
1+{-# OPTIONS_GHC -funbox-strict-fields #-}
2+{-# LANGUAGE BangPatterns #-}
3+module Type.UnionFind
4+ ( Point
5+ , fresh
6+ , repr
7+ , union
8+ , equivalent
9+ , redundant
10+ , descriptor
11+ , setDescriptor
12+ , modifyDescriptor
13+ )
14+ where
15+
16+
17+{- This is based on the following implementations:
18+
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
21+
22+It seems like the OCaml one came first, but I am not sure.
23+
24+Compared to the Haskell implementation, the major changes here include:
25+
26+ 1. No more reallocating PointInfo when changing the weight
27+ 2. Using the strict modifyIORef
28+
29+-}
30+
31+
32+import Control.Monad ( when )
33+import Data.IORef (IORef, modifyIORef', newIORef, readIORef, writeIORef)
34+import Data.Word (Word32)
35+
36+
37+
38+-- POINT
39+
40+
41+newtype Point a =
42+ Pt (IORef (PointInfo a))
43+ deriving Eq
44+
45+
46+data PointInfo a
47+ = Info {-# UNPACK #-} !(IORef Word32) {-# UNPACK #-} !(IORef a)
48+ | Link {-# UNPACK #-} !(Point a)
49+
50+
51+
52+-- HELPERS
53+
54+
55+fresh :: a -> IO (Point a)
56+fresh value =
57+ do weight <- newIORef 1
58+ desc <- newIORef value
59+ link <- newIORef (Info weight desc)
60+ return (Pt link)
61+
62+
63+repr :: Point a -> IO (Point a)
64+repr point@(Pt ref) =
65+ do pInfo <- readIORef ref
66+ case pInfo of
67+ Info _ _ ->
68+ return point
69+
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
76+
77+
78+descriptor :: Point a -> IO a
79+descriptor point@(Pt ref) =
80+ do pInfo <- readIORef ref
81+ case pInfo of
82+ Info _ descRef ->
83+ readIORef descRef
84+
85+ Link (Pt ref1) ->
86+ do link' <- readIORef ref1
87+ case link' of
88+ Info _ descRef ->
89+ readIORef descRef
90+
91+ Link _ ->
92+ descriptor =<< repr point
93+
94+
95+setDescriptor :: Point a -> a -> IO ()
96+setDescriptor point@(Pt ref) newDesc =
97+ do pInfo <- readIORef ref
98+ case pInfo of
99+ Info _ descRef ->
100+ writeIORef descRef newDesc
101+
102+ Link (Pt ref1) ->
103+ do link' <- readIORef ref1
104+ case link' of
105+ Info _ descRef ->
106+ writeIORef descRef newDesc
107+
108+ Link _ ->
109+ do newPoint <- repr point
110+ setDescriptor newPoint newDesc
111+
112+
113+modifyDescriptor :: Point a -> (a -> a) -> IO ()
114+modifyDescriptor point@(Pt ref) func =
115+ do pInfo <- readIORef ref
116+ case pInfo of
117+ Info _ descRef ->
118+ modifyIORef' descRef func
119+
120+ Link (Pt ref1) ->
121+ do link' <- readIORef ref1
122+ case link' of
123+ Info _ descRef ->
124+ modifyIORef' descRef func
125+
126+ Link _ ->
127+ do newPoint <- repr point
128+ modifyDescriptor newPoint func
129+
130+
131+union :: Point a -> Point a -> a -> IO ()
132+union p1 p2 newDesc =
133+ do point1@(Pt ref1) <- repr p1
134+ point2@(Pt ref2) <- repr p2
135+
136+ -- assert (point1 /= point2)
137+
138+ Info w1 d1 <- readIORef ref1
139+ Info w2 d2 <- readIORef ref2
140+
141+ weight1 <- readIORef w1
142+ weight2 <- readIORef w2
143+
144+ let !newWeight = weight1 + weight2
145+
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
155+
156+
157+equivalent :: Point a -> Point a -> IO Bool
158+equivalent p1 p2 =
159+ do v1 <- repr p1
160+ v2 <- repr p2
161+ return (v1 == v2)
162+
163+
164+redundant :: Point a -> IO Bool
165+redundant (Pt ref) =
166+ do pInfo <- readIORef ref
167+ case pInfo of
168+ Info _ _ ->
169+ return False
170+
171+ Link _ ->
172+ return True