Separate length from data in array representation
[packages/dph.git] / dph-common / Data / Array / Parallel / Lifted / Closure.hs
1 module Data.Array.Parallel.Lifted.Closure (
2 (:->)(..), PArray(..),
3 mkClosure, mkClosureP, ($:), ($:^),
4 closure, liftedClosure, liftedApply,
5 dPA_Clo, dPR_Clo,
6
7 closure1, closure2, closure3
8 ) where
9
10 import Data.Array.Parallel.Lifted.PArray
11 import Data.Array.Parallel.Lifted.Instances (dPA_Unit, dPA_2, dPA_3)
12 import Data.Array.Parallel.Lifted.Repr
13
14 import GHC.Exts (Int#)
15
16 infixr 0 :->
17 infixl 0 $:, $:^
18
19 -- |The type of closures
20 --
21 data a :-> b = forall e. Clo (PA e)
22 !(e -> a -> b)
23 !(Int# -> PData e -> PData a -> PData b)
24 e
25
26 lifted :: (PArray e -> PArray a -> PArray b)
27 -> Int# -> PData e -> PData a -> PData b
28 {-# INLINE lifted #-}
29 lifted f n# es as = case f (PArray n# es) (PArray n# as) of PArray _ bs -> bs
30
31 -- |Closure construction
32 --
33 mkClosure :: forall a b e.
34 PA e -> (e -> a -> b)
35 -> (PArray e -> PArray a -> PArray b)
36 -> e -> (a :-> b)
37 {-# INLINE CONLIKE mkClosure #-}
38 mkClosure pa fv fl e = Clo pa fv (lifted fl) e
39
40 closure :: forall a b e.
41 PA e -> (e -> a -> b)
42 -> (Int# -> PData e -> PData a -> PData b)
43 -> e
44 -> (a :-> b)
45 {-# INLINE closure #-}
46 closure pa fv fl e = Clo pa fv fl e
47
48 -- |Closure application
49 --
50 ($:) :: forall a b. (a :-> b) -> a -> b
51 {-# INLINE ($:) #-}
52 Clo _ f _ e $: a = f e a
53
54 {-# RULES
55
56 "mkClosure/($:)" forall pa fv fl e x.
57 mkClosure pa fv fl e $: x = fv e x
58
59 #-}
60
61 -- |Arrays of closures (aka array closures)
62 --
63 data instance PData (a :-> b)
64 = forall e. AClo (PA e)
65 !(e -> a -> b)
66 !(Int# -> PData e -> PData a -> PData b)
67 (PData e)
68
69 -- |Lifted closure construction
70 --
71 mkClosureP :: forall a b e.
72 PA e -> (e -> a -> b)
73 -> (PArray e -> PArray a -> PArray b)
74 -> PArray e -> PArray (a :-> b)
75 {-# INLINE mkClosureP #-}
76 mkClosureP pa fv fl (PArray n# es) = PArray n# (AClo pa fv (lifted fl) es)
77
78 liftedClosure :: forall a b e.
79 PA e -> (e -> a -> b)
80 -> (Int# -> PData e -> PData a -> PData b)
81 -> PData e
82 -> PData (a :-> b)
83 {-# INLINE liftedClosure #-}
84 liftedClosure pa fv fl es = AClo pa fv fl es
85
86 -- |Lifted closure application
87 --
88 ($:^) :: forall a b. PArray (a :-> b) -> PArray a -> PArray b
89 {-# INLINE ($:^) #-}
90 PArray n# (AClo _ _ f es) $:^ PArray _ as = PArray n# (f n# es as)
91
92 liftedApply :: forall a b. Int# -> PData (a :-> b) -> PData a -> PData b
93 {-# INLINE liftedApply #-}
94 liftedApply n# (AClo _ _ f es) as = f n# es as
95
96 type instance PRepr (a :-> b) = a :-> b
97
98 dPA_Clo :: PA a -> PA b -> PA (a :-> b)
99 {-# INLINE dPA_Clo #-}
100 dPA_Clo _ _ = PA {
101 toPRepr = id
102 , fromPRepr = id
103 , toArrPRepr = id
104 , fromArrPRepr = id
105 , dictPRepr = dPR_Clo
106 }
107
108 dPR_Clo :: PR (a :-> b)
109 {-# INLINE dPR_Clo #-}
110 dPR_Clo = PR {
111 emptyPR = emptyPR_Clo
112 , replicatePR = replicatePR_Clo
113 , replicatelPR = replicatelPR_Clo
114 , indexPR = indexPR_Clo
115 , bpermutePR = bpermutePR_Clo
116 , packPR = packPR_Clo
117 }
118
119 {-# INLINE emptyPR_Clo #-}
120 emptyPR_Clo = AClo dPA_Unit (\e a -> error "empty array closure")
121 (\es as -> error "empty array closure")
122 (emptyPD dPA_Unit)
123
124 {-# INLINE replicatePR_Clo #-}
125 replicatePR_Clo n# (Clo pa f f' e) = AClo pa f f' (replicatePD pa n# e)
126
127 {-# INLINE replicatelPR_Clo #-}
128 replicatelPR_Clo segd (AClo pa f f' es)
129 = AClo pa f f' (replicatelPD pa segd es)
130
131 {-# INLINE indexPR_Clo #-}
132 indexPR_Clo (AClo pa f f' es) i# = Clo pa f f' (indexPD pa es i#)
133
134 {-# INLINE bpermutePR_Clo #-}
135 bpermutePR_Clo (AClo pa f f' es) n# is = AClo pa f f' (bpermutePD pa es n# is)
136
137 {-# INLINE packPR_Clo #-}
138 packPR_Clo (AClo pa f f' es) n# sel = AClo pa f f' (packPD pa es n# sel)
139
140 -- Closure construction
141
142 closure1 :: (a -> b) -> (PArray a -> PArray b) -> (a :-> b)
143 {-# INLINE closure1 #-}
144 closure1 fv fl = mkClosure dPA_Unit (\_ -> fv) (\_ -> fl) ()
145
146 closure2 :: PA a
147 -> (a -> b -> c)
148 -> (PArray a -> PArray b -> PArray c)
149 -> (a :-> b :-> c)
150 {-# INLINE closure2 #-}
151 closure2 pa fv fl = mkClosure dPA_Unit fv_1 fl_1 ()
152 where
153 fv_1 _ x = mkClosure pa fv fl x
154 fl_1 _ xs = mkClosureP pa fv fl xs
155
156 closure3 :: PA a -> PA b
157 -> (a -> b -> c -> d)
158 -> (PArray a -> PArray b -> PArray c -> PArray d)
159 -> (a :-> b :-> c :-> d)
160 {-# INLINE closure3 #-}
161 closure3 pa pb fv fl = mkClosure dPA_Unit fv_1 fl_1 ()
162 where
163 fv_1 _ x = mkClosure pa fv_2 fl_2 x
164 fl_1 _ xs = mkClosureP pa fv_2 fl_2 xs
165
166 fv_2 x y = mkClosure (dPA_2 pa pb) fv_3 fl_3 (x,y)
167 fl_2 xs ys = mkClosureP (dPA_2 pa pb) fv_3 fl_3 (zipPA# xs ys)
168
169 fv_3 (x,y) z = fv x y z
170 fl_3 ps zs = case unzipPA# ps of (xs,ys) -> fl xs ys zs
171