8f55f3e30b1b5c9c69c3c9d09093a29de7f5d9c2
[darcs-mirrors/vector.git] / internal / GenUnboxTuple.hs
1 {-# LANGUAGE ParallelListComp #-}
2 module Main where
3
4 import Text.PrettyPrint
5
6 import System.Environment ( getArgs )
7
8 main = do
9 [s] <- getArgs
10 let n = read s
11 mapM_ (putStrLn . render . generate) [2..n]
12
13 generate :: Int -> Doc
14 generate n =
15 vcat [ text "#ifdef DEFINE_INSTANCES"
16 , data_instance "MVector s" "MV"
17 , data_instance "Vector" "V"
18 , class_instance "Unbox"
19 , class_instance "M.MVector MVector" <+> text "where"
20 , nest 2 $ vcat $ map method methods_MVector
21 , class_instance "G.Vector Vector" <+> text "where"
22 , nest 2 $ vcat $ map method methods_Vector
23 , text "#endif"
24 , text "#ifdef DEFINE_MUTABLE"
25 , define_zip "MVector s" "MV"
26 , define_unzip "MVector s" "MV"
27 , text "#endif"
28 , text "#ifdef DEFINE_IMMUTABLE"
29 , define_zip "Vector" "V"
30 , define_unzip "Vector" "V"
31 , text "#endif"
32 ]
33
34 where
35 vars = map char $ take n ['a'..]
36 varss = map (<> char 's') vars
37 tuple xs = parens $ hsep $ punctuate comma xs
38 vtuple xs = parens $ sep $ punctuate comma xs
39 con s = text s <> char '_' <> int n
40 var c = text (c : "_")
41
42 data_instance ty c
43 = hang (hsep [text "data instance", text ty, tuple vars])
44 4
45 (hsep [char '=', con c, text "{-# UNPACK #-} !Int"
46 , vcat $ map (\v -> parens (text ty <+> v)) vars])
47
48 class_instance cls
49 = text "instance" <+> vtuple [text "Unbox" <+> v | v <- vars]
50 <+> text "=>" <+> text cls <+> tuple vars
51
52
53 define_zip ty c
54 = sep [name <+> text "::"
55 <+> vtuple [text "Unbox" <+> v | v <- vars]
56 <+> text "=>"
57 <+> sep (punctuate (text " ->") [text ty <+> v | v <- vars])
58 <+> text "->"
59 <+> text ty <+> tuple vars
60 ,text "{-# INLINE" <+> name <+> text "#-}"
61 ,name <+> sep varss
62 <+> text "="
63 <+> con c
64 <+> text "len"
65 <+> sep [parens $ text "unsafeSlice"
66 <+> vs
67 <+> char '0'
68 <+> text "len" | vs <- varss]
69 ,nest 2 $ hang (text "where")
70 2
71 $ text "len ="
72 <+> sep (punctuate (text " `min`")
73 [text "length" <+> vs | vs <- varss])
74 ]
75 where
76 name | n == 2 = text "zip"
77 | otherwise = text "zip" <> int n
78
79 define_unzip ty c
80 = sep [name <+> text "::"
81 <+> vtuple [text "Unbox" <+> v | v <- vars]
82 <+> text "=>"
83 <+> text ty <+> tuple vars
84 <+> text "->" <+> vtuple [text ty <+> v | v <- vars]
85 ,text "{-# INLINE" <+> name <+> text "#-}"
86 ,name <+> pat c <+> text "="
87 <+> vtuple varss
88 ]
89 where
90 name | n == 2 = text "unzip"
91 | otherwise = text "unzip" <> int n
92
93 pat c = parens $ con c <+> var 'n' <+> sep varss
94 patn c n = parens $ con c <+> (var 'n' <> int n)
95 <+> sep [v <> int n | v <- varss]
96
97 qM s = text "M." <> text s
98 qG s = text "G." <> text s
99
100 gen_length c _ = (pat c, var 'n')
101
102 gen_unsafeSlice mod c rec
103 = (pat c <+> var 'i' <+> var 'm',
104 con c <+> var 'm'
105 <+> vcat [parens
106 $ text mod <> char '.' <> text rec
107 <+> vs <+> var 'i' <+> var 'm'
108 | vs <- varss])
109
110
111 gen_overlaps rec = (patn "MV" 1 <+> patn "MV" 2,
112 vcat $ r : [text "||" <+> r | r <- rs])
113 where
114 r : rs = [qM rec <+> v <> char '1' <+> v <> char '2' | v <- varss]
115
116 gen_unsafeNew rec
117 = (var 'n',
118 mk_do [v <+> text "<-" <+> qM rec <+> var 'n' | v <- varss]
119 $ text "return $" <+> con "MV" <+> var 'n' <+> sep varss)
120
121 gen_unsafeNewWith rec
122 = (var 'n' <+> tuple vars,
123 mk_do [vs <+> text "<-" <+> qM rec <+> var 'n' <+> v
124 | v <- vars | vs <- varss]
125 $ text "return $" <+> con "MV" <+> var 'n' <+> sep varss)
126
127 gen_unsafeRead rec
128 = (pat "MV" <+> var 'i',
129 mk_do [v <+> text "<-" <+> qM rec <+> vs <+> var 'i' | v <- vars
130 | vs <- varss]
131 $ text "return" <+> tuple vars)
132
133 gen_unsafeWrite rec
134 = (pat "MV" <+> var 'i' <+> tuple vars,
135 mk_do [qM rec <+> vs <+> var 'i' <+> v | v <- vars | vs <- varss]
136 empty)
137
138 gen_clear rec
139 = (pat "MV", mk_do [qM rec <+> vs | vs <- varss] empty)
140
141 gen_set rec
142 = (pat "MV" <+> tuple vars,
143 mk_do [qM rec <+> vs <+> v | vs <- varss | v <- vars] empty)
144
145 gen_unsafeCopy rec
146 = (patn "MV" 1 <+> patn "MV" 2,
147 mk_do [qM rec <+> vs <> char '1' <+> vs <> char '2' | vs <- varss]
148 empty)
149
150 gen_unsafeGrow rec
151 = (pat "MV" <+> var 'm',
152 mk_do [qM rec <+> vs <+> var 'm' | vs <- varss]
153 $ text "return $" <+> con "MV"
154 <+> parens (var 'm' <> char '+' <> var 'n')
155 <+> sep varss)
156
157 gen_unsafeFreeze rec
158 = (pat "MV",
159 mk_do [vs <> char '\'' <+> text "<-" <+> qG rec <+> vs | vs <- varss]
160 $ text "return $" <+> con "V" <+> var 'n'
161 <+> sep [vs <> char '\'' | vs <- varss])
162
163 gen_basicUnsafeIndexM rec
164 = (pat "V" <+> var 'i',
165 mk_do [v <+> text "<-" <+> qG rec <+> vs <+> var 'i'
166 | vs <- varss | v <- vars]
167 $ text "return" <+> tuple vars)
168
169
170
171
172 mk_do cmds ret = hang (text "do")
173 2
174 $ vcat $ cmds ++ [ret]
175
176 method (s, f) = case f s of
177 (p,e) -> text "{-# INLINE" <+> text s <+> text " #-}"
178 $$ hang (text s <+> p)
179 4
180 (char '=' <+> e)
181
182
183 methods_MVector = [("basicLength", gen_length "MV")
184 ,("basicUnsafeSlice", gen_unsafeSlice "M" "MV")
185 ,("basicOverlaps", gen_overlaps)
186 ,("basicUnsafeNew", gen_unsafeNew)
187 ,("basicUnsafeNewWith", gen_unsafeNewWith)
188 ,("basicUnsafeRead", gen_unsafeRead)
189 ,("basicUnsafeWrite", gen_unsafeWrite)
190 ,("basicClear", gen_clear)
191 ,("basicSet", gen_set)
192 ,("basicUnsafeCopy", gen_unsafeCopy)
193 ,("basicUnsafeGrow", gen_unsafeGrow)]
194
195 methods_Vector = [("unsafeFreeze", gen_unsafeFreeze)
196 ,("basicLength", gen_length "V")
197 ,("basicUnsafeSlice", gen_unsafeSlice "G" "V")
198 ,("basicUnsafeIndexM", gen_basicUnsafeIndexM)]