@@ -23,6 +23,7 @@ module Streamly.Internal.Data.Unbox.TH
2323
2424import Data.Word (Word16, Word32, Word64, Word8)
2525import Data.Proxy (Proxy(..))
26+import Data.List (elemIndex)
2627
2728import Language.Haskell.TH
2829import Language.Haskell.TH.Syntax
@@ -438,13 +439,31 @@ deriveUnbox mDecs = do
438439 case dec of
439440 [InstanceD mo preds headTyWC []] -> do
440441 let headTy = unwrap dec headTyWC
441− dt <- reifyDataType (getMainTypeName dec headTy)
442− let cons = dtCons dt
442+ (mainTyName, subs) = getMainTypeName dec headTy
443+ dt <- reifyDataType mainTyName
444+ let tyVars = dtTvs dt
445+ mapper = mapperWith (VarT <$> tyVars) subs
446+ cons = map (modifyConVariables mapper) (dtCons dt)
443447 deriveUnboxInternal headTy cons (mkInst mo preds headTyWC)
444448 _ -> errorMessage dec
445449
446450 where
447451
452+ mapperWith l1 l2 a =
453+ case elemIndex a l1 of
454+ Nothing -> a
455+ -- XXX Capture this case and give a relavant error.
456+ Just i -> l2 !! i
457+
458+ mapType f (AppT t1 t2) = AppT (mapType f t1) (mapType f t2)
459+ mapType f (InfixT t1 n t2) = InfixT (mapType f t1) n (mapType f t2)
460+ mapType f (UInfixT t1 n t2) = UInfixT (mapType f t1) n (mapType f t2)
461+ mapType f (ParensT t) = ParensT (mapType f t)
462+ mapType f v = f v
463+
464+ modifyConVariables f con =
465+ con { dcFields = map (\(a, b) -> (a, mapType f b)) (dcFields con) }
466+
448467 mkInst mo preds headTyWC methods =
449468 pure [InstanceD mo preds headTyWC methods]
450469
@@ -466,10 +485,10 @@ deriveUnbox mDecs = do
466485 unwrap _ (AppT (ConT _) r) = r
467486 unwrap dec _ = errorMessage dec
468487
469− getMainTypeName dec = go
488+ getMainTypeName dec = go []
470489
471490 where
472491
473− go (ConT nm) = nm
474− go (AppT l _) = go l
475− go _ = errorMessage dec
492+ go xs (ConT nm) = (nm, xs)
493+ go xs (AppT l r) = go (r:xs) l
494+ go _ _ = errorMessage dec