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