compiler/vectorise/VectType.hs

   1 module VectType ( vectTyCon, vectAndLiftType, vectType, vectTypeEnv,
   2                   -- arrSumArity, pdataCompTys, pdataCompVars,
   3                   buildPADict,
   4                   fromVect )
   5 where
   6
   7 import VectMonad
   8 import VectUtils
   9 import VectCore
  10
  11 import HscTypes          ( TypeEnv, extendTypeEnvList, typeEnvTyCons )
  12 import CoreSyn
  13 import CoreUtils
  14 import MkCore            ( mkWildCase )
  15 import BuildTyCl
  16 import DataCon
  17 import TyCon
  18 import Type
  19 import TypeRep
  20 import Coercion
  21 import FamInstEnv        ( FamInst, mkLocalFamInst )
  22 import OccName
  23 import MkId
  24 import BasicTypes        ( StrictnessMark(..), boolToRecFlag )
  25 import Var               ( Var, TyVar )
  26 import Name              ( Name, getOccName )
  27 import NameEnv
  28
  29 import Unique
  30 import UniqFM
  31 import UniqSet
  32 import Util
  33 import Digraph           ( SCC(..), stronglyConnCompFromEdgedVertices )
  34
  35 import Outputable
  36 import FastString
  37
  38 import Control.Monad  ( liftM, liftM2, zipWithM, zipWithM_, mapAndUnzipM )
  39 import Data.List      ( inits, tails, zipWith4, zipWith5 )
  40
  41 -- ----------------------------------------------------------------------------
  42 -- Types
  43
  44 vectTyCon :: TyCon -> VM TyCon
  45 vectTyCon tc
  46   | isFunTyCon tc        = builtin closureTyCon
  47   | isBoxedTupleTyCon tc = return tc
  48   | isUnLiftedTyCon tc   = return tc
  49   | otherwise            = maybeCantVectoriseM "Tycon not vectorised:" (ppr tc)
  50                          $ lookupTyCon tc
  51
  52 vectAndLiftType :: Type -> VM (Type, Type)
  53 vectAndLiftType ty | Just ty' <- coreView ty = vectAndLiftType ty'
  54 vectAndLiftType ty
  55   = do
  56       mdicts   <- mapM paDictArgType tyvars
  57       let dicts = [dict | Just dict <- mdicts]
  58       vmono_ty <- vectType mono_ty
  59       lmono_ty <- mkPDataType vmono_ty
  60       return (abstractType tyvars dicts vmono_ty,
  61               abstractType tyvars dicts lmono_ty)
  62   where
  63     (tyvars, mono_ty) = splitForAllTys ty
  64
  65
  66 vectType :: Type -> VM Type
  67 vectType ty | Just ty' <- coreView ty = vectType ty'
  68 vectType (TyVarTy tv) = return $ TyVarTy tv
  69 vectType (AppTy ty1 ty2) = liftM2 AppTy (vectType ty1) (vectType ty2)
  70 vectType (TyConApp tc tys) = liftM2 TyConApp (vectTyCon tc) (mapM vectType tys)
  71 vectType (FunTy ty1 ty2)   = liftM2 TyConApp (builtin closureTyCon)
  72                                              (mapM vectAndBoxType [ty1,ty2])
  73 vectType ty@(ForAllTy _ _)
  74   = do
  75       mdicts   <- mapM paDictArgType tyvars
  76       mono_ty' <- vectType mono_ty
  77       return $ abstractType tyvars [dict | Just dict <- mdicts] mono_ty'
  78   where
  79     (tyvars, mono_ty) = splitForAllTys ty
  80
  81 vectType ty = cantVectorise "Can't vectorise type" (ppr ty)
  82
  83 vectAndBoxType :: Type -> VM Type
  84 vectAndBoxType ty = vectType ty >>= boxType
  85
  86 abstractType :: [TyVar] -> [Type] -> Type -> Type
  87 abstractType tyvars dicts = mkForAllTys tyvars . mkFunTys dicts
  88
  89 -- ----------------------------------------------------------------------------
  90 -- Boxing
  91
  92 boxType :: Type -> VM Type
  93 boxType ty
  94   | Just (tycon, []) <- splitTyConApp_maybe ty
  95   , isUnLiftedTyCon tycon
  96   = do
  97       r <- lookupBoxedTyCon tycon
  98       case r of
  99         Just tycon' -> return $ mkTyConApp tycon' []
 100         Nothing     -> return ty
 101 boxType ty = return ty
 102
 103 -- ----------------------------------------------------------------------------
 104 -- Type definitions
 105
 106 type TyConGroup = ([TyCon], UniqSet TyCon)
 107
 108 vectTypeEnv :: TypeEnv -> VM (TypeEnv, [FamInst], [(Var, CoreExpr)])
 109 vectTypeEnv env
 110   = do
 111       cs <- readGEnv $ mk_map . global_tycons
 112       let (conv_tcs, keep_tcs) = classifyTyCons cs groups
 113           keep_dcs             = concatMap tyConDataCons keep_tcs
 114       zipWithM_ defTyCon   keep_tcs keep_tcs
 115       zipWithM_ defDataCon keep_dcs keep_dcs
 116       new_tcs <- vectTyConDecls conv_tcs
 117
 118       let orig_tcs = keep_tcs ++ conv_tcs
 119           vect_tcs = keep_tcs ++ new_tcs
 120
 121       repr_tcs  <- zipWithM buildPReprTyCon orig_tcs vect_tcs
 122       pdata_tcs <- zipWithM buildPDataTyCon orig_tcs vect_tcs
 123       dfuns     <- mapM mkPADFun vect_tcs
 124       defTyConPAs (zip vect_tcs dfuns)
 125       binds    <- sequence (zipWith5 buildTyConBindings orig_tcs
 126                                                         vect_tcs
 127                                                         repr_tcs
 128                                                         pdata_tcs
 129                                                         dfuns)
 130
 131       let all_new_tcs = new_tcs ++ repr_tcs ++ pdata_tcs
 132
 133       let new_env = extendTypeEnvList env
 134                        (map ATyCon all_new_tcs
 135                         ++ [ADataCon dc | tc <- all_new_tcs
 136                                         , dc <- tyConDataCons tc])
 137
 138       return (new_env, map mkLocalFamInst (repr_tcs ++ pdata_tcs), concat binds)
 139   where
 140     tycons = typeEnvTyCons env
 141     groups = tyConGroups tycons
 142
 143     mk_map env = listToUFM_Directly [(u, getUnique n /= u) | (u,n) <- nameEnvUniqueElts env]
 144
 145
 146 vectTyConDecls :: [TyCon] -> VM [TyCon]
 147 vectTyConDecls tcs = fixV $ \tcs' ->
 148   do
 149     mapM_ (uncurry defTyCon) (zipLazy tcs tcs')
 150     mapM vectTyConDecl tcs
 151
 152 vectTyConDecl :: TyCon -> VM TyCon
 153 vectTyConDecl tc
 154   = do
 155       name' <- cloneName mkVectTyConOcc name
 156       rhs'  <- vectAlgTyConRhs tc (algTyConRhs tc)
 157
 158       liftDs $ buildAlgTyCon name'
 159                              tyvars
 160                              []           -- no stupid theta
 161                              rhs'
 162                              rec_flag     -- FIXME: is this ok?
 163                              False        -- FIXME: no generics
 164                              False        -- not GADT syntax
 165                              Nothing      -- not a family instance
 166   where
 167     name   = tyConName tc
 168     tyvars = tyConTyVars tc
 169     rec_flag = boolToRecFlag (isRecursiveTyCon tc)
 170
 171 vectAlgTyConRhs :: TyCon -> AlgTyConRhs -> VM AlgTyConRhs
 172 vectAlgTyConRhs _ (DataTyCon { data_cons = data_cons
 173                              , is_enum   = is_enum
 174                              })
 175   = do
 176       data_cons' <- mapM vectDataCon data_cons
 177       zipWithM_ defDataCon data_cons data_cons'
 178       return $ DataTyCon { data_cons = data_cons'
 179                          , is_enum   = is_enum
 180                          }
 181 vectAlgTyConRhs tc _ = cantVectorise "Can't vectorise type definition:" (ppr tc)
 182
 183 vectDataCon :: DataCon -> VM DataCon
 184 vectDataCon dc
 185   | not . null $ dataConExTyVars dc
 186         = cantVectorise "Can't vectorise constructor (existentials):" (ppr dc)
 187   | not . null $ dataConEqSpec   dc
 188         = cantVectorise "Can't vectorise constructor (eq spec):" (ppr dc)
 189   | otherwise
 190   = do
 191       name'    <- cloneName mkVectDataConOcc name
 192       tycon'   <- vectTyCon tycon
 193       arg_tys  <- mapM vectType rep_arg_tys
 194
 195       liftDs $ buildDataCon name'
 196                             False           -- not infix
 197                             (map (const NotMarkedStrict) arg_tys)
 198                             []              -- no labelled fields
 199                             univ_tvs
 200                             []              -- no existential tvs for now
 201                             []              -- no eq spec for now
 202                             []              -- no context
 203                             arg_tys
 204                             (mkFamilyTyConApp tycon' (mkTyVarTys univ_tvs))
 205                             tycon'
 206   where
 207     name        = dataConName dc
 208     univ_tvs    = dataConUnivTyVars dc
 209     rep_arg_tys = dataConRepArgTys dc
 210     tycon       = dataConTyCon dc
 211
 212 mk_fam_inst :: TyCon -> TyCon -> (TyCon, [Type])
 213 mk_fam_inst fam_tc arg_tc
 214   = (fam_tc, [mkTyConApp arg_tc . mkTyVarTys $ tyConTyVars arg_tc])
 215
 216 buildPReprTyCon :: TyCon -> TyCon -> VM TyCon
 217 buildPReprTyCon orig_tc vect_tc
 218   = do
 219       name     <- cloneName mkPReprTyConOcc (tyConName orig_tc)
 220       rhs_ty   <- buildPReprType vect_tc
 221       prepr_tc <- builtin preprTyCon
 222       liftDs $ buildSynTyCon name
 223                              tyvars
 224                              (SynonymTyCon rhs_ty)
 225                              (typeKind rhs_ty)
 226                              (Just $ mk_fam_inst prepr_tc vect_tc)
 227   where
 228     tyvars = tyConTyVars vect_tc
 229
 230 buildPReprType :: TyCon -> VM Type
 231 buildPReprType vect_tc = sum_type . map dataConRepArgTys $ tyConDataCons vect_tc
 232   where
 233     sum_type []    = voidType
 234     sum_type [tys] = prod_type tys
 235     sum_type _     = do
 236                        (sum_tc, _, _, args) <- reprSumTyCons vect_tc
 237                        return $ mkTyConApp sum_tc args
 238
 239     prod_type []   = voidType
 240     prod_type [ty] = return ty
 241     prod_type tys  = do
 242                        prod_tc <- builtin (prodTyCon (length tys))
 243                        return $ mkTyConApp prod_tc tys
 244
 245 reprSumTyCons :: TyCon -> VM (TyCon, TyCon, Type, [Type])
 246 reprSumTyCons vect_tc
 247   = do
 248       tc   <- builtin (sumTyCon arity)
 249       args <- mapM (prod . dataConRepArgTys) cons
 250       (pdata_tc, _) <- pdataReprTyCon (mkTyConApp tc args)
 251       sel_ty <- builtin (selTy arity)
 252       return (tc, pdata_tc, sel_ty, args)
 253   where
 254     cons = tyConDataCons vect_tc
 255     arity = length cons
 256
 257     prod []   = voidType
 258     prod [ty] = return ty
 259     prod tys  = do
 260                   prod_tc <- builtin (prodTyCon (length tys))
 261                   return $ mkTyConApp prod_tc tys
 262
 263 buildToPRepr :: TyCon -> TyCon -> TyCon -> VM CoreExpr
 264 buildToPRepr vect_tc repr_tc _
 265   = do
 266       let arg_ty = mkTyConApp vect_tc ty_args
 267       res_ty <- mkPReprType arg_ty
 268       arg    <- newLocalVar (fsLit "x") arg_ty
 269       result <- to_sum (Var arg) arg_ty res_ty (tyConDataCons vect_tc)
 270       return $ Lam arg result
 271   where
 272     ty_args = mkTyVarTys (tyConTyVars vect_tc)
 273
 274     wrap = wrapFamInstBody repr_tc ty_args
 275
 276     to_sum _ _ _ []
 277       = do
 278           void <- builtin voidVar
 279           return $ wrap (Var void)
 280
 281     to_sum arg arg_ty res_ty [con]
 282       = do
 283           (prod, vars) <- to_prod (dataConRepArgTys con)
 284           return $ mkWildCase arg arg_ty res_ty
 285                    [(DataAlt con, vars, wrap prod)]
 286
 287     to_sum arg arg_ty res_ty cons
 288       = do
 289           (prods, vars) <- mapAndUnzipM (to_prod . dataConRepArgTys) cons
 290           (sum_tc, _, _, sum_ty_args) <- reprSumTyCons vect_tc
 291           let sum_cons = [mkConApp con (map Type sum_ty_args)
 292                             | con <- tyConDataCons sum_tc]
 293           return . mkWildCase arg arg_ty res_ty
 294                  $ zipWith4 mk_alt cons vars sum_cons prods
 295       where
 296         mk_alt con vars sum_con expr
 297           = (DataAlt con, vars, wrap $ sum_con `App` expr)
 298
 299     to_prod []
 300       = do
 301           void <- builtin voidVar
 302           return (Var void, [])
 303     to_prod [ty]
 304       = do
 305           var <- newLocalVar (fsLit "x") ty
 306           return (Var var, [var])
 307     to_prod tys
 308       = do
 309           prod_con <- builtin (prodDataCon (length tys))
 310           vars <- newLocalVars (fsLit "x") tys
 311           return (mkConApp prod_con (map Type tys ++ map Var vars), vars)
 312
 313 buildFromPRepr :: TyCon -> TyCon -> TyCon -> VM CoreExpr
 314 buildFromPRepr vect_tc repr_tc _
 315   = do
 316       arg_ty <- mkPReprType res_ty
 317       arg <- newLocalVar (fsLit "x") arg_ty
 318
 319       result <- from_sum (unwrapFamInstScrut repr_tc ty_args (Var arg))
 320                          (tyConDataCons vect_tc)
 321       return $ Lam arg result
 322   where
 323     ty_args = mkTyVarTys (tyConTyVars vect_tc)
 324     res_ty  = mkTyConApp vect_tc ty_args
 325
 326     from_sum _    []    = pprPanic "buildFromPRepr" (ppr vect_tc)
 327     from_sum expr [con] = from_prod expr con
 328     from_sum expr cons
 329       = do
 330           (sum_tc, _, _, sum_ty_args) <- reprSumTyCons vect_tc
 331           let sum_cons = tyConDataCons sum_tc
 332           vars <- newLocalVars (fsLit "x") sum_ty_args
 333           prods <- zipWithM from_prod (map Var vars) cons
 334           return . mkWildCase expr (exprType expr) res_ty
 335                  $ zipWith3 mk_alt sum_cons vars prods
 336       where
 337         mk_alt con var expr = (DataAlt con, [var], expr)
 338
 339     from_prod expr con
 340       = case dataConRepArgTys con of
 341           []   -> return $ apply_con []
 342           [_]  -> return $ apply_con [expr]
 343           tys  -> do
 344                     prod_con <- builtin (prodDataCon (length tys))
 345                     vars <- newLocalVars (fsLit "y") tys
 346                     return $ mkWildCase expr (exprType expr) res_ty
 347                              [(DataAlt prod_con, vars, apply_con (map Var vars))]
 348       where
 349         apply_con exprs = mkConApp con (map Type ty_args) `mkApps` exprs
 350
 351 buildToArrPRepr :: TyCon -> TyCon -> TyCon -> VM CoreExpr
 352 buildToArrPRepr vect_tc prepr_tc pdata_tc
 353   = do
 354       arg_ty <- mkPDataType el_ty
 355       res_ty <- mkPDataType =<< mkPReprType el_ty
 356       arg    <- newLocalVar (fsLit "xs") arg_ty
 357
 358       pdata_co <- mkBuiltinCo pdataTyCon
 359       let Just repr_co = tyConFamilyCoercion_maybe prepr_tc
 360           co           = mkAppCoercion pdata_co
 361                        . mkSymCoercion
 362                        $ mkTyConApp repr_co ty_args
 363
 364           scrut   = unwrapFamInstScrut pdata_tc ty_args (Var arg)
 365
 366       (vars, result) <- to_sum (tyConDataCons vect_tc)
 367
 368       return . Lam arg
 369              $ mkWildCase scrut (mkTyConApp pdata_tc ty_args) res_ty
 370                [(DataAlt pdata_dc, vars, mkCoerce co result)]
 371   where
 372     ty_args = mkTyVarTys $ tyConTyVars vect_tc
 373     el_ty   = mkTyConApp vect_tc ty_args
 374
 375     [pdata_dc] = tyConDataCons pdata_tc
 376
 377     to_sum []    = do
 378                      pvoid <- builtin pvoidVar
 379                      return ([], Var pvoid)
 380     to_sum [con] = to_prod con
 381     to_sum cons  = do
 382                      (vars, exprs) <- mapAndUnzipM to_prod cons
 383                      (_, pdata_tc, sel_ty, arg_tys) <- reprSumTyCons vect_tc
 384                      sel <- newLocalVar (fsLit "sel") sel_ty
 385                      let [pdata_con] = tyConDataCons pdata_tc
 386                          result = wrapFamInstBody pdata_tc arg_tys
 387                                 . mkConApp pdata_con
 388                                 $ map Type arg_tys ++ (Var sel : exprs)
 389                      return (sel : concat vars, result)
 390
 391     to_prod con
 392       | [] <- tys = do
 393                       pvoid <- builtin pvoidVar
 394                       return ([], Var pvoid)
 395       | [ty] <- tys = do
 396                         var <- newLocalVar (fsLit "x") ty
 397                         return ([var], Var var)
 398       | otherwise
 399         = do
 400             vars <- newLocalVars (fsLit "x") tys
 401             prod_tc <- builtin (prodTyCon (length tys))
 402             (pdata_prod_tc, _) <- pdataReprTyCon (mkTyConApp prod_tc tys)
 403             let [pdata_prod_con] = tyConDataCons pdata_prod_tc
 404                 result = wrapFamInstBody pdata_prod_tc tys
 405                        . mkConApp pdata_prod_con
 406                        $ map Type tys ++ map Var vars
 407             return (vars, result)
 408       where
 409         tys = dataConRepArgTys con
 410
 411 buildFromArrPRepr :: TyCon -> TyCon -> TyCon -> VM CoreExpr
 412 buildFromArrPRepr vect_tc prepr_tc pdata_tc
 413   = do
 414       arg_ty <- mkPDataType =<< mkPReprType el_ty
 415       res_ty <- mkPDataType el_ty
 416       arg    <- newLocalVar (fsLit "xs") arg_ty
 417
 418       pdata_co <- mkBuiltinCo pdataTyCon
 419       let Just repr_co = tyConFamilyCoercion_maybe prepr_tc
 420           co           = mkAppCoercion pdata_co
 421                        $ mkTyConApp repr_co var_tys
 422
 423           scrut  = mkCoerce co (Var arg)
 424
 425       (args, mk) <- from_sum res_ty scrut (tyConDataCons vect_tc)
 426
 427       let result = wrapFamInstBody pdata_tc var_tys
 428                  . mkConApp pdata_dc
 429                  $ map Type var_tys ++ args
 430
 431       return $ Lam arg (mk result)
 432   where
 433     var_tys = mkTyVarTys $ tyConTyVars vect_tc
 434     el_ty   = mkTyConApp vect_tc var_tys
 435
 436     [pdata_dc] = tyConDataCons pdata_tc
 437
 438     from_sum res_ty expr [] = return ([], mk)
 439       where
 440         mk body = mkWildCase expr (exprType expr) res_ty [(DEFAULT, [], body)]
 441     from_sum res_ty expr [con] = from_prod res_ty expr con
 442     from_sum res_ty expr cons
 443       = do
 444           (_, pdata_tc, sel_ty, arg_tys) <- reprSumTyCons vect_tc
 445           sel  <- newLocalVar (fsLit "sel") sel_ty
 446           vars <- newLocalVars (fsLit "xs") arg_tys
 447           rs   <- zipWithM (from_prod res_ty) (map Var vars) cons
 448           let (prods, mks) = unzip rs
 449               [pdata_con]  = tyConDataCons pdata_tc
 450               scrut        = unwrapFamInstScrut pdata_tc arg_tys expr
 451
 452               mk body = mkWildCase scrut (exprType scrut) res_ty
 453                         [(DataAlt pdata_con, sel : vars, foldr ($) body mks)]
 454           return (Var sel : concat prods, mk)
 455
 456
 457     from_prod res_ty expr con
 458       | []  <- tys = return ([], id)
 459       | [_] <- tys = return ([expr], id)
 460       | otherwise
 461         = do
 462             prod_tc <- builtin (prodTyCon (length tys))
 463             (pdata_tc, _) <- pdataReprTyCon (mkTyConApp prod_tc tys)
 464             pdata_tys <- mapM mkPDataType tys
 465             vars <- newLocalVars (fsLit "ys") pdata_tys
 466             let [pdata_con] = tyConDataCons pdata_tc
 467                 scrut       = unwrapFamInstScrut pdata_tc tys expr
 468
 469                 mk body = mkWildCase scrut (exprType scrut) res_ty
 470                           [(DataAlt pdata_con, vars, body)]
 471
 472             return (map Var vars, mk)
 473       where
 474         tys = dataConRepArgTys con
 475
 476 buildPRDict :: TyCon -> TyCon -> TyCon -> VM CoreExpr
 477 buildPRDict vect_tc prepr_tc _
 478   = do
 479       dict <- sum_dict (tyConDataCons vect_tc)
 480       pr_co <- mkBuiltinCo prTyCon
 481       let co = mkAppCoercion pr_co
 482              . mkSymCoercion
 483              $ mkTyConApp arg_co ty_args
 484       return (mkCoerce co dict)
 485   where
 486     ty_args = mkTyVarTys (tyConTyVars vect_tc)
 487     Just arg_co = tyConFamilyCoercion_maybe prepr_tc
 488
 489     sum_dict []    = prDFunOfTyCon =<< builtin voidTyCon
 490     sum_dict [con] = prod_dict con
 491     sum_dict cons  = do
 492                        dicts <- mapM prod_dict cons
 493                        (sum_tc, _, _, sum_ty_args) <- reprSumTyCons vect_tc
 494                        dfun <- prDFunOfTyCon sum_tc
 495                        return $ dfun `mkTyApps` sum_ty_args `mkApps` dicts
 496
 497     prod_dict con
 498       | []   <- tys = prDFunOfTyCon =<< builtin voidTyCon
 499       | [ty] <- tys = mkPR ty
 500       | otherwise   = do
 501                         dicts <- mapM mkPR tys
 502                         prod_tc <- builtin (prodTyCon (length tys))
 503                         dfun <- prDFunOfTyCon prod_tc
 504                         return $ dfun `mkTyApps` tys `mkApps` dicts
 505       where
 506         tys = dataConRepArgTys con
 507
 508 buildPDataTyCon :: TyCon -> TyCon -> VM TyCon
 509 buildPDataTyCon orig_tc vect_tc = fixV $ \repr_tc ->
 510   do
 511     name' <- cloneName mkPDataTyConOcc orig_name
 512     rhs   <- buildPDataTyConRhs orig_name vect_tc repr_tc
 513     pdata <- builtin pdataTyCon
 514
 515     liftDs $ buildAlgTyCon name'
 516                            tyvars
 517                            []          -- no stupid theta
 518                            rhs
 519                            rec_flag    -- FIXME: is this ok?
 520                            False       -- FIXME: no generics
 521                            False       -- not GADT syntax
 522                            (Just $ mk_fam_inst pdata vect_tc)
 523   where
 524     orig_name = tyConName orig_tc
 525     tyvars = tyConTyVars vect_tc
 526     rec_flag = boolToRecFlag (isRecursiveTyCon vect_tc)
 527
 528
 529 buildPDataTyConRhs :: Name -> TyCon -> TyCon -> VM AlgTyConRhs
 530 buildPDataTyConRhs orig_name vect_tc repr_tc
 531   = do
 532       data_con <- buildPDataDataCon orig_name vect_tc repr_tc
 533       return $ DataTyCon { data_cons = [data_con], is_enum = False }
 534
 535 buildPDataDataCon :: Name -> TyCon -> TyCon -> VM DataCon
 536 buildPDataDataCon orig_name vect_tc repr_tc
 537   = do
 538       dc_name  <- cloneName mkPDataDataConOcc orig_name
 539       comp_tys <- components
 540
 541       liftDs $ buildDataCon dc_name
 542                             False                  -- not infix
 543                             (map (const NotMarkedStrict) comp_tys)
 544                             []                     -- no field labels
 545                             tvs
 546                             []                     -- no existentials
 547                             []                     -- no eq spec
 548                             []                     -- no context
 549                             comp_tys
 550                             (mkFamilyTyConApp repr_tc (mkTyVarTys tvs))
 551                             repr_tc
 552   where
 553     tvs   = tyConTyVars vect_tc
 554     cons  = tyConDataCons vect_tc
 555     arity = length cons
 556
 557     components
 558       | arity > 1 = liftM2 (:) (builtin (selTy arity)) data_components
 559       | otherwise = data_components
 560
 561     data_components = mapM mkPDataType
 562                     . concat
 563                     $ map dataConRepArgTys cons
 564
 565 mkPADFun :: TyCon -> VM Var
 566 mkPADFun vect_tc
 567   = newExportedVar (mkPADFunOcc $ getOccName vect_tc) =<< paDFunType vect_tc
 568
 569 buildTyConBindings :: TyCon -> TyCon -> TyCon -> TyCon -> Var
 570                    -> VM [(Var, CoreExpr)]
 571 buildTyConBindings orig_tc vect_tc prepr_tc pdata_tc dfun
 572   = do
 573       vectDataConWorkers orig_tc vect_tc pdata_tc
 574       dict <- buildPADict vect_tc prepr_tc pdata_tc dfun
 575       binds <- takeHoisted
 576       return $ (dfun, dict) : binds
 577
 578 vectDataConWorkers :: TyCon -> TyCon -> TyCon -> VM ()
 579 vectDataConWorkers orig_tc vect_tc arr_tc
 580   = do
 581       bs <- sequence
 582           . zipWith3 def_worker  (tyConDataCons orig_tc) rep_tys
 583           $ zipWith4 mk_data_con (tyConDataCons vect_tc)
 584                                  rep_tys
 585                                  (inits rep_tys)
 586                                  (tail $ tails rep_tys)
 587       mapM_ (uncurry hoistBinding) bs
 588   where
 589     tyvars   = tyConTyVars vect_tc
 590     var_tys  = mkTyVarTys tyvars
 591     ty_args  = map Type var_tys
 592     res_ty   = mkTyConApp vect_tc var_tys
 593
 594     cons     = tyConDataCons vect_tc
 595     arity    = length cons
 596     [arr_dc] = tyConDataCons arr_tc
 597
 598     rep_tys  = map dataConRepArgTys $ tyConDataCons vect_tc
 599
 600
 601     mk_data_con con tys pre post
 602       = liftM2 (,) (vect_data_con con)
 603                    (lift_data_con tys pre post (mkDataConTag con))
 604
 605     sel_replicate len tag
 606       | arity > 1 = do
 607                       rep <- builtin (selReplicate arity)
 608                       return [rep `mkApps` [len, tag]]
 609
 610       | otherwise = return []
 611
 612     vect_data_con con = return $ mkConApp con ty_args
 613     lift_data_con tys pre_tys post_tys tag
 614       = do
 615           len  <- builtin liftingContext
 616           args <- mapM (newLocalVar (fsLit "xs"))
 617                   =<< mapM mkPDataType tys
 618
 619           sel  <- sel_replicate (Var len) tag
 620
 621           pre   <- mapM emptyPD (concat pre_tys)
 622           post  <- mapM emptyPD (concat post_tys)
 623
 624           return . mkLams (len : args)
 625                  . wrapFamInstBody arr_tc var_tys
 626                  . mkConApp arr_dc
 627                  $ ty_args ++ sel ++ pre ++ map Var args ++ post
 628
 629     def_worker data_con arg_tys mk_body
 630       = do
 631           body <- closedV
 632                 . inBind orig_worker
 633                 . polyAbstract tyvars $ \abstract ->
 634                   liftM (abstract . vectorised)
 635                 $ buildClosures tyvars [] arg_tys res_ty mk_body
 636
 637           vect_worker <- cloneId mkVectOcc orig_worker (exprType body)
 638           defGlobalVar orig_worker vect_worker
 639           return (vect_worker, body)
 640       where
 641         orig_worker = dataConWorkId data_con
 642
 643 buildPADict :: TyCon -> TyCon -> TyCon -> Var -> VM CoreExpr
 644 buildPADict vect_tc prepr_tc arr_tc _
 645   = polyAbstract tvs $ \abstract ->
 646     do
 647       meth_binds <- mapM mk_method paMethods
 648       let meth_exprs = map (Var . fst) meth_binds
 649
 650       pa_dc <- builtin paDataCon
 651       let dict = mkConApp pa_dc (Type (mkTyConApp vect_tc arg_tys) : meth_exprs)
 652           body = Let (Rec meth_binds) dict
 653       return . mkInlineMe $ abstract body
 654   where
 655     tvs = tyConTyVars arr_tc
 656     arg_tys = mkTyVarTys tvs
 657
 658     mk_method (name, build)
 659       = localV
 660       $ do
 661           body <- build vect_tc prepr_tc arr_tc
 662           var  <- newLocalVar name (exprType body)
 663           return (var, mkInlineMe body)
 664
 665 paMethods :: [(FastString, TyCon -> TyCon -> TyCon -> VM CoreExpr)]
 666 paMethods = [(fsLit "toPRepr",      buildToPRepr),
 667              (fsLit "fromPRepr",    buildFromPRepr),
 668              (fsLit "toArrPRepr",   buildToArrPRepr),
 669              (fsLit "fromArrPRepr", buildFromArrPRepr),
 670              (fsLit "dictPRepr",    buildPRDict)]
 671
 672 -- | Split the given tycons into two sets depending on whether they have to be
 673 -- converted (first list) or not (second list). The first argument contains
 674 -- information about the conversion status of external tycons:
 675 --
 676 --   * tycons which have converted versions are mapped to True
 677 --   * tycons which are not changed by vectorisation are mapped to False
 678 --   * tycons which can't be converted are not elements of the map
 679 --
 680 classifyTyCons :: UniqFM Bool -> [TyConGroup] -> ([TyCon], [TyCon])
 681 classifyTyCons = classify [] []
 682   where
 683     classify conv keep _  [] = (conv, keep)
 684     classify conv keep cs ((tcs, ds) : rs)
 685       | can_convert && must_convert
 686         = classify (tcs ++ conv) keep (cs `addListToUFM` [(tc,True) | tc <- tcs]) rs
 687       | can_convert
 688         = classify conv (tcs ++ keep) (cs `addListToUFM` [(tc,False) | tc <- tcs]) rs
 689       | otherwise
 690         = classify conv keep cs rs
 691       where
 692         refs = ds `delListFromUniqSet` tcs
 693
 694         can_convert  = isNullUFM (refs `minusUFM` cs) && all convertable tcs
 695         must_convert = foldUFM (||) False (intersectUFM_C const cs refs)
 696
 697         convertable tc = isDataTyCon tc && all isVanillaDataCon (tyConDataCons tc)
 698
 699 -- | Compute mutually recursive groups of tycons in topological order
 700 --
 701 tyConGroups :: [TyCon] -> [TyConGroup]
 702 tyConGroups tcs = map mk_grp (stronglyConnCompFromEdgedVertices edges)
 703   where
 704     edges = [((tc, ds), tc, uniqSetToList ds) | tc <- tcs
 705                                 , let ds = tyConsOfTyCon tc]
 706
 707     mk_grp (AcyclicSCC (tc, ds)) = ([tc], ds)
 708     mk_grp (CyclicSCC els)       = (tcs, unionManyUniqSets dss)
 709       where
 710         (tcs, dss) = unzip els
 711
 712 tyConsOfTyCon :: TyCon -> UniqSet TyCon
 713 tyConsOfTyCon
 714   = tyConsOfTypes . concatMap dataConRepArgTys . tyConDataCons
 715
 716 tyConsOfType :: Type -> UniqSet TyCon
 717 tyConsOfType ty
 718   | Just ty' <- coreView ty    = tyConsOfType ty'
 719 tyConsOfType (TyVarTy _)       = emptyUniqSet
 720 tyConsOfType (TyConApp tc tys) = extend (tyConsOfTypes tys)
 721   where
 722     extend | isUnLiftedTyCon tc
 723            || isTupleTyCon   tc = id
 724
 725            | otherwise          = (`addOneToUniqSet` tc)
 726
 727 tyConsOfType (AppTy a b)       = tyConsOfType a `unionUniqSets` tyConsOfType b
 728 tyConsOfType (FunTy a b)       = (tyConsOfType a `unionUniqSets` tyConsOfType b)
 729                                  `addOneToUniqSet` funTyCon
 730 tyConsOfType (ForAllTy _ ty)   = tyConsOfType ty
 731 tyConsOfType other             = pprPanic "ClosureConv.tyConsOfType" $ ppr other
 732
 733 tyConsOfTypes :: [Type] -> UniqSet TyCon
 734 tyConsOfTypes = unionManyUniqSets . map tyConsOfType
 735
 736
 737 -- ----------------------------------------------------------------------------
 738 -- Conversions
 739
 740 fromVect :: Type -> CoreExpr -> VM CoreExpr
 741 fromVect ty expr | Just ty' <- coreView ty = fromVect ty' expr
 742 fromVect (FunTy arg_ty res_ty) expr
 743   = do
 744       arg     <- newLocalVar (fsLit "x") arg_ty
 745       varg    <- toVect arg_ty (Var arg)
 746       varg_ty <- vectType arg_ty
 747       vres_ty <- vectType res_ty
 748       apply   <- builtin applyVar
 749       body    <- fromVect res_ty
 750                $ Var apply `mkTyApps` [varg_ty, vres_ty] `mkApps` [expr, varg]
 751       return $ Lam arg body
 752 fromVect ty expr
 753   = identityConv ty >> return expr
 754
 755 toVect :: Type -> CoreExpr -> VM CoreExpr
 756 toVect ty expr = identityConv ty >> return expr
 757
 758 identityConv :: Type -> VM ()
 759 identityConv ty | Just ty' <- coreView ty = identityConv ty'
 760 identityConv (TyConApp tycon tys)
 761   = do
 762       mapM_ identityConv tys
 763       identityConvTyCon tycon
 764 identityConv _ = noV
 765
 766 identityConvTyCon :: TyCon -> VM ()
 767 identityConvTyCon tc
 768   | isBoxedTupleTyCon tc = return ()
 769   | isUnLiftedTyCon tc   = return ()
 770   | otherwise            = do
 771                              tc' <- maybeV (lookupTyCon tc)
 772                              if tc == tc' then return () else noV
 773