Testsuite: tabs -> spaces [skip ci]
[ghc.git] / testsuite / tests / arrows / should_run / arrowrun002.hs
1 {-# LANGUAGE Arrows #-}
2
3 -- Homogeneous (or depth-preserving) functions over perfectly balanced trees.
4
5 module Main where
6
7 import Control.Arrow
8 import Control.Category
9 import Data.Complex
10 import Prelude hiding (id, (.))
11
12 infixr 4 :&:
13
14 -- Consider the following non-regular type of perfectly balanced trees,
15 -- or `powertrees' (cf Jayadev Misra's powerlists):
16
17 data Pow a = Zero a | Succ (Pow (Pair a))
18 deriving Show
19
20 type Pair a = (a, a)
21
22 -- Here are some example elements:
23
24 tree0 = Zero 1
25 tree1 = Succ (Zero (1, 2))
26 tree2 = Succ (Succ (Zero ((1, 2), (3, 4))))
27 tree3 = Succ (Succ (Succ (Zero (((1, 2), (3, 4)), ((5, 6), (7, 8))))))
28
29 -- The elements of this type have a string of constructors expressing
30 -- a depth n as a Peano numeral, enclosing a nested pair tree of 2^n
31 -- elements. The type definition ensures that all elements of this type
32 -- are perfectly balanced binary trees of this form. (Such things arise
33 -- in circuit design, eg Ruby, and descriptions of parallel algorithms.)
34 -- And the type system will ensure that all legal programs preserve
35 -- this structural invariant.
36 --
37 -- The only problem is that the type constraint is too restrictive, rejecting
38 -- many of the standard operations on these trees. Typically you want to
39 -- split a tree into two subtrees, do some processing on the subtrees and
40 -- combine the results. But the type system cannot discover that the two
41 -- results are of the same depth (and thus combinable). We need a type
42 -- that says a function preserves depth. Here it is:
43
44 data Hom a b = (a -> b) :&: Hom (Pair a) (Pair b)
45
46 -- A homogeneous (or depth-preserving) function is an infinite sequence of
47 -- functions of type Pair^n a -> Pair^n b, one for each depth n. We can
48 -- apply a homogeneous function to a powertree by selecting the function
49 -- for the required depth:
50
51 apply :: Hom a b -> Pow a -> Pow b
52 apply (f :&: fs) (Zero x) = Zero (f x)
53 apply (f :&: fs) (Succ t) = Succ (apply fs t)
54
55 -- Having defined apply, we can forget about powertrees and do all our
56 -- programming with Hom's. Firstly, Hom is an arrow:
57
58 instance Category Hom where
59 id = id :&: id
60 (f :&: fs) . (g :&: gs) = (f . g) :&: (fs . gs)
61
62 instance Arrow Hom where
63 arr f = f :&: arr (f *** f)
64 first (f :&: fs) =
65 first f :&: (arr transpose >>> first fs >>> arr transpose)
66
67 transpose :: ((a,b), (c,d)) -> ((a,c), (b,d))
68 transpose ((a,b), (c,d)) = ((a,c), (b,d))
69
70 -- arr maps f over the leaves of a powertree.
71
72 -- The composition >>> composes sequences of functions pairwise.
73 --
74 -- The *** operator unriffles a powertree of pairs into a pair of powertrees,
75 -- applies the appropriate function to each and riffles the results.
76 -- It defines a categorical product for this arrow category.
77
78 -- When describing algorithms, one often provides a pure function for the
79 -- base case (trees of one element) and a (usually recursive) expression
80 -- for trees of pairs.
81
82 -- For example, a common divide-and-conquer pattern is the butterfly, where
83 -- one recursive call processes the odd-numbered elements and the other
84 -- processes the even ones (cf Geraint Jones and Mary Sheeran's Ruby papers):
85
86 butterfly :: (Pair a -> Pair a) -> Hom a a
87 butterfly f = id :&: proc (x, y) -> do
88 x' <- butterfly f -< x
89 y' <- butterfly f -< y
90 returnA -< f (x', y')
91
92 -- The recursive calls operate on halves of the original tree, so the
93 -- recursion is well-defined.
94
95 -- Some examples of butterflies:
96
97 rev :: Hom a a
98 rev = butterfly swap
99 where swap (x, y) = (y, x)
100
101 unriffle :: Hom (Pair a) (Pair a)
102 unriffle = butterfly transpose
103
104 -- Batcher's sorter for bitonic sequences:
105
106 bisort :: Ord a => Hom a a
107 bisort = butterfly cmp
108 where cmp (x, y) = (min x y, max x y)
109
110 -- This can be used (with rev) as the merge phase of a merge sort.
111 --
112 sort :: Ord a => Hom a a
113 sort = id :&: proc (x, y) -> do
114 x' <- sort -< x
115 y' <- sort -< y
116 yr <- rev -< y'
117 p <- unriffle -< (x', yr)
118 bisort2 -< p
119 where _ :&: bisort2 = bisort
120
121 -- Here is the scan operation, using the algorithm of Ladner and Fischer:
122
123 scan :: (a -> a -> a) -> a -> Hom a a
124 scan op b = id :&: proc (x, y) -> do
125 y' <- scan op b -< op x y
126 l <- rsh b -< y'
127 returnA -< (op l x, y')
128
129 -- The auxiliary function rsh b shifts each element in the tree one place to
130 -- the right, placing b in the now-vacant leftmost position, and discarding
131 -- the old rightmost element:
132
133 rsh :: a -> Hom a a
134 rsh b = const b :&: proc (x, y) -> do
135 w <- rsh b -< y
136 returnA -< (w, x)
137
138 -- Finally, here is the Fast Fourier Transform:
139
140 type C = Complex Double
141
142 fft :: Hom C C
143 fft = id :&: proc (x, y) -> do
144 x' <- fft -< x
145 y' <- fft -< y
146 r <- roots (-1) -< ()
147 let z = r*y'
148 unriffle -< (x' + z, x' - z)
149
150 -- The auxiliary function roots r (where r is typically a root of unity)
151 -- populates a tree of size n (necessarily a power of 2) with the values
152 -- 1, w, w^2, ..., w^(n-1), where w^n = r.
153
154 roots :: C -> Hom () C
155 roots r = const 1 :&: proc _ -> do
156 x <- roots r' -< ()
157 unriffle -< (x, x*r')
158 where r' = if imagPart s >= 0 then -s else s
159 s = sqrt r
160
161 -- Miscellaneous functions:
162
163 rrot :: Hom a a
164 rrot = id :&: proc (x, y) -> do
165 w <- rrot -< y
166 returnA -< (w, x)
167
168 ilv :: Hom a a -> Hom (Pair a) (Pair a)
169 ilv f = proc (x, y) -> do
170 x' <- f -< x
171 y' <- f -< y
172 returnA -< (x', y')
173
174 scan' :: (a -> a -> a) -> a -> Hom a a
175 scan' op b = proc x -> do
176 l <- rsh b -< x
177 (id :&: ilv (scan' op b)) -< op l x
178
179 riffle :: Hom (Pair a) (Pair a)
180 riffle = id :&: proc ((x1, y1), (x2, y2)) -> do
181 x <- riffle -< (x1, x2)
182 y <- riffle -< (y1, y2)
183 returnA -< (x, y)
184
185 invert :: Hom a a
186 invert = id :&: proc (x, y) -> do
187 x' <- invert -< x
188 y' <- invert -< y
189 unriffle -< (x', y')
190
191 carryLookaheadAdder :: Hom (Bool, Bool) Bool
192 carryLookaheadAdder = proc (x, y) -> do
193 carryOut <- rsh (Just False) -<
194 if x == y then Just x else Nothing
195 Just carryIn <- scan plusMaybe Nothing -< carryOut
196 returnA -< x `xor` y `xor` carryIn
197 where plusMaybe x Nothing = x
198 plusMaybe x (Just y) = Just y
199 False `xor` b = b
200 True `xor` b = not b
201
202 -- Global conditional for SIMD
203
204 ifAll :: Hom a b -> Hom a b -> Hom (a, Bool) b
205 ifAll fs gs = ifAllAux snd (arr fst >>> fs) (arr fst >>> gs)
206 where ifAllAux :: (a -> Bool) -> Hom a b -> Hom a b -> Hom a b
207 ifAllAux p (f :&: fs) (g :&: gs) =
208 liftIf p f g :&: ifAllAux (liftAnd p) fs gs
209 liftIf p f g x = if p x then f x else g x
210 liftAnd p (x, y) = p x && p y
211
212 maybeAll :: Hom a c -> Hom (a, b) c -> Hom (a, Maybe b) c
213 maybeAll (n :&: ns) (j :&: js) =
214 choose :&: (arr dist >>> maybeAll ns (arr transpose >>> js))
215 where choose (a, Nothing) = n a
216 choose (a, Just b) = j (a, b)
217 dist ((a1, b1), (a2, b2)) = ((a1, a2), zipMaybe b1 b2)
218 zipMaybe (Just x) (Just y) = Just (x, y)
219 zipMaybe _ _ = Nothing
220
221 main = do
222 print (apply rev tree3)
223 print (apply invert tree3)
224 print (apply (invert >>> sort) tree3)
225 print (apply (scan (+) 0) tree3)