@@ -0,0 +1,487 @@
1+{-# LANGUAGE TemplateHaskell #-}
2+
3+module Streamly.Internal.Data.Unbox.TH
4+ ( deriveUnbox
5+ , deriveUnboxWith
6+ ) where
7+
8+--------------------------------------------------------------------------------
9+-- Imports
10+--------------------------------------------------------------------------------
11+
12+import Data.Word (Word16, Word32, Word64, Word8)
13+import Data.Proxy (Proxy(..))
14+
15+import Language.Haskell.TH
16+import Language.Haskell.TH.Syntax
17+import Streamly.Internal.Data.Unbox
18+
19+--------------------------------------------------------------------------------
20+-- th-utilities
21+--------------------------------------------------------------------------------
22+
23+-- The following are copied to remove the dependency on th-utilities.
24+-- The code has been copied from th-abstraction and th-utilities.
25+
26+-- Some CPP macros in the following code are not required but are kept
27+-- anyway. They can be removed if deemed as a maintainance burden.
28+
29+#if MIN_VERSION_template_haskell(2,17,0)
30+type TyVarBndr_ flag = TyVarBndr flag
31+#else
32+type TyVarBndr_ flag = TyVarBndr
33+#endif
34+
35+-- | Case analysis for a 'TyVarBndr'. If the value is a @'PlainTV' n _@, apply
36+-- the first function to @n@; if it is @'KindedTV' n _ k@, apply the second
37+-- function to @n@ and @k@.
38+elimTV :: (Name -> r) -> (Name -> Kind -> r) -> TyVarBndr_ flag -> r
39+#if MIN_VERSION_template_haskell(2,17,0)
40+elimTV ptv _ktv (PlainTV n _) = ptv n
41+elimTV _ptv ktv (KindedTV n _ k) = ktv n k
42+#else
43+elimTV ptv _ktv (PlainTV n) = ptv n
44+elimTV _ptv ktv (KindedTV n k) = ktv n k
45+#endif
46+
47+-- | Extract the type variable name from a 'TyVarBndr', ignoring the
48+-- kind signature if one exists.
49+tvName :: TyVarBndr_ flag -> Name
50+tvName = elimTV id (\n _ -> n)
51+
52+-- | Get the 'Name' of a 'TyVarBndr'
53+tyVarBndrName :: TyVarBndr_ flag -> Name
54+tyVarBndrName = tvName
55+
56+appsT :: Type -> [Type] -> Type
57+appsT x [] = x
58+appsT x (y:xs) = appsT (AppT x y) xs
59+
60+-- | Utility to conveniently handle change to 'InstanceD' API in
61+-- template-haskell-2.11.0
62+plainInstanceD :: Cxt -> Type -> [Dec] -> Dec
63+plainInstanceD =
64+#if MIN_VERSION_template_haskell(2,11,0)
65+ InstanceD Nothing
66+#else
67+ InstanceD
68+#endif
69+
70+-- | Simplified info about a 'DataD'. Omits deriving, strictness,
71+-- kind info, and whether it's @data@ or @newtype@.
72+data DataType = DataType
73+ { dtName :: Name
74+ , dtTvs :: [Name]
75+ , dtCxt :: Cxt
76+ , dtCons :: [DataCon]
77+ } deriving (Eq, Show, Ord) --, Data, Typeable, Generic)
78+
79+-- | Simplified info about a 'Con'. Omits deriving, strictness, and kind
80+-- info. This is much nicer than consuming 'Con' directly, because it
81+-- unifies all the constructors into one.
82+data DataCon = DataCon
83+ { dcName :: Name
84+ , dcTvs :: [Name]
85+ , dcCxt :: Cxt
86+ , dcFields :: [(Maybe Name, Type)]
87+ } deriving (Eq, Show, Ord) --, Data, Typeable, Generic)
88+
89+
90+-- | Convert a 'Con' to a list of 'DataCon'. The result is a list
91+-- because 'GadtC' and 'RecGadtC' can define multiple constructors.
92+conToDataCons :: Con -> [DataCon]
93+conToDataCons = \case
94+ NormalC name slots ->
95+ [DataCon name [] [] (map (\(_, ty) -> (Nothing, ty)) slots)]
96+ RecC name fields ->
97+ [DataCon name [] [] (map (\(n, _, ty) -> (Just n, ty)) fields)]
98+ InfixC (_, ty1) name (_, ty2) ->
99+ [DataCon name [] [] [(Nothing, ty1), (Nothing, ty2)]]
100+ ForallC tvs preds con ->
101+ map (\(DataCon name tvs0 preds0 fields) ->
102+ DataCon name (tvs0 ++ map tyVarBndrName tvs) (preds0 ++ preds) fields) (conToDataCons con)
103+#if MIN_VERSION_template_haskell(2,11,0)
104+ GadtC ns slots _ ->
105+ map (\dn -> DataCon dn [] [] (map (\(_, ty) -> (Nothing, ty)) slots)) ns
106+ RecGadtC ns fields _ ->
107+ map (\dn -> DataCon dn [] [] (map (\(fn, _, ty) -> (Just fn, ty)) fields)) ns
108+#endif
109+
110+-- | Reify the given data or newtype declaration, and yields its
111+-- 'DataType' representation.
112+reifyDataType :: Name -> Q DataType
113+reifyDataType name = do
114+ info <- reify name
115+ case infoToDataType info of
116+ Nothing -> fail $ "Expected to reify a datatype. Instead got:\n" ++ pprint info
117+ Just x -> return x
118+
119+infoToDataType :: Info -> Maybe DataType
120+infoToDataType info = case info of
121+#if MIN_VERSION_template_haskell(2,11,0)
122+ TyConI (DataD preds name tvs _kind cons _deriving) ->
123+#else
124+ TyConI (DataD preds name tvs cons _deriving) ->
125+#endif
126+ Just $ DataType name (map tyVarBndrName tvs) preds (concatMap conToDataCons cons)
127+#if MIN_VERSION_template_haskell(2,11,0)
128+ TyConI (NewtypeD preds name tvs _kind con _deriving) ->
129+#else
130+ TyConI (NewtypeD preds name tvs con _deriving) ->
131+#endif
132+ Just $ DataType name (map tyVarBndrName tvs) preds (conToDataCons con)
133+ _ -> Nothing
134+
135+--------------------------------------------------------------------------------
136+-- Helpers
137+--------------------------------------------------------------------------------
138+
139+type Field = (Maybe Name, Type)
140+
141+_arr :: Name
142+_arr = mkName "arr"
143+
144+_tag :: Name
145+_tag = mkName "tag"
146+
147+_initialOffset :: Name
148+_initialOffset = mkName "initialOffset"
149+
150+_val :: Name
151+_val = mkName "val"
152+
153+mkOffsetName :: Int -> Name
154+mkOffsetName i = mkName ("offset" ++ show i)
155+
156+mkFieldName :: Int -> Name
157+mkFieldName i = mkName ("field" ++ show i)
158+
159+--------------------------------------------------------------------------------
160+-- Domain specific helpers
161+--------------------------------------------------------------------------------
162+
163+exprGetSize :: Type -> Q Exp
164+exprGetSize ty = appE (varE 'sizeOf) [|Proxy :: Proxy $(pure ty)|]
165+
166+getTagSize :: Int -> Int
167+getTagSize numConstructors
168+ | numConstructors == 1 = 0
169+ | fromIntegral (maxBound :: Word8) >= numConstructors = 1
170+ | fromIntegral (maxBound :: Word16) >= numConstructors = 2
171+ | fromIntegral (maxBound :: Word32) >= numConstructors = 4
172+ | fromIntegral (maxBound :: Word64) >= numConstructors = 8
173+ | otherwise = error "Too many constructors"
174+
175+getTagType :: Int -> Name
176+getTagType numConstructors
177+ | numConstructors == 1 = error "No tag for 1 constructor"
178+ | fromIntegral (maxBound :: Word8) >= numConstructors = ''Word8
179+ | fromIntegral (maxBound :: Word16) >= numConstructors = ''Word16
180+ | fromIntegral (maxBound :: Word32) >= numConstructors = ''Word32
181+ | fromIntegral (maxBound :: Word64) >= numConstructors = ''Word64
182+ | otherwise = error "Too many constructors"
183+
184+mkOffsetDecls :: Int -> [Field] -> [Q Dec]
185+mkOffsetDecls tagSize fields =
186+ init
187+ ((:) (valD
188+ (varP (mkOffsetName 0))
189+ (normalB
190+ [|$(litE (IntegerL (fromIntegral tagSize))) +
191+ $(varE _initialOffset)|])
192+ [])
193+ (map mkOffsetExpr (zip [1 ..] fields)))
194+
195+ where
196+
197+ mkOffsetExpr (i, (_, ty)) =
198+ valD
199+ (varP (mkOffsetName i))
200+ (normalB [|$(varE (mkOffsetName (i - 1))) + $(exprGetSize ty)|])
201+ []
202+
203+--------------------------------------------------------------------------------
204+-- Size
205+--------------------------------------------------------------------------------
206+
207+isUnitType :: [DataCon] -> Bool
208+isUnitType [DataCon _ _ _ []] = True
209+isUnitType _ = False
210+
211+mkSizeOfExpr :: Type -> [DataCon] -> Q Exp
212+mkSizeOfExpr headTy constructors =
213+ case constructors of
214+ [] ->
215+ [|error
216+ ("Attempting to get size with no constructors (" ++
217+ $(lift (pprint headTy)) ++ ")")|]
218+ -- One constructor with no fields is a unit type. Size of a unit type is
219+ -- 1.
220+ [con@(DataCon _ _ _ fields)] ->
221+ case fields of
222+ [] -> litE (IntegerL 1)
223+ _ -> [|$(sizeOfConstructor con)|]
224+ _ -> [|$(litE (IntegerL (fromIntegral tagSize))) + $(sizeOfHeadDt)|]
225+
226+ where
227+
228+ tagSize = getTagSize (length constructors)
229+ sizeOfField (_, ty) = exprGetSize ty
230+ sizeOfConstructor (DataCon _ _ _ fields) =
231+ appE (varE 'sum) (listE (map sizeOfField fields))
232+ -- The size of any Unbox type is atleast 1
233+ sizeOfHeadDt =
234+ appE (varE 'maximum) (listE (map sizeOfConstructor constructors))
235+
236+--------------------------------------------------------------------------------
237+-- Peek
238+--------------------------------------------------------------------------------
239+
240+mkPeekExprOne :: Int -> DataCon -> Q Exp
241+mkPeekExprOne tagSize (DataCon cname _ _ fields) =
242+ case fields of
243+ [] -> [|pure $(conE cname)|]
244+ _ ->
245+ letE
246+ (mkOffsetDecls tagSize fields)
247+ (foldl
248+ (\acc i -> [|$(acc) <*> $(peekField i)|])
249+ [|$(conE cname) <$> $(peekField 0)|]
250+ [1 .. (length fields - 1)])
251+
252+ where
253+
254+ peekField i = [|peekByteIndex $(varE (mkOffsetName i)) $(varE _arr)|]
255+
256+mkPeekExpr :: Type -> [DataCon] -> Q Exp
257+mkPeekExpr headTy cons =
258+ case cons of
259+ [] ->
260+ [|error
261+ ("Attempting to peek type with no constructors (" ++
262+ $(lift (pprint headTy)) ++ ")")|]
263+ [con] -> mkPeekExprOne 0 con
264+ _ ->
265+ doE
266+ [ bindS
267+ (varP _tag)
268+ [|peekByteIndex $(varE _initialOffset) $(varE _arr)|]
269+ , noBindS
270+ (caseE
271+ (sigE (varE _tag) (conT tagType))
272+ (map peekMatch (zip [0 ..] cons) ++ [peekErr]))
273+ ]
274+
275+ where
276+
277+ lenCons = length cons
278+ tagType = getTagType lenCons
279+ tagSize = getTagSize lenCons
280+ peekMatch (i, con) =
281+ match (litP (IntegerL i)) (normalB (mkPeekExprOne tagSize con)) []
282+ peekErr =
283+ match
284+ wildP
285+ (normalB
286+ [|error
287+ ("Found invalid tag while peeking (" ++
288+ $(lift (pprint headTy)) ++ ")")|])
289+ []
290+
291+--------------------------------------------------------------------------------
292+-- Poke
293+--------------------------------------------------------------------------------
294+
295+mkPokeExprTag :: Name -> Int -> Q Exp
296+mkPokeExprTag tagType tagVal = pokeTag
297+
298+ where
299+
300+ pokeTag =
301+ [|pokeByteIndex
302+ $(varE _initialOffset)
303+ $(varE _arr)
304+ $((sigE (litE (IntegerL (fromIntegral tagVal))) (conT tagType)))|]
305+
306+mkPokeExprFields :: Int -> [Field] -> Q Exp
307+mkPokeExprFields tagSize fields = do
308+ case fields of
309+ [] -> [|pure ()|]
310+ _ ->
311+ letE
312+ (mkOffsetDecls tagSize fields)
313+ (doE $ map (noBindS . pokeField) [0 .. (numFields - 1)])
314+
315+ where
316+
317+ numFields = length fields
318+ pokeField i =
319+ [|pokeByteIndex
320+ $(varE (mkOffsetName i))
321+ $(varE _arr)
322+ $(varE (mkFieldName i))|]
323+
324+mkPokeMatch :: Name -> Int -> Q Exp -> Q Match
325+mkPokeMatch cname numFields exp0 =
326+ match
327+ (conP cname (map varP (map mkFieldName [0 .. (numFields - 1)])))
328+ (normalB exp0)
329+ []
330+
331+mkPokeExpr :: Type -> [DataCon] -> Q Exp
332+mkPokeExpr headTy cons =
333+ case cons of
334+ [] ->
335+ [|error
336+ ("Attempting to poke type with no constructors (" ++
337+ $(lift (pprint headTy)) ++ ")")|]
338+ [(DataCon _ _ _ [])] -> [|pure ()|]
339+ [(DataCon cname _ _ fields)] ->
340+ caseE
341+ (varE _val)
342+ [mkPokeMatch cname (length fields) (mkPokeExprFields 0 fields)]
343+ _ ->
344+ caseE
345+ (varE _val)
346+ (map (\(tagVal, (DataCon cname _ _ fields)) ->
347+ mkPokeMatch
348+ cname
349+ (length fields)
350+ (doE [ noBindS $ mkPokeExprTag tagType tagVal
351+ , noBindS $ mkPokeExprFields tagSize fields
352+ ]))
353+ (zip [0 ..] cons))
354+
355+ where
356+
357+ lenCons = length cons
358+ tagType = getTagType lenCons
359+ tagSize = getTagSize lenCons
360+
361+--------------------------------------------------------------------------------
362+-- Main
363+--------------------------------------------------------------------------------
364+
365+-- | A general function to derive Unbox instances where you can control which
366+-- Constructors of the datatype to consider and what the Context for the 'Unbox'
367+-- instance would be.
368+--
369+-- Consider the datatype:
370+-- @
371+-- data CustomDataType a b
372+-- = CDTConstructor1
373+-- | CDTConstructor2 Bool
374+-- | CDTConstructor3 Bool b
375+-- deriving (Show, Eq)
376+-- @
377+--
378+-- Usage:
379+-- @
380+-- $(deriveUnboxInternal
381+-- [AppT (ConT ''Unbox) (VarT (mkName "b"))]
382+-- (AppT
383+-- (AppT (ConT ''CustomDataType) (VarT (mkName "a")))
384+-- (VarT (mkName "b")))
385+-- [ DataCon 'CDTConstructor1 [] [] []
386+-- , DataCon 'CDTConstructor2 [] [] [(Nothing, (ConT ''Bool))]
387+-- , DataCon
388+-- 'CDTConstructor3
389+-- []
390+-- []
391+-- [(Nothing, (ConT ''Bool)), (Nothing, (VarT (mkName "b")))]
392+-- ])
393+-- @
394+deriveUnboxInternal :: Cxt -> Type -> [DataCon] -> Q [Dec]
395+deriveUnboxInternal preds headTy cons = do
396+ sizeOfMethod <- mkSizeOfExpr headTy cons
397+ peekMethod <- mkPeekExpr headTy cons
398+ pokeMethod <- mkPokeExpr headTy cons
399+ let methods =
400+ [ FunD 'sizeOf [Clause [WildP] (NormalB sizeOfMethod) []]
401+ , FunD
402+ 'peekByteIndex
403+ [ Clause
404+ (if isUnitType cons
405+ then [WildP, WildP]
406+ else [VarP _initialOffset, VarP _arr])
407+ (NormalB peekMethod)
408+ []
409+ ]
410+ , FunD
411+ 'pokeByteIndex
412+ [ Clause
413+ (if isUnitType cons
414+ then [WildP, WildP, WildP]
415+ else [VarP _initialOffset, VarP _arr, VarP _val])
416+ (NormalB pokeMethod)
417+ []
418+ ]
419+ ]
420+ return [plainInstanceD preds (AppT (ConT ''Unbox) headTy) methods]
421+
422+-- | Template haskell helper to create instances of 'Unbox' automatically.
423+--
424+-- Consider the datatype:
425+-- @
426+-- data CustomDataType a b c = ...
427+-- @
428+--
429+-- Usage: @$(deriveUnbox ''CustomDataType)@
430+--
431+-- Note: All type variables automatcally get an "Unbox" constraint.
432+-- The derived code will look like the following,
433+-- @
434+-- instance (Unbox a, Unbox b, Unbox c) => Unbox (CustomDataType a b c) where
435+-- ...
436+-- @
437+--
438+-- To control which type variables get the Unbox constraint, use
439+-- 'deriveUnboxWith'.
440+deriveUnbox :: Name -> Q [Dec]
441+deriveUnbox name = do
442+ dt <- reifyDataType name
443+ let preds = map (unboxPred . VarT) (dtTvs dt)
444+ headTy = appsT (ConT name) (map VarT (dtTvs dt))
445+ cons = dtCons dt
446+ deriveUnboxInternal preds headTy cons
447+
448+ where
449+
450+ unboxPred ty =
451+#if MIN_VERSION_template_haskell(2,10,0)
452+ AppT (ConT ''Unbox) ty
453+#else
454+ ClassP ''Unbox [ty]
455+#endif
456+
457+-- | Like 'deriveUnbox' but control which type variables get the 'Unbox'
458+-- constraint.
459+--
460+-- Consider the datatype:
461+-- @
462+-- data CustomDataType a b c = ...
463+-- @
464+--
465+-- Usage: @$(deriveUnboxWith ["a", "c"] ''CustomDataType)@
466+--
467+-- @
468+-- instance (Unbox a, Unbox c) => Unbox (CustomDataType a b c) where
469+-- ...
470+-- @
471+--
472+deriveUnboxWith :: [String] -> Name -> Q [Dec]
473+deriveUnboxWith vars name = do
474+ dt <- reifyDataType name
475+ let preds = map (unboxPred . VarT) (fmap mkName vars)
476+ headTy = appsT (ConT name) (map VarT (dtTvs dt))
477+ cons = dtCons dt
478+ deriveUnboxInternal preds headTy cons
479+
480+ where
481+
482+ unboxPred ty =
483+#if MIN_VERSION_template_haskell(2,10,0)
484+ AppT (ConT ''Unbox) ty
485+#else
486+ ClassP ''Unbox [ty]
487+#endif