- tvs = tyConTyVars vect_tc
- arg_tys = mkTyVarTys tvs
- res_ty = mkTyConApp vect_tc arg_tys
-
- orig_worker = dataConWorkId orig_dc
-
- mk_vect = return . mkConApp vect_dc $ map Type arg_tys
- mk_lift = do
- len <- newLocalVar FSLIT("n") intPrimTy
- arr_tys <- mapM mkPArrayType dc_tys
- args <- mapM (newLocalVar FSLIT("xs")) arr_tys
- shapes <- shapeReplicate shape
- (Var len)
- (mkDataConTag vect_dc)
-
- empty_pre <- mapM emptyPA (concat pre)
- empty_post <- mapM emptyPA (concat post)
-
- return . mkLams (len : args)
- . wrapFamInstBody arr_tc arg_tys
- . mkConApp arr_dc
- $ map Type arg_tys ++ shapes
- ++ empty_pre
- ++ map Var args
- ++ empty_post
+ tyvars = tyConTyVars vect_tc
+ var_tys = mkTyVarTys tyvars
+ ty_args = map Type var_tys
+
+ res_ty = mkTyConApp vect_tc var_tys
+
+ rep_tys = map dataConRepArgTys $ tyConDataCons vect_tc
+
+ [arr_dc] = tyConDataCons arr_tc
+
+ mk_data_con con tys pre post
+ = liftM2 (,) (vect_data_con con)
+ (lift_data_con tys (concat pre)
+ (concat post)
+ (mkDataConTag con))
+
+ vect_data_con con = return $ mkConApp con ty_args
+ lift_data_con tys pre_tys post_tys tag
+ = do
+ len <- builtin liftingContext
+ args <- mapM (newLocalVar FSLIT("xs"))
+ =<< mapM mkPArrayType tys
+
+ shape <- replicateShape repr (Var len) tag
+ repr <- mk_arr_repr (Var len) (map Var args)
+
+ pre <- mapM emptyPA pre_tys
+ post <- mapM emptyPA post_tys
+
+ return . mkLams (len : args)
+ . wrapFamInstBody arr_tc var_tys
+ . mkConApp arr_dc
+ $ ty_args ++ shape ++ pre ++ repr ++ post
+
+ mk_arr_repr len []
+ = do
+ units <- replicatePA len (Var unitDataConId)
+ return [units]
+
+ mk_arr_repr len arrs = return arrs
+
+ def_worker data_con arg_tys mk_body
+ = do
+ body <- closedV
+ . inBind orig_worker
+ . polyAbstract tyvars $ \abstract ->
+ liftM (abstract . vectorised)
+ $ buildClosures tyvars [] arg_tys res_ty mk_body
+
+ vect_worker <- cloneId mkVectOcc orig_worker (exprType body)
+ defGlobalVar orig_worker vect_worker
+ return (vect_worker, body)
+ where
+ orig_worker = dataConWorkId data_con