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