1 {-# OPTIONS -fno-warn-missing-signatures #-}
3 module VectType ( vectTyCon, vectAndLiftType, vectType, vectTypeEnv,
4 -- arrSumArity, pdataCompTys, pdataCompVars,
12 import Vectorise.Monad
13 import Vectorise.Builtins
14 import Vectorise.Type.Type
15 import Vectorise.Type.TyConDecl
17 import HscTypes ( TypeEnv, extendTypeEnvList, typeEnvTyCons )
22 import MkCore ( mkWildCase )
29 import FamInstEnv ( FamInst, mkLocalFamInst )
34 import Name ( Name, getOccName )
41 import Digraph ( SCC(..), stronglyConnCompFromEdgedVertices )
46 import MonadUtils ( zipWith3M, foldrM, concatMapM )
47 import Control.Monad ( liftM, liftM2, zipWithM, zipWithM_, mapAndUnzipM )
51 dtrace s x = if debug then pprTrace "VectType" s x else x
53 -- ----------------------------------------------------------------------------
57 -- ----------------------------------------------------------------------------
60 type TyConGroup = ([TyCon], UniqSet TyCon)
62 -- | Vectorise a type environment.
63 -- The type environment contains all the type things defined in a module.
64 vectTypeEnv :: TypeEnv -> VM (TypeEnv, [FamInst], [(Var, CoreExpr)])
68 cs <- readGEnv $ mk_map . global_tycons
70 -- Split the list of TyCons into the ones we have to vectorise vs the
71 -- ones we can pass through unchanged. We also pass through algebraic
72 -- types that use non Haskell98 features, as we don't handle those.
73 let (conv_tcs, keep_tcs) = classifyTyCons cs groups
74 keep_dcs = concatMap tyConDataCons keep_tcs
76 zipWithM_ defTyCon keep_tcs keep_tcs
77 zipWithM_ defDataCon keep_dcs keep_dcs
79 new_tcs <- vectTyConDecls conv_tcs
81 let orig_tcs = keep_tcs ++ conv_tcs
83 -- We don't need to make new representation types for dictionary
84 -- constructors. The constructors are always fully applied, and we don't
85 -- need to lift them to arrays as a dictionary of a particular type
86 -- always has the same value.
87 let vect_tcs = filter (not . isClassTyCon)
90 (_, binds, inst_tcs) <- fixV $ \ ~(dfuns', _, _) ->
92 defTyConPAs (zipLazy vect_tcs dfuns')
93 reprs <- mapM tyConRepr vect_tcs
94 repr_tcs <- zipWith3M buildPReprTyCon orig_tcs vect_tcs reprs
95 pdata_tcs <- zipWith3M buildPDataTyCon orig_tcs vect_tcs reprs
98 $ zipWith5 buildTyConBindings
106 return (dfuns, binds, repr_tcs ++ pdata_tcs)
108 let all_new_tcs = new_tcs ++ inst_tcs
110 let new_env = extendTypeEnvList env
111 (map ATyCon all_new_tcs
112 ++ [ADataCon dc | tc <- all_new_tcs
113 , dc <- tyConDataCons tc])
115 return (new_env, map mkLocalFamInst inst_tcs, binds)
117 tycons = typeEnvTyCons env
118 groups = tyConGroups tycons
120 mk_map env = listToUFM_Directly [(u, getUnique n /= u) | (u,n) <- nameEnvUniqueElts env]
123 mk_fam_inst :: TyCon -> TyCon -> (TyCon, [Type])
124 mk_fam_inst fam_tc arg_tc
125 = (fam_tc, [mkTyConApp arg_tc . mkTyVarTys $ tyConTyVars arg_tc])
128 buildPReprTyCon :: TyCon -> TyCon -> SumRepr -> VM TyCon
129 buildPReprTyCon orig_tc vect_tc repr
131 name <- cloneName mkPReprTyConOcc (tyConName orig_tc)
132 -- rhs_ty <- buildPReprType vect_tc
133 rhs_ty <- sumReprType repr
134 prepr_tc <- builtin preprTyCon
135 liftDs $ buildSynTyCon name
137 (SynonymTyCon rhs_ty)
139 (Just $ mk_fam_inst prepr_tc vect_tc)
141 tyvars = tyConTyVars vect_tc
143 data CompRepr = Keep Type
144 CoreExpr -- PR dictionary for the type
147 data ProdRepr = EmptyProd
149 | Prod { repr_tup_tc :: TyCon -- representation tuple tycon
150 , repr_ptup_tc :: TyCon -- PData representation tycon
151 , repr_comp_tys :: [Type] -- representation types of
152 , repr_comps :: [CompRepr] -- components
154 data ConRepr = ConRepr DataCon ProdRepr
156 data SumRepr = EmptySum
158 | Sum { repr_sum_tc :: TyCon -- representation sum tycon
159 , repr_psum_tc :: TyCon -- PData representation tycon
160 , repr_sel_ty :: Type -- type of selector
161 , repr_con_tys :: [Type] -- representation types of
162 , repr_cons :: [ConRepr] -- components
165 tyConRepr :: TyCon -> VM SumRepr
166 tyConRepr tc = sum_repr (tyConDataCons tc)
168 sum_repr [] = return EmptySum
169 sum_repr [con] = liftM UnarySum (con_repr con)
171 rs <- mapM con_repr cons
172 sum_tc <- builtin (sumTyCon arity)
173 tys <- mapM conReprType rs
174 (psum_tc, _) <- pdataReprTyCon (mkTyConApp sum_tc tys)
175 sel_ty <- builtin (selTy arity)
176 return $ Sum { repr_sum_tc = sum_tc
177 , repr_psum_tc = psum_tc
178 , repr_sel_ty = sel_ty
185 con_repr con = liftM (ConRepr con) (prod_repr (dataConRepArgTys con))
187 prod_repr [] = return EmptyProd
188 prod_repr [ty] = liftM UnaryProd (comp_repr ty)
190 rs <- mapM comp_repr tys
191 tup_tc <- builtin (prodTyCon arity)
192 tys' <- mapM compReprType rs
193 (ptup_tc, _) <- pdataReprTyCon (mkTyConApp tup_tc tys')
194 return $ Prod { repr_tup_tc = tup_tc
195 , repr_ptup_tc = ptup_tc
196 , repr_comp_tys = tys'
202 comp_repr ty = liftM (Keep ty) (prDictOfType ty)
203 `orElseV` return (Wrap ty)
205 sumReprType :: SumRepr -> VM Type
206 sumReprType EmptySum = voidType
207 sumReprType (UnarySum r) = conReprType r
208 sumReprType (Sum { repr_sum_tc = sum_tc, repr_con_tys = tys })
209 = return $ mkTyConApp sum_tc tys
211 conReprType :: ConRepr -> VM Type
212 conReprType (ConRepr _ r) = prodReprType r
214 prodReprType :: ProdRepr -> VM Type
215 prodReprType EmptyProd = voidType
216 prodReprType (UnaryProd r) = compReprType r
217 prodReprType (Prod { repr_tup_tc = tup_tc, repr_comp_tys = tys })
218 = return $ mkTyConApp tup_tc tys
220 compReprType :: CompRepr -> VM Type
221 compReprType (Keep ty _) = return ty
222 compReprType (Wrap ty) = do
223 wrap_tc <- builtin wrapTyCon
224 return $ mkTyConApp wrap_tc [ty]
226 compOrigType :: CompRepr -> Type
227 compOrigType (Keep ty _) = ty
228 compOrigType (Wrap ty) = ty
230 buildToPRepr :: TyCon -> TyCon -> TyCon -> SumRepr -> VM CoreExpr
231 buildToPRepr vect_tc repr_tc _ repr
233 let arg_ty = mkTyConApp vect_tc ty_args
234 res_ty <- mkPReprType arg_ty
235 arg <- newLocalVar (fsLit "x") arg_ty
236 result <- to_sum (Var arg) arg_ty res_ty repr
237 return $ Lam arg result
239 ty_args = mkTyVarTys (tyConTyVars vect_tc)
241 wrap_repr_inst = wrapFamInstBody repr_tc ty_args
243 to_sum _ _ _ EmptySum
245 void <- builtin voidVar
246 return $ wrap_repr_inst $ Var void
248 to_sum arg arg_ty res_ty (UnarySum r)
250 (pat, vars, body) <- con_alt r
251 return $ mkWildCase arg arg_ty res_ty
252 [(pat, vars, wrap_repr_inst body)]
254 to_sum arg arg_ty res_ty (Sum { repr_sum_tc = sum_tc
256 , repr_cons = cons })
258 alts <- mapM con_alt cons
259 let alts' = [(pat, vars, wrap_repr_inst
260 $ mkConApp sum_con (map Type tys ++ [body]))
261 | ((pat, vars, body), sum_con)
262 <- zip alts (tyConDataCons sum_tc)]
263 return $ mkWildCase arg arg_ty res_ty alts'
265 con_alt (ConRepr con r)
267 (vars, body) <- to_prod r
268 return (DataAlt con, vars, body)
272 void <- builtin voidVar
273 return ([], Var void)
275 to_prod (UnaryProd comp)
277 var <- newLocalVar (fsLit "x") (compOrigType comp)
278 body <- to_comp (Var var) comp
281 to_prod(Prod { repr_tup_tc = tup_tc
282 , repr_comp_tys = tys
283 , repr_comps = comps })
285 vars <- newLocalVars (fsLit "x") (map compOrigType comps)
286 exprs <- zipWithM to_comp (map Var vars) comps
287 return (vars, mkConApp tup_con (map Type tys ++ exprs))
289 [tup_con] = tyConDataCons tup_tc
291 to_comp expr (Keep _ _) = return expr
292 to_comp expr (Wrap ty) = do
293 wrap_tc <- builtin wrapTyCon
294 return $ wrapNewTypeBody wrap_tc [ty] expr
297 buildFromPRepr :: TyCon -> TyCon -> TyCon -> SumRepr -> VM CoreExpr
298 buildFromPRepr vect_tc repr_tc _ repr
300 arg_ty <- mkPReprType res_ty
301 arg <- newLocalVar (fsLit "x") arg_ty
303 result <- from_sum (unwrapFamInstScrut repr_tc ty_args (Var arg))
305 return $ Lam arg result
307 ty_args = mkTyVarTys (tyConTyVars vect_tc)
308 res_ty = mkTyConApp vect_tc ty_args
312 dummy <- builtin fromVoidVar
313 return $ Var dummy `App` Type res_ty
315 from_sum expr (UnarySum r) = from_con expr r
316 from_sum expr (Sum { repr_sum_tc = sum_tc
318 , repr_cons = cons })
320 vars <- newLocalVars (fsLit "x") tys
321 es <- zipWithM from_con (map Var vars) cons
322 return $ mkWildCase expr (exprType expr) res_ty
323 [(DataAlt con, [var], e)
324 | (con, var, e) <- zip3 (tyConDataCons sum_tc) vars es]
326 from_con expr (ConRepr con r)
327 = from_prod expr (mkConApp con $ map Type ty_args) r
329 from_prod _ con EmptyProd = return con
330 from_prod expr con (UnaryProd r)
332 e <- from_comp expr r
335 from_prod expr con (Prod { repr_tup_tc = tup_tc
336 , repr_comp_tys = tys
340 vars <- newLocalVars (fsLit "y") tys
341 es <- zipWithM from_comp (map Var vars) comps
342 return $ mkWildCase expr (exprType expr) res_ty
343 [(DataAlt tup_con, vars, con `mkApps` es)]
345 [tup_con] = tyConDataCons tup_tc
347 from_comp expr (Keep _ _) = return expr
348 from_comp expr (Wrap ty)
350 wrap <- builtin wrapTyCon
351 return $ unwrapNewTypeBody wrap [ty] expr
354 buildToArrPRepr :: TyCon -> TyCon -> TyCon -> SumRepr -> VM CoreExpr
355 buildToArrPRepr vect_tc prepr_tc pdata_tc r
357 arg_ty <- mkPDataType el_ty
358 res_ty <- mkPDataType =<< mkPReprType el_ty
359 arg <- newLocalVar (fsLit "xs") arg_ty
361 pdata_co <- mkBuiltinCo pdataTyCon
362 let Just repr_co = tyConFamilyCoercion_maybe prepr_tc
363 co = mkAppCoercion pdata_co
365 $ mkTyConApp repr_co ty_args
367 scrut = unwrapFamInstScrut pdata_tc ty_args (Var arg)
369 (vars, result) <- to_sum r
372 $ mkWildCase scrut (mkTyConApp pdata_tc ty_args) res_ty
373 [(DataAlt pdata_dc, vars, mkCoerce co result)]
375 ty_args = mkTyVarTys $ tyConTyVars vect_tc
376 el_ty = mkTyConApp vect_tc ty_args
378 [pdata_dc] = tyConDataCons pdata_tc
382 pvoid <- builtin pvoidVar
383 return ([], Var pvoid)
384 to_sum (UnarySum r) = to_con r
385 to_sum (Sum { repr_psum_tc = psum_tc
386 , repr_sel_ty = sel_ty
391 (vars, exprs) <- mapAndUnzipM to_con cons
392 sel <- newLocalVar (fsLit "sel") sel_ty
393 return (sel : concat vars, mk_result (Var sel) exprs)
395 [psum_con] = tyConDataCons psum_tc
396 mk_result sel exprs = wrapFamInstBody psum_tc tys
398 $ map Type tys ++ (sel : exprs)
400 to_con (ConRepr _ r) = to_prod r
402 to_prod EmptyProd = do
403 pvoid <- builtin pvoidVar
404 return ([], Var pvoid)
405 to_prod (UnaryProd r)
407 pty <- mkPDataType (compOrigType r)
408 var <- newLocalVar (fsLit "x") pty
409 expr <- to_comp (Var var) r
412 to_prod (Prod { repr_ptup_tc = ptup_tc
413 , repr_comp_tys = tys
414 , repr_comps = comps })
416 ptys <- mapM (mkPDataType . compOrigType) comps
417 vars <- newLocalVars (fsLit "x") ptys
418 es <- zipWithM to_comp (map Var vars) comps
419 return (vars, mk_result es)
421 [ptup_con] = tyConDataCons ptup_tc
422 mk_result exprs = wrapFamInstBody ptup_tc tys
424 $ map Type tys ++ exprs
426 to_comp expr (Keep _ _) = return expr
428 -- FIXME: this is bound to be wrong!
429 to_comp expr (Wrap ty)
431 wrap_tc <- builtin wrapTyCon
432 (pwrap_tc, _) <- pdataReprTyCon (mkTyConApp wrap_tc [ty])
433 return $ wrapNewTypeBody pwrap_tc [ty] expr
436 buildFromArrPRepr :: TyCon -> TyCon -> TyCon -> SumRepr -> VM CoreExpr
437 buildFromArrPRepr vect_tc prepr_tc pdata_tc r
439 arg_ty <- mkPDataType =<< mkPReprType el_ty
440 res_ty <- mkPDataType el_ty
441 arg <- newLocalVar (fsLit "xs") arg_ty
443 pdata_co <- mkBuiltinCo pdataTyCon
444 let Just repr_co = tyConFamilyCoercion_maybe prepr_tc
445 co = mkAppCoercion pdata_co
446 $ mkTyConApp repr_co var_tys
448 scrut = mkCoerce co (Var arg)
450 mk_result args = wrapFamInstBody pdata_tc var_tys
452 $ map Type var_tys ++ args
454 (expr, _) <- fixV $ \ ~(_, args) ->
455 from_sum res_ty (mk_result args) scrut r
457 return $ Lam arg expr
459 -- (args, mk) <- from_sum res_ty scrut r
461 -- let result = wrapFamInstBody pdata_tc var_tys
462 -- . mkConApp pdata_dc
463 -- $ map Type var_tys ++ args
465 -- return $ Lam arg (mk result)
467 var_tys = mkTyVarTys $ tyConTyVars vect_tc
468 el_ty = mkTyConApp vect_tc var_tys
470 [pdata_con] = tyConDataCons pdata_tc
472 from_sum _ res _ EmptySum = return (res, [])
473 from_sum res_ty res expr (UnarySum r) = from_con res_ty res expr r
474 from_sum res_ty res expr (Sum { repr_psum_tc = psum_tc
475 , repr_sel_ty = sel_ty
477 , repr_cons = cons })
479 sel <- newLocalVar (fsLit "sel") sel_ty
480 ptys <- mapM mkPDataType tys
481 vars <- newLocalVars (fsLit "xs") ptys
482 (res', args) <- fold from_con res_ty res (map Var vars) cons
483 let scrut = unwrapFamInstScrut psum_tc tys expr
484 body = mkWildCase scrut (exprType scrut) res_ty
485 [(DataAlt psum_con, sel : vars, res')]
486 return (body, Var sel : args)
488 [psum_con] = tyConDataCons psum_tc
491 from_con res_ty res expr (ConRepr _ r) = from_prod res_ty res expr r
493 from_prod _ res _ EmptyProd = return (res, [])
494 from_prod res_ty res expr (UnaryProd r)
495 = from_comp res_ty res expr r
496 from_prod res_ty res expr (Prod { repr_ptup_tc = ptup_tc
497 , repr_comp_tys = tys
498 , repr_comps = comps })
500 ptys <- mapM mkPDataType tys
501 vars <- newLocalVars (fsLit "ys") ptys
502 (res', args) <- fold from_comp res_ty res (map Var vars) comps
503 let scrut = unwrapFamInstScrut ptup_tc tys expr
504 body = mkWildCase scrut (exprType scrut) res_ty
505 [(DataAlt ptup_con, vars, res')]
508 [ptup_con] = tyConDataCons ptup_tc
510 from_comp _ res expr (Keep _ _) = return (res, [expr])
511 from_comp _ res expr (Wrap ty)
513 wrap_tc <- builtin wrapTyCon
514 (pwrap_tc, _) <- pdataReprTyCon (mkTyConApp wrap_tc [ty])
515 return (res, [unwrapNewTypeBody pwrap_tc [ty]
516 $ unwrapFamInstScrut pwrap_tc [ty] expr])
518 fold f res_ty res exprs rs = foldrM f' (res, []) (zip exprs rs)
520 f' (expr, r) (res, args) = do
521 (res', args') <- f res_ty res expr r
522 return (res', args' ++ args)
524 buildPRDict :: TyCon -> TyCon -> TyCon -> SumRepr -> VM CoreExpr
525 buildPRDict vect_tc prepr_tc _ r
528 pr_co <- mkBuiltinCo prTyCon
529 let co = mkAppCoercion pr_co
531 $ mkTyConApp arg_co ty_args
532 return (mkCoerce co dict)
534 ty_args = mkTyVarTys (tyConTyVars vect_tc)
535 Just arg_co = tyConFamilyCoercion_maybe prepr_tc
537 sum_dict EmptySum = prDFunOfTyCon =<< builtin voidTyCon
538 sum_dict (UnarySum r) = con_dict r
539 sum_dict (Sum { repr_sum_tc = sum_tc
544 dicts <- mapM con_dict cons
545 dfun <- prDFunOfTyCon sum_tc
546 return $ dfun `mkTyApps` tys `mkApps` dicts
548 con_dict (ConRepr _ r) = prod_dict r
550 prod_dict EmptyProd = prDFunOfTyCon =<< builtin voidTyCon
551 prod_dict (UnaryProd r) = comp_dict r
552 prod_dict (Prod { repr_tup_tc = tup_tc
553 , repr_comp_tys = tys
554 , repr_comps = comps })
556 dicts <- mapM comp_dict comps
557 dfun <- prDFunOfTyCon tup_tc
558 return $ dfun `mkTyApps` tys `mkApps` dicts
560 comp_dict (Keep _ pr) = return pr
561 comp_dict (Wrap ty) = wrapPR ty
564 buildPDataTyCon :: TyCon -> TyCon -> SumRepr -> VM TyCon
565 buildPDataTyCon orig_tc vect_tc repr = fixV $ \repr_tc ->
567 name' <- cloneName mkPDataTyConOcc orig_name
568 rhs <- buildPDataTyConRhs orig_name vect_tc repr_tc repr
569 pdata <- builtin pdataTyCon
571 liftDs $ buildAlgTyCon name'
573 [] -- no stupid theta
575 rec_flag -- FIXME: is this ok?
576 False -- FIXME: no generics
577 False -- not GADT syntax
578 (Just $ mk_fam_inst pdata vect_tc)
580 orig_name = tyConName orig_tc
581 tyvars = tyConTyVars vect_tc
582 rec_flag = boolToRecFlag (isRecursiveTyCon vect_tc)
585 buildPDataTyConRhs :: Name -> TyCon -> TyCon -> SumRepr -> VM AlgTyConRhs
586 buildPDataTyConRhs orig_name vect_tc repr_tc repr
588 data_con <- buildPDataDataCon orig_name vect_tc repr_tc repr
589 return $ DataTyCon { data_cons = [data_con], is_enum = False }
591 buildPDataDataCon :: Name -> TyCon -> TyCon -> SumRepr -> VM DataCon
592 buildPDataDataCon orig_name vect_tc repr_tc repr
594 dc_name <- cloneName mkPDataDataConOcc orig_name
595 comp_tys <- sum_tys repr
597 liftDs $ buildDataCon dc_name
599 (map (const HsNoBang) comp_tys)
600 [] -- no field labels
602 [] -- no existentials
606 (mkFamilyTyConApp repr_tc (mkTyVarTys tvs))
609 tvs = tyConTyVars vect_tc
611 sum_tys EmptySum = return []
612 sum_tys (UnarySum r) = con_tys r
613 sum_tys (Sum { repr_sel_ty = sel_ty
614 , repr_cons = cons })
615 = liftM (sel_ty :) (concatMapM con_tys cons)
617 con_tys (ConRepr _ r) = prod_tys r
619 prod_tys EmptyProd = return []
620 prod_tys (UnaryProd r) = liftM singleton (comp_ty r)
621 prod_tys (Prod { repr_comps = comps }) = mapM comp_ty comps
623 comp_ty r = mkPDataType (compOrigType r)
626 buildTyConBindings :: TyCon -> TyCon -> TyCon -> TyCon -> SumRepr
628 buildTyConBindings orig_tc vect_tc prepr_tc pdata_tc repr
630 vectDataConWorkers orig_tc vect_tc pdata_tc
631 buildPADict vect_tc prepr_tc pdata_tc repr
633 vectDataConWorkers :: TyCon -> TyCon -> TyCon -> VM ()
634 vectDataConWorkers orig_tc vect_tc arr_tc
637 . zipWith3 def_worker (tyConDataCons orig_tc) rep_tys
638 $ zipWith4 mk_data_con (tyConDataCons vect_tc)
641 (tail $ tails rep_tys)
642 mapM_ (uncurry hoistBinding) bs
644 tyvars = tyConTyVars vect_tc
645 var_tys = mkTyVarTys tyvars
646 ty_args = map Type var_tys
647 res_ty = mkTyConApp vect_tc var_tys
649 cons = tyConDataCons vect_tc
651 [arr_dc] = tyConDataCons arr_tc
653 rep_tys = map dataConRepArgTys $ tyConDataCons vect_tc
656 mk_data_con con tys pre post
657 = liftM2 (,) (vect_data_con con)
658 (lift_data_con tys pre post (mkDataConTag con))
660 sel_replicate len tag
662 rep <- builtin (selReplicate arity)
663 return [rep `mkApps` [len, tag]]
665 | otherwise = return []
667 vect_data_con con = return $ mkConApp con ty_args
668 lift_data_con tys pre_tys post_tys tag
670 len <- builtin liftingContext
671 args <- mapM (newLocalVar (fsLit "xs"))
672 =<< mapM mkPDataType tys
674 sel <- sel_replicate (Var len) tag
676 pre <- mapM emptyPD (concat pre_tys)
677 post <- mapM emptyPD (concat post_tys)
679 return . mkLams (len : args)
680 . wrapFamInstBody arr_tc var_tys
682 $ ty_args ++ sel ++ pre ++ map Var args ++ post
684 def_worker data_con arg_tys mk_body
686 arity <- polyArity tyvars
689 . polyAbstract tyvars $ \args ->
690 liftM (mkLams (tyvars ++ args) . vectorised)
691 $ buildClosures tyvars [] arg_tys res_ty mk_body
693 raw_worker <- cloneId mkVectOcc orig_worker (exprType body)
694 let vect_worker = raw_worker `setIdUnfolding`
695 mkInlineRule body (Just arity)
696 defGlobalVar orig_worker vect_worker
697 return (vect_worker, body)
699 orig_worker = dataConWorkId data_con
701 buildPADict :: TyCon -> TyCon -> TyCon -> SumRepr -> VM Var
702 buildPADict vect_tc prepr_tc arr_tc repr
703 = polyAbstract tvs $ \args ->
705 method_ids <- mapM (method args) paMethods
707 pa_tc <- builtin paTyCon
708 pa_dc <- builtin paDataCon
709 let dict = mkLams (tvs ++ args)
711 $ Type inst_ty : map (method_call args) method_ids
713 dfun_ty = mkForAllTys tvs
714 $ mkFunTys (map varType args) (mkTyConApp pa_tc [inst_ty])
716 raw_dfun <- newExportedVar dfun_name dfun_ty
717 let dfun = raw_dfun `setIdUnfolding` mkDFunUnfolding dfun_ty (map Var method_ids)
718 `setInlinePragma` dfunInlinePragma
720 hoistBinding dfun dict
723 tvs = tyConTyVars vect_tc
724 arg_tys = mkTyVarTys tvs
725 inst_ty = mkTyConApp vect_tc arg_tys
727 dfun_name = mkPADFunOcc (getOccName vect_tc)
729 method args (name, build)
732 expr <- build vect_tc prepr_tc arr_tc repr
733 let body = mkLams (tvs ++ args) expr
734 raw_var <- newExportedVar (method_name name) (exprType body)
736 `setIdUnfolding` mkInlineRule body (Just (length args))
737 `setInlinePragma` alwaysInlinePragma
738 hoistBinding var body
741 method_call args id = mkApps (Var id) (map Type arg_tys ++ map Var args)
743 method_name name = mkVarOcc $ occNameString dfun_name ++ ('$' : name)
746 paMethods :: [(String, TyCon -> TyCon -> TyCon -> SumRepr -> VM CoreExpr)]
747 paMethods = [("dictPRepr", buildPRDict),
748 ("toPRepr", buildToPRepr),
749 ("fromPRepr", buildFromPRepr),
750 ("toArrPRepr", buildToArrPRepr),
751 ("fromArrPRepr", buildFromArrPRepr)]
754 -- | Split the given tycons into two sets depending on whether they have to be
755 -- converted (first list) or not (second list). The first argument contains
756 -- information about the conversion status of external tycons:
758 -- * tycons which have converted versions are mapped to True
759 -- * tycons which are not changed by vectorisation are mapped to False
760 -- * tycons which can't be converted are not elements of the map
762 classifyTyCons :: UniqFM Bool -> [TyConGroup] -> ([TyCon], [TyCon])
763 classifyTyCons = classify [] []
765 classify conv keep _ [] = (conv, keep)
766 classify conv keep cs ((tcs, ds) : rs)
767 | can_convert && must_convert
768 = classify (tcs ++ conv) keep (cs `addListToUFM` [(tc,True) | tc <- tcs]) rs
770 = classify conv (tcs ++ keep) (cs `addListToUFM` [(tc,False) | tc <- tcs]) rs
772 = classify conv keep cs rs
774 refs = ds `delListFromUniqSet` tcs
776 can_convert = isNullUFM (refs `minusUFM` cs) && all convertable tcs
777 must_convert = foldUFM (||) False (intersectUFM_C const cs refs)
779 convertable tc = isDataTyCon tc && all isVanillaDataCon (tyConDataCons tc)
781 -- | Compute mutually recursive groups of tycons in topological order
783 tyConGroups :: [TyCon] -> [TyConGroup]
784 tyConGroups tcs = map mk_grp (stronglyConnCompFromEdgedVertices edges)
786 edges = [((tc, ds), tc, uniqSetToList ds) | tc <- tcs
787 , let ds = tyConsOfTyCon tc]
789 mk_grp (AcyclicSCC (tc, ds)) = ([tc], ds)
790 mk_grp (CyclicSCC els) = (tcs, unionManyUniqSets dss)
792 (tcs, dss) = unzip els
794 tyConsOfTyCon :: TyCon -> UniqSet TyCon
796 = tyConsOfTypes . concatMap dataConRepArgTys . tyConDataCons
798 tyConsOfType :: Type -> UniqSet TyCon
800 | Just ty' <- coreView ty = tyConsOfType ty'
801 tyConsOfType (TyVarTy _) = emptyUniqSet
802 tyConsOfType (TyConApp tc tys) = extend (tyConsOfTypes tys)
804 extend | isUnLiftedTyCon tc
805 || isTupleTyCon tc = id
807 | otherwise = (`addOneToUniqSet` tc)
809 tyConsOfType (AppTy a b) = tyConsOfType a `unionUniqSets` tyConsOfType b
810 tyConsOfType (FunTy a b) = (tyConsOfType a `unionUniqSets` tyConsOfType b)
811 `addOneToUniqSet` funTyCon
812 tyConsOfType (ForAllTy _ ty) = tyConsOfType ty
813 tyConsOfType other = pprPanic "ClosureConv.tyConsOfType" $ ppr other
815 tyConsOfTypes :: [Type] -> UniqSet TyCon
816 tyConsOfTypes = unionManyUniqSets . map tyConsOfType
819 -- ----------------------------------------------------------------------------
822 -- | Build an expression that calls the vectorised version of some
823 -- function from a `Closure`.
829 -- ($v_foo $: x) $: y
832 -- We use the type of the original binding to work out how many
833 -- outer lambdas to add.
836 :: Type -- ^ The type of the original binding.
837 -> CoreExpr -- ^ Expression giving the closure to use, eg @$v_foo@.
840 -- Convert the type to the core view if it isn't already.
842 | Just ty' <- coreView ty
845 -- For each function constructor in the original type we add an outer
846 -- lambda to bind the parameter variable, and an inner application of it.
847 fromVect (FunTy arg_ty res_ty) expr
849 arg <- newLocalVar (fsLit "x") arg_ty
850 varg <- toVect arg_ty (Var arg)
851 varg_ty <- vectType arg_ty
852 vres_ty <- vectType res_ty
853 apply <- builtin applyVar
854 body <- fromVect res_ty
855 $ Var apply `mkTyApps` [varg_ty, vres_ty] `mkApps` [expr, varg]
856 return $ Lam arg body
858 -- If the type isn't a function then it's time to call on the closure.
860 = identityConv ty >> return expr
863 toVect :: Type -> CoreExpr -> VM CoreExpr
864 toVect ty expr = identityConv ty >> return expr
867 identityConv :: Type -> VM ()
868 identityConv ty | Just ty' <- coreView ty = identityConv ty'
869 identityConv (TyConApp tycon tys)
871 mapM_ identityConv tys
872 identityConvTyCon tycon
875 identityConvTyCon :: TyCon -> VM ()
877 | isBoxedTupleTyCon tc = return ()
878 | isUnLiftedTyCon tc = return ()
880 tc' <- maybeV (lookupTyCon tc)
881 if tc == tc' then return () else noV