-- Vectorise a modules type and class declarations.---- This produces new type constructors and family instances top be included in the module toplevel-- as well as bindings for worker functions, dfuns, and the like.moduleVectorise.Type.Env(vectTypeEnv,)where#include "HsVersions.h"importVectorise.EnvimportVectorise.VectimportVectorise.MonadimportVectorise.BuiltinsimportVectorise.Type.TyConDeclimportVectorise.Type.ClassifyimportVectorise.Generic.PADictimportVectorise.Generic.PAMethodsimportVectorise.Generic.PDataimportVectorise.Generic.DescriptionimportVectorise.UtilsimportCoreSynimportCoreUtilsimportCoreUnfoldimportDataConimportTyConimportTypeimportFamInstEnvimportIdimportMkIdimportNameEnvimportNameSetimportOccNameimportUtilimportOutputableimportFastStringimportMonadUtilsimportControl.MonadimportData.MaybeimportData.List-- Note [Pragmas to vectorise tycons]-- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~---- VECTORISE pragmas for type constructors cover three different flavours of vectorising data type-- constructors:---- (1) Data type constructor 'T' that may be used in vectorised code, where 'T' represents itself,-- but the representation of 'T' is opaque in vectorised code. ---- An example is the treatment of 'Int'. 'Int's can be used in vectorised code and remain-- unchanged by vectorisation. However, the representation of 'Int' by the 'I#' data-- constructor wrapping an 'Int#' is not exposed in vectorised code. Instead, computations-- involving the representation need to be confined to scalar code.---- 'PData' and 'PRepr' instances need to be explicitly supplied for 'T' (they are not generated-- by the vectoriser).---- Type constructors declared with {-# VECTORISE SCALAR type T #-} are treated in this manner.-- (The vectoriser never treats a type constructor automatically in this manner.)---- (2) Data type constructor 'T' that may be used in vectorised code, where 'T' is represented by an-- explicitly given 'Tv', but the representation of 'T' is opaque in vectorised code. ---- An example is the treatment of '[::]'. '[::]'s can be used in vectorised code and is-- vectorised to 'PArray'. However, the representation of '[::]' is not exposed in vectorised-- code. Instead, computations involving the representation need to be confined to scalar code.---- 'PData' and 'PRepr' instances need to be explicitly supplied for 'T' (they are not generated-- by the vectoriser).---- Type constructors declared with {-# VECTORISE SCALAR type T = T' #-} are treated in this -- manner. (The vectoriser never treats a type constructor automatically in this manner.)---- (3) Data type constructor 'T' that together with its constructors 'Cn' may be used in vectorised-- code, where 'T' and the 'Cn' are automatically vectorised in the same manner as data types-- declared in a vectorised module. This includes the case where the vectoriser determines that-- the original representation of 'T' may be used in vectorised code (as it does not embed any-- parallel arrays.) This case is for type constructors that are *imported* from a non--- vectorised module, but that we want to use with full vectorisation support.---- An example is the treatment of 'Ordering' and '[]'. The former remains unchanged by-- vectorisation, whereas the latter is fully vectorised.-- 'PData' and 'PRepr' instances are automatically generated by the vectoriser.---- Type constructors declared with {-# VECTORISE type T #-} are treated in this manner.---- (4) Data type constructor 'T' that together with its constructors 'Cn' may be used in vectorised-- code, where 'T' is represented by an explicitly given 'Tv' whose constructors 'Cvn' represent-- the original constructors in vectorised code. As a special case, we can have 'Tv = T'---- An example is the treatment of 'Bool', which is represented by itself in vectorised code-- (as it cannot embed any parallel arrays). However, we do not want any automatic generation-- of class and family instances, which is why Case (3) does not apply.---- 'PData' and 'PRepr' instances need to be explicitly supplied for 'T' (they are not generated-- by the vectoriser).---- Type constructors declared with {-# VECTORISE type T = T' #-} are treated in this manner.---- In addition, we have also got a single pragma form for type classes: {-# VECTORISE class C #-}.-- It implies that the class type constructor may be used in vectorised code together with its data-- constructor. We generally produce a vectorised version of the data type and data constructor.-- We do not generate 'PData' and 'PRepr' instances for class type constructors. This pragma is the-- default for all type classes declared in this module, but the pragma can also be used explitly on-- imported classes.-- Note [Vectorising classes]-- ~~~~~~~~~~~~~~~~~~~~~~~~~~---- We vectorise classes essentially by just vectorising their desugared Core representation, but we-- do generate a 'Class' structure along the way (see 'Vectorise.Type.TyConDecl.vectTyConDecl').---- Here is an example illustrating the mapping — assume---- class Num a where-- (+) :: a -> a -> a---- It desugars to---- data Num a = D:Num { (+) :: a -> a -> a }---- which we vectorise to---- data V:Num a = D:V:Num { ($v+) :: PArray a :-> PArray a :-> PArray a }---- while adding the following entries to the vectorisation map:---- tycon : Num --> V:Num-- datacon: D:Num --> D:V:Num-- var : (+) --> ($v+)-- |Vectorise type constructor including class type constructors.--vectTypeEnv::[TyCon]-- Type constructors defined in this module->[CoreVect]-- All 'VECTORISE [SCALAR] type' declarations in this module->[CoreVect]-- All 'VECTORISE class' declarations in this module->VM([TyCon]-- old TyCons ++ new TyCons,[FamInst]-- New type family instances.,[(Var,CoreExpr)])-- New top level bindings.vectTypeEnvtyconsvectTypeDeclsvectClassDecls=do{traceVt"** vectTypeEnv"$pprtycons;let-- {-# VECTORISE SCALAR type T -#} (imported and local tycons)localAbstractTyCons=[tycon|VectTypeTruetyconNothing<-vectTypeDecls]-- {-# VECTORISE type T -#} (ONLY the imported tycons)impVectTyCons=([tycon|VectTypeFalsetyconNothing<-vectTypeDecls]++[tycon|VectClasstycon<-vectClassDecls])\\tycons-- {-# VECTORISE [SCALAR] type T = T' -#} (imported and local tycons)vectTyConsWithRHS=[(tycon,rhs,isAbstract)|VectTypeisAbstracttycon(Justrhs)<-vectTypeDecls]-- filter VECTORISE SCALAR tycons and VECTORISE tycons with explicit rhsesvectSpecialTyConNames=mkNameSet.maptyConName$localAbstractTyCons++mapfst3vectTyConsWithRHSnotVectSpecialTyContc=not$(tyConNametc)`elemNameSet`vectSpecialTyConNames-- Build a map containing all vectorised type constructor. If they are scalar, they are-- mapped to 'False' (vectorised type constructor == original type constructor).;allScalarTyConNames<-globalScalarTyCons-- covers both current and imported modules;vectTyCons<-globalVectTyCons;letvectTyConBase=mapNameEnv(constTrue)vectTyCons-- by default fully vectorisedvectTyConFlavour=vectTyConBase`plusNameEnv`mkNameEnv[(tyConNametycon,True)|(tycon,_,_)<-vectTyConsWithRHS]`plusNameEnv`mkNameEnv[(tcName,False)-- original representation|tcName<-nameSetToListallScalarTyConNames]`plusNameEnv`mkNameEnv[(tyConNametycon,False)-- original representation|tycon<-localAbstractTyCons]-- Split the list of 'TyCons' into the ones (1) that we must vectorise and those (2)-- that we could, but don't need to vectorise. Type constructors that are not data-- type constructors or use non-Haskell98 features are being dropped. They may not-- appear in vectorised code. (We also drop the local type constructors appearing in a-- VECTORISE SCALAR pragma or a VECTORISE pragma with an explicit right-hand side, as-- these are being handled separately. NB: Some type constructors may be marked SCALAR-- /and/ have an explicit right-hand side.)---- Furthermore, 'drop_tcs' are those type constructors that we cannot vectorise.;letmaybeVectoriseTyCons=filternotVectSpecialTyContycons++impVectTyCons(conv_tcs,keep_tcs,drop_tcs)=classifyTyConsvectTyConFlavourmaybeVectoriseTyCons;traceVt" VECT SCALAR : "$pprlocalAbstractTyCons;traceVt" VECT [class] : "$pprimpVectTyCons;traceVt" VECT with rhs : "$ppr(mapfst3vectTyConsWithRHS);traceVt" -- after classification (local and VECT [class] tycons) --"empty;traceVt" reuse : "$pprkeep_tcs;traceVt" convert : "$pprconv_tcs-- warn the user about unvectorised type constructors;letexplanation=ptext(sLit"(They use unsupported language extensions")$$ptext(sLit"or depend on type constructors that are not vectorised)")drop_tcs_nosyn=filter(not.isSynTyCon)drop_tcs;unless(nulldrop_tcs_nosyn)$emitVt"Warning: cannot vectorise these type constructors:"$pprQuotedListdrop_tcs_nosyn$$explanation;mapM_addGlobalScalarTyConkeep_tcs;letmapping=-- Type constructors that we don't need to vectorise, use the same-- representation in both unvectorised and vectorised code; they are not-- abstract.[(tycon,tycon,False)|tycon<-keep_tcs]-- We do the same for type constructors declared VECTORISE SCALAR /without/-- an explicit right-hand side, but ignore their representation (data-- constructors) as they are abstract.++[(tycon,tycon,True)|tycon<-localAbstractTyCons]-- Type constructors declared VECTORISE /with/ an explicit vectorised type,-- we map from the original to the given type; whether they are abstract depends-- on whether the vectorisation declaration was SCALAR.++vectTyConsWithRHS;syn_tcs<-catMaybes<$>mapMdefTyConDataConsmapping-- Vectorise all the data type declarations that we can and must vectorise (enter the-- type and data constructors into the vectorisation map on-the-fly.);new_tcs<-vectTyConDeclsconv_tcs;letdumpTctcvTc=traceVt"---"(pprtc<+>text"::"<+>ppr(dataConSigtc)$$pprvTc<+>text"::"<+>ppr(dataConSigvTc))dataConSigtc|Justdc<-tyConSingleDataCon_maybetc=dataConRepTypedc|otherwise=panic"dataConSig";zipWithM_dumpTc(filterisClassTyConconv_tcs)(filterisClassTyConnew_tcs)-- We don't need new representation types for dictionary constructors. The constructors-- are always fully applied, and we don't need to lift them to arrays as a dictionary-- of a particular type always has the same value.;letorig_tcs=filter(not.isClassTyCon)$keep_tcs++conv_tcsvect_tcs=filter(not.isClassTyCon)$keep_tcs++new_tcs-- Build 'PRepr' and 'PData' instance type constructors and family instances for all-- type constructors with vectorised representations.;reprs<-mapMtyConReprvect_tcs;repr_fis<-zipWith3MbuildPReprTyConorig_tcsvect_tcsreprs;pdata_fis<-zipWith3MbuildPDataTyConorig_tcsvect_tcsreprs;pdatas_fis<-zipWith3MbuildPDatasTyConorig_tcsvect_tcsreprs;letfam_insts=repr_fis++pdata_fis++pdatas_fisrepr_axs=mapfamInstAxiomrepr_fispdata_tcs=famInstsRepTyConspdata_fispdatas_tcs=famInstsRepTyConspdatas_fis;updGEnv$extendFamEnvfam_insts-- Generate workers for the vectorised data constructors, dfuns for the 'PA' instances of-- the vectorised type constructors, and associate the type constructors with their dfuns-- in the global environment. We get back the dfun bindings (which we will subsequently-- inject into the modules toplevel).;(_,binds)<-fixV$\~(dfuns,_)->do{defTyConPAs(zipLazyvect_tcsdfuns)-- Query the 'PData' instance type constructors for type constructors that have a-- VECTORISE pragma with an explicit right-hand side (this is Item (4) of-- "Note [Pragmas to vectorise tycons]" above).;let(withRHS_non_abstract,vwithRHS_non_abstract)=unzip[(tycon,vtycon)|(tycon,vtycon,False)<-vectTyConsWithRHS];pdata_withRHS_tcs<-mapMpdataReprTyConExactwithRHS_non_abstract-- Build workers for all vectorised data constructors (except abstract ones);sequence_$zipWith3vectDataConWorkers(orig_tcs++withRHS_non_abstract)(vect_tcs++vwithRHS_non_abstract)(pdata_tcs++pdata_withRHS_tcs)-- Build a 'PA' dictionary for all type constructors (except abstract ones & those-- defined with an explicit right-hand side where the dictionary is user-supplied);dfuns<-sequence$zipWith4buildTyConPADictvect_tcsrepr_axspdata_tcspdatas_tcs;binds<-takeHoisted;return(dfuns,binds)}-- Return the vectorised variants of type constructors as well as the generated instance-- type constructors, family instances, and dfun bindings.;return(new_tcs++pdata_tcs++pdatas_tcs++syn_tcs,fam_insts,binds)}wherefst3(a,_,_)=a-- Add a mapping from the original to vectorised type constructor to the vectorisation map. -- Unless the type constructor is abstract, also mappings from the orignal's data constructors-- to the vectorised type's data constructors.---- We have three cases: (1) original and vectorised type constructor are the same, (2) the-- name of the vectorised type constructor is canonical (as prescribed by 'mkVectTyConOcc'), or-- (3) the name is not canonical. In the third case, we additionally introduce a type synonym-- with the canonical name that is set equal to the non-canonical name (so that we find the-- right type constructor when reading vectorisation information from interface files).--defTyConDataCons(origTyCon,vectTyCon,isAbstract)=do{canonName<-mkLocalisedNamemkVectTyConOccorigName;iforigName==vectName-- Case (1)||vectName==canonName-- Case (2)thendo{defTyConorigTyConvectTyCon-- T --> vT;defDataCons-- Ci --> vCi;returnNothing}elsedo-- Case (3){letsynTyCon=mkSyncanonName(mkTyConTyvectTyCon)-- type S = vT;defTyConorigTyConsynTyCon-- T --> S;defDataCons-- Ci --> vCi;return$JustsynTyCon}}whereorigName=tyConNameorigTyConvectName=tyConNamevectTyConmkSyncanonNamety=mkSynTyConcanonName(typeKindty)[](SynonymTyConty)NoParentTyCondefDataCons|isAbstract=return()|otherwise=do{MASSERT(length(tyConDataConsorigTyCon)==length(tyConDataConsvectTyCon));zipWithM_defDataCon(tyConDataConsorigTyCon)(tyConDataConsvectTyCon)}-- Helpers --------------------------------------------------------------------buildTyConPADict::TyCon->CoAxiom->TyCon->TyCon->VMVarbuildTyConPADictvect_tcprepr_axpdata_tcpdatas_tc=tyConReprvect_tc>>=buildPADictvect_tcprepr_axpdata_tcpdatas_tc-- Produce a custom-made worker for the data constructors of a vectorised data type. This includes-- all data constructors that may be used in vetcorised code — i.e., all data constructors of data-- types other than scalar ones. Also adds a mapping from the original to vectorised worker into-- the vectorisation map.---- FIXME: It's not nice that we need create a special worker after the data constructors has-- already been constructed. Also, I don't think the worker is properly added to the data-- constructor. Seems messy.vectDataConWorkers::TyCon->TyCon->TyCon->VM()vectDataConWorkersorig_tcvect_tcarr_tc=do{traceVt"Building vectorised worker for datatype"(pprorig_tc);bs<-sequence.zipWith3def_worker(tyConDataConsorig_tc)rep_tys$zipWith4mk_data_con(tyConDataConsvect_tc)rep_tys(initsrep_tys)(tail$tailsrep_tys);mapM_(uncurryhoistBinding)bs}wheretyvars=tyConTyVarsvect_tcvar_tys=mkTyVarTystyvarsty_args=mapTypevar_tysres_ty=mkTyConAppvect_tcvar_tyscons=tyConDataConsvect_tcarity=lengthcons[arr_dc]=tyConDataConsarr_tcrep_tys=mapdataConRepArgTys$tyConDataConsvect_tcmk_data_concontysprepost=liftM2(,)(vect_data_concon)(lift_data_contysprepost(mkDataConTagcon))sel_replicatelentag|arity>1=dorep<-builtin(selReplicatearity)return[rep`mkApps`[len,tag]]|otherwise=return[]vect_data_concon=return$mkConAppconty_argslift_data_contyspre_tyspost_tystag=dolen<-builtinliftingContextargs<-mapM(newLocalVar(fsLit"xs"))=<<mapMmkPDataTypetyssel<-sel_replicate(Varlen)tagpre<-mapMemptyPD(concatpre_tys)post<-mapMemptyPD(concatpost_tys)return.mkLams(len:args).wrapFamInstBodyarr_tcvar_tys.mkConApparr_dc$ty_args++sel++pre++mapVarargs++postdef_workerdata_conarg_tysmk_body=doarity<-polyAritytyvarsbody<-closedV.inBindorig_worker.polyAbstracttyvars$\args->liftM(mkLams(tyvars++args).vectorised)$buildClosurestyvars[][]arg_tysres_tymk_bodyraw_worker<-mkVectIdorig_worker(exprTypebody)letvect_worker=raw_worker`setIdUnfolding`mkInlineUnfolding(Justarity)bodydefGlobalVarorig_workervect_workerreturn(vect_worker,body)whereorig_worker=dataConWorkIddata_con