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 CoreUnfold
  15 import MkCore            ( mkWildCase )
  16 import BuildTyCl
  17 import DataCon
  18 import TyCon
  19 import Type
  20 import TypeRep
  21 import Coercion
  22 import FamInstEnv        ( FamInst, mkLocalFamInst )
  23 import OccName
  24 import Id
  25 import MkId
  26 import BasicTypes        ( HsBang(..), boolToRecFlag,
  27                            alwaysInlinePragma, dfunInlinePragma )
  28 import Var               ( Var, TyVar, varType )
  29 import Name              ( Name, getOccName )
  30 import NameEnv
  31
  32 import Unique
  33 import UniqFM
  34 import UniqSet
  35 import Util
  36 import Digraph           ( SCC(..), stronglyConnCompFromEdgedVertices )
  37
  38 import Outputable
  39 import FastString
  40
  41 import MonadUtils     ( zipWith3M, foldrM, concatMapM )
  42 import Control.Monad  ( liftM, liftM2, zipWithM, zipWithM_, mapAndUnzipM )
  43 import Data.List      ( inits, tails, zipWith4, zipWith5 )
  44
  45 -- ----------------------------------------------------------------------------
  46 -- Types
  47
  48 vectTyCon :: TyCon -> VM TyCon
  49 vectTyCon tc
  50   | isFunTyCon tc        = builtin closureTyCon
  51   | isBoxedTupleTyCon tc = return tc
  52   | isUnLiftedTyCon tc   = return tc
  53   | otherwise            = maybeCantVectoriseM "Tycon not vectorised:" (ppr tc)
  54                          $ lookupTyCon tc
  55
  56 vectAndLiftType :: Type -> VM (Type, Type)
  57 vectAndLiftType ty | Just ty' <- coreView ty = vectAndLiftType ty'
  58 vectAndLiftType ty
  59   = do
  60       mdicts   <- mapM paDictArgType tyvars
  61       let dicts = [dict | Just dict <- mdicts]
  62       vmono_ty <- vectType mono_ty
  63       lmono_ty <- mkPDataType vmono_ty
  64       return (abstractType tyvars dicts vmono_ty,
  65               abstractType tyvars dicts lmono_ty)
  66   where
  67     (tyvars, mono_ty) = splitForAllTys ty
  68
  69
  70 vectType :: Type -> VM Type
  71 vectType ty | Just ty' <- coreView ty = vectType ty'
  72 vectType (TyVarTy tv) = return $ TyVarTy tv
  73 vectType (AppTy ty1 ty2) = liftM2 AppTy (vectType ty1) (vectType ty2)
  74 vectType (TyConApp tc tys) = liftM2 TyConApp (vectTyCon tc) (mapM vectType tys)
  75 vectType (FunTy ty1 ty2)   = liftM2 TyConApp (builtin closureTyCon)
  76                                              (mapM vectAndBoxType [ty1,ty2])
  77 vectType ty@(ForAllTy _ _)
  78   = do
  79       mdicts   <- mapM paDictArgType tyvars
  80       mono_ty' <- vectType mono_ty
  81       return $ abstractType tyvars [dict | Just dict <- mdicts] mono_ty'
  82   where
  83     (tyvars, mono_ty) = splitForAllTys ty
  84
  85 vectType ty = cantVectorise "Can't vectorise type" (ppr ty)
  86
  87 vectAndBoxType :: Type -> VM Type
  88 vectAndBoxType ty = vectType ty >>= boxType
  89
  90 abstractType :: [TyVar] -> [Type] -> Type -> Type
  91 abstractType tyvars dicts = mkForAllTys tyvars . mkFunTys dicts
  92
  93 -- ----------------------------------------------------------------------------
  94 -- Boxing
  95
  96 boxType :: Type -> VM Type
  97 boxType ty
  98   | Just (tycon, []) <- splitTyConApp_maybe ty
  99   , isUnLiftedTyCon tycon
 100   = do
 101       r <- lookupBoxedTyCon tycon
 102       case r of
 103         Just tycon' -> return $ mkTyConApp tycon' []
 104         Nothing     -> return ty
 105 boxType ty = return ty
 106
 107 -- ----------------------------------------------------------------------------
 108 -- Type definitions
 109
 110 type TyConGroup = ([TyCon], UniqSet TyCon)
 111
 112 vectTypeEnv :: TypeEnv -> VM (TypeEnv, [FamInst], [(Var, CoreExpr)])
 113 vectTypeEnv env
 114   = do
 115       cs <- readGEnv $ mk_map . global_tycons
 116       let (conv_tcs, keep_tcs) = classifyTyCons cs groups
 117           keep_dcs             = concatMap tyConDataCons keep_tcs
 118       zipWithM_ defTyCon   keep_tcs keep_tcs
 119       zipWithM_ defDataCon keep_dcs keep_dcs
 120       new_tcs <- vectTyConDecls conv_tcs
 121
 122       let orig_tcs = keep_tcs ++ conv_tcs
 123           vect_tcs = keep_tcs ++ new_tcs
 124
 125       (_, binds, inst_tcs) <- fixV $ \ ~(dfuns', _, _) ->
 126         do
 127           defTyConPAs (zipLazy vect_tcs dfuns')
 128           reprs <- mapM tyConRepr vect_tcs
 129           repr_tcs  <- zipWith3M buildPReprTyCon orig_tcs vect_tcs reprs
 130           pdata_tcs <- zipWith3M buildPDataTyCon orig_tcs vect_tcs reprs
 131           dfuns <- sequence $ zipWith5 buildTyConBindings orig_tcs
 132                                                           vect_tcs
 133                                                           repr_tcs
 134                                                           pdata_tcs
 135                                                           reprs
 136           binds <- takeHoisted
 137           return (dfuns, binds, repr_tcs ++ pdata_tcs)
 138
 139       let all_new_tcs = new_tcs ++ inst_tcs
 140
 141       let new_env = extendTypeEnvList env
 142                        (map ATyCon all_new_tcs
 143                         ++ [ADataCon dc | tc <- all_new_tcs
 144                                         , dc <- tyConDataCons tc])
 145
 146       return (new_env, map mkLocalFamInst inst_tcs, binds)
 147   where
 148     tycons = typeEnvTyCons env
 149     groups = tyConGroups tycons
 150
 151     mk_map env = listToUFM_Directly [(u, getUnique n /= u) | (u,n) <- nameEnvUniqueElts env]
 152
 153
 154 vectTyConDecls :: [TyCon] -> VM [TyCon]
 155 vectTyConDecls tcs = fixV $ \tcs' ->
 156   do
 157     mapM_ (uncurry defTyCon) (zipLazy tcs tcs')
 158     mapM vectTyConDecl tcs
 159
 160 vectTyConDecl :: TyCon -> VM TyCon
 161 vectTyConDecl tc
 162   = do
 163       name' <- cloneName mkVectTyConOcc name
 164       rhs'  <- vectAlgTyConRhs tc (algTyConRhs tc)
 165
 166       liftDs $ buildAlgTyCon name'
 167                              tyvars
 168                              []           -- no stupid theta
 169                              rhs'
 170                              rec_flag     -- FIXME: is this ok?
 171                              False        -- FIXME: no generics
 172                              False        -- not GADT syntax
 173                              Nothing      -- not a family instance
 174   where
 175     name   = tyConName tc
 176     tyvars = tyConTyVars tc
 177     rec_flag = boolToRecFlag (isRecursiveTyCon tc)
 178
 179 vectAlgTyConRhs :: TyCon -> AlgTyConRhs -> VM AlgTyConRhs
 180 vectAlgTyConRhs _ (DataTyCon { data_cons = data_cons
 181                              , is_enum   = is_enum
 182                              })
 183   = do
 184       data_cons' <- mapM vectDataCon data_cons
 185       zipWithM_ defDataCon data_cons data_cons'
 186       return $ DataTyCon { data_cons = data_cons'
 187                          , is_enum   = is_enum
 188                          }
 189 vectAlgTyConRhs tc _ = cantVectorise "Can't vectorise type definition:" (ppr tc)
 190
 191 vectDataCon :: DataCon -> VM DataCon
 192 vectDataCon dc
 193   | not . null $ dataConExTyVars dc
 194         = cantVectorise "Can't vectorise constructor (existentials):" (ppr dc)
 195   | not . null $ dataConEqSpec   dc
 196         = cantVectorise "Can't vectorise constructor (eq spec):" (ppr dc)
 197   | otherwise
 198   = do
 199       name'    <- cloneName mkVectDataConOcc name
 200       tycon'   <- vectTyCon tycon
 201       arg_tys  <- mapM vectType rep_arg_tys
 202
 203       liftDs $ buildDataCon name'
 204                             False           -- not infix
 205                             (map (const HsNoBang) arg_tys)
 206                             []              -- no labelled fields
 207                             univ_tvs
 208                             []              -- no existential tvs for now
 209                             []              -- no eq spec for now
 210                             []              -- no context
 211                             arg_tys
 212                             (mkFamilyTyConApp tycon' (mkTyVarTys univ_tvs))
 213                             tycon'
 214   where
 215     name        = dataConName dc
 216     univ_tvs    = dataConUnivTyVars dc
 217     rep_arg_tys = dataConRepArgTys dc
 218     tycon       = dataConTyCon dc
 219
 220 mk_fam_inst :: TyCon -> TyCon -> (TyCon, [Type])
 221 mk_fam_inst fam_tc arg_tc
 222   = (fam_tc, [mkTyConApp arg_tc . mkTyVarTys $ tyConTyVars arg_tc])
 223
 224
 225 buildPReprTyCon :: TyCon -> TyCon -> SumRepr -> VM TyCon
 226 buildPReprTyCon orig_tc vect_tc repr
 227   = do
 228       name     <- cloneName mkPReprTyConOcc (tyConName orig_tc)
 229       -- rhs_ty   <- buildPReprType vect_tc
 230       rhs_ty   <- sumReprType repr
 231       prepr_tc <- builtin preprTyCon
 232       liftDs $ buildSynTyCon name
 233                              tyvars
 234                              (SynonymTyCon rhs_ty)
 235                              (typeKind rhs_ty)
 236                              (Just $ mk_fam_inst prepr_tc vect_tc)
 237   where
 238     tyvars = tyConTyVars vect_tc
 239
 240 data CompRepr = Keep Type
 241                      CoreExpr     -- PR dictionary for the type
 242               | Wrap Type
 243
 244 data ProdRepr = EmptyProd
 245               | UnaryProd CompRepr
 246               | Prod { repr_tup_tc   :: TyCon  -- representation tuple tycon
 247                      , repr_ptup_tc  :: TyCon  -- PData representation tycon
 248                      , repr_comp_tys :: [Type] -- representation types of
 249                      , repr_comps    :: [CompRepr]          -- components
 250                      }
 251 data ConRepr  = ConRepr DataCon ProdRepr
 252
 253 data SumRepr  = EmptySum
 254               | UnarySum ConRepr
 255               | Sum  { repr_sum_tc   :: TyCon  -- representation sum tycon
 256                      , repr_psum_tc  :: TyCon  -- PData representation tycon
 257                      , repr_sel_ty   :: Type   -- type of selector
 258                      , repr_con_tys :: [Type]  -- representation types of
 259                      , repr_cons     :: [ConRepr]           -- components
 260                      }
 261
 262 tyConRepr :: TyCon -> VM SumRepr
 263 tyConRepr tc = sum_repr (tyConDataCons tc)
 264   where
 265     sum_repr []    = return EmptySum
 266     sum_repr [con] = liftM UnarySum (con_repr con)
 267     sum_repr cons  = do
 268                        rs     <- mapM con_repr cons
 269                        sum_tc <- builtin (sumTyCon arity)
 270                        tys    <- mapM conReprType rs
 271                        (psum_tc, _) <- pdataReprTyCon (mkTyConApp sum_tc tys)
 272                        sel_ty <- builtin (selTy arity)
 273                        return $ Sum { repr_sum_tc  = sum_tc
 274                                     , repr_psum_tc = psum_tc
 275                                     , repr_sel_ty  = sel_ty
 276                                     , repr_con_tys = tys
 277                                     , repr_cons    = rs
 278                                     }
 279       where
 280         arity = length cons
 281
 282     con_repr con = liftM (ConRepr con) (prod_repr (dataConRepArgTys con))
 283
 284     prod_repr []   = return EmptyProd
 285     prod_repr [ty] = liftM UnaryProd (comp_repr ty)
 286     prod_repr tys  = do
 287                        rs <- mapM comp_repr tys
 288                        tup_tc <- builtin (prodTyCon arity)
 289                        tys'    <- mapM compReprType rs
 290                        (ptup_tc, _) <- pdataReprTyCon (mkTyConApp tup_tc tys')
 291                        return $ Prod { repr_tup_tc   = tup_tc
 292                                      , repr_ptup_tc  = ptup_tc
 293                                      , repr_comp_tys = tys'
 294                                      , repr_comps    = rs
 295                                      }
 296       where
 297         arity = length tys
 298
 299     comp_repr ty = liftM (Keep ty) (prDictOfType ty)
 300                    `orElseV` return (Wrap ty)
 301
 302 sumReprType :: SumRepr -> VM Type
 303 sumReprType EmptySum = voidType
 304 sumReprType (UnarySum r) = conReprType r
 305 sumReprType (Sum { repr_sum_tc  = sum_tc, repr_con_tys = tys })
 306   = return $ mkTyConApp sum_tc tys
 307
 308 conReprType :: ConRepr -> VM Type
 309 conReprType (ConRepr _ r) = prodReprType r
 310
 311 prodReprType :: ProdRepr -> VM Type
 312 prodReprType EmptyProd = voidType
 313 prodReprType (UnaryProd r) = compReprType r
 314 prodReprType (Prod { repr_tup_tc = tup_tc, repr_comp_tys = tys })
 315   = return $ mkTyConApp tup_tc tys
 316
 317 compReprType :: CompRepr -> VM Type
 318 compReprType (Keep ty _) = return ty
 319 compReprType (Wrap ty) = do
 320                              wrap_tc <- builtin wrapTyCon
 321                              return $ mkTyConApp wrap_tc [ty]
 322
 323 compOrigType :: CompRepr -> Type
 324 compOrigType (Keep ty _) = ty
 325 compOrigType (Wrap ty) = ty
 326
 327 buildToPRepr :: TyCon -> TyCon -> TyCon -> SumRepr -> VM CoreExpr
 328 buildToPRepr vect_tc repr_tc _ repr
 329   = do
 330       let arg_ty = mkTyConApp vect_tc ty_args
 331       res_ty <- mkPReprType arg_ty
 332       arg    <- newLocalVar (fsLit "x") arg_ty
 333       result <- to_sum (Var arg) arg_ty res_ty repr
 334       return $ Lam arg result
 335   where
 336     ty_args = mkTyVarTys (tyConTyVars vect_tc)
 337
 338     wrap_repr_inst = wrapFamInstBody repr_tc ty_args
 339
 340     to_sum _ _ _ EmptySum
 341       = do
 342           void <- builtin voidVar
 343           return $ wrap_repr_inst $ Var void
 344
 345     to_sum arg arg_ty res_ty (UnarySum r)
 346       = do
 347           (pat, vars, body) <- con_alt r
 348           return $ mkWildCase arg arg_ty res_ty
 349                    [(pat, vars, wrap_repr_inst body)]
 350
 351     to_sum arg arg_ty res_ty (Sum { repr_sum_tc  = sum_tc
 352                                   , repr_con_tys = tys
 353                                   , repr_cons    =  cons })
 354       = do
 355           alts <- mapM con_alt cons
 356           let alts' = [(pat, vars, wrap_repr_inst
 357                                    $ mkConApp sum_con (map Type tys ++ [body]))
 358                         | ((pat, vars, body), sum_con)
 359                             <- zip alts (tyConDataCons sum_tc)]
 360           return $ mkWildCase arg arg_ty res_ty alts'
 361
 362     con_alt (ConRepr con r)
 363       = do
 364           (vars, body) <- to_prod r
 365           return (DataAlt con, vars, body)
 366
 367     to_prod EmptyProd
 368       = do
 369           void <- builtin voidVar
 370           return ([], Var void)
 371
 372     to_prod (UnaryProd comp)
 373       = do
 374           var  <- newLocalVar (fsLit "x") (compOrigType comp)
 375           body <- to_comp (Var var) comp
 376           return ([var], body)
 377
 378     to_prod(Prod { repr_tup_tc   = tup_tc
 379                  , repr_comp_tys = tys
 380                  , repr_comps    = comps })
 381       = do
 382           vars  <- newLocalVars (fsLit "x") (map compOrigType comps)
 383           exprs <- zipWithM to_comp (map Var vars) comps
 384           return (vars, mkConApp tup_con (map Type tys ++ exprs))
 385       where
 386         [tup_con] = tyConDataCons tup_tc
 387
 388     to_comp expr (Keep _ _) = return expr
 389     to_comp expr (Wrap ty)  = do
 390                                 wrap_tc <- builtin wrapTyCon
 391                                 return $ wrapNewTypeBody wrap_tc [ty] expr
 392
 393
 394 buildFromPRepr :: TyCon -> TyCon -> TyCon -> SumRepr -> VM CoreExpr
 395 buildFromPRepr vect_tc repr_tc _ repr
 396   = do
 397       arg_ty <- mkPReprType res_ty
 398       arg <- newLocalVar (fsLit "x") arg_ty
 399
 400       result <- from_sum (unwrapFamInstScrut repr_tc ty_args (Var arg))
 401                          repr
 402       return $ Lam arg result
 403   where
 404     ty_args = mkTyVarTys (tyConTyVars vect_tc)
 405     res_ty  = mkTyConApp vect_tc ty_args
 406
 407     from_sum _ EmptySum
 408       = do
 409           dummy <- builtin fromVoidVar
 410           return $ Var dummy `App` Type res_ty
 411
 412     from_sum expr (UnarySum r) = from_con expr r
 413     from_sum expr (Sum { repr_sum_tc  = sum_tc
 414                        , repr_con_tys = tys
 415                        , repr_cons    = cons })
 416       = do
 417           vars  <- newLocalVars (fsLit "x") tys
 418           es    <- zipWithM from_con (map Var vars) cons
 419           return $ mkWildCase expr (exprType expr) res_ty
 420                    [(DataAlt con, [var], e)
 421                       | (con, var, e) <- zip3 (tyConDataCons sum_tc) vars es]
 422
 423     from_con expr (ConRepr con r)
 424       = from_prod expr (mkConApp con $ map Type ty_args) r
 425
 426     from_prod _ con EmptyProd = return con
 427     from_prod expr con (UnaryProd r)
 428       = do
 429           e <- from_comp expr r
 430           return $ con `App` e
 431
 432     from_prod expr con (Prod { repr_tup_tc   = tup_tc
 433                              , repr_comp_tys = tys
 434                              , repr_comps    = comps
 435                              })
 436       = do
 437           vars <- newLocalVars (fsLit "y") tys
 438           es   <- zipWithM from_comp (map Var vars) comps
 439           return $ mkWildCase expr (exprType expr) res_ty
 440                    [(DataAlt tup_con, vars, con `mkApps` es)]
 441       where
 442         [tup_con] = tyConDataCons tup_tc
 443
 444     from_comp expr (Keep _ _) = return expr
 445     from_comp expr (Wrap ty)
 446       = do
 447           wrap <- builtin wrapTyCon
 448           return $ unwrapNewTypeBody wrap [ty] expr
 449
 450
 451 buildToArrPRepr :: TyCon -> TyCon -> TyCon -> SumRepr -> VM CoreExpr
 452 buildToArrPRepr vect_tc prepr_tc pdata_tc r
 453   = do
 454       arg_ty <- mkPDataType el_ty
 455       res_ty <- mkPDataType =<< mkPReprType el_ty
 456       arg    <- newLocalVar (fsLit "xs") arg_ty
 457
 458       pdata_co <- mkBuiltinCo pdataTyCon
 459       let Just repr_co = tyConFamilyCoercion_maybe prepr_tc
 460           co           = mkAppCoercion pdata_co
 461                        . mkSymCoercion
 462                        $ mkTyConApp repr_co ty_args
 463
 464           scrut   = unwrapFamInstScrut pdata_tc ty_args (Var arg)
 465
 466       (vars, result) <- to_sum r
 467
 468       return . Lam arg
 469              $ mkWildCase scrut (mkTyConApp pdata_tc ty_args) res_ty
 470                [(DataAlt pdata_dc, vars, mkCoerce co result)]
 471   where
 472     ty_args = mkTyVarTys $ tyConTyVars vect_tc
 473     el_ty   = mkTyConApp vect_tc ty_args
 474
 475     [pdata_dc] = tyConDataCons pdata_tc
 476
 477
 478     to_sum EmptySum = do
 479                         pvoid <- builtin pvoidVar
 480                         return ([], Var pvoid)
 481     to_sum (UnarySum r) = to_con r
 482     to_sum (Sum { repr_psum_tc = psum_tc
 483                 , repr_sel_ty  = sel_ty
 484                 , repr_con_tys = tys
 485                 , repr_cons    = cons
 486                 })
 487       = do
 488           (vars, exprs) <- mapAndUnzipM to_con cons
 489           sel <- newLocalVar (fsLit "sel") sel_ty
 490           return (sel : concat vars, mk_result (Var sel) exprs)
 491       where
 492         [psum_con] = tyConDataCons psum_tc
 493         mk_result sel exprs = wrapFamInstBody psum_tc tys
 494                             $ mkConApp psum_con
 495                             $ map Type tys ++ (sel : exprs)
 496
 497     to_con (ConRepr _ r) = to_prod r
 498
 499     to_prod EmptyProd = do
 500                           pvoid <- builtin pvoidVar
 501                           return ([], Var pvoid)
 502     to_prod (UnaryProd r)
 503       = do
 504           pty  <- mkPDataType (compOrigType r)
 505           var  <- newLocalVar (fsLit "x") pty
 506           expr <- to_comp (Var var) r
 507           return ([var], expr)
 508
 509     to_prod (Prod { repr_ptup_tc  = ptup_tc
 510                   , repr_comp_tys = tys
 511                   , repr_comps    = comps })
 512       = do
 513           ptys <- mapM (mkPDataType . compOrigType) comps
 514           vars <- newLocalVars (fsLit "x") ptys
 515           es   <- zipWithM to_comp (map Var vars) comps
 516           return (vars, mk_result es)
 517       where
 518         [ptup_con] = tyConDataCons ptup_tc
 519         mk_result exprs = wrapFamInstBody ptup_tc tys
 520                         $ mkConApp ptup_con
 521                         $ map Type tys ++ exprs
 522
 523     to_comp expr (Keep _ _) = return expr
 524
 525     -- FIXME: this is bound to be wrong!
 526     to_comp expr (Wrap ty)
 527       = do
 528           wrap_tc  <- builtin wrapTyCon
 529           (pwrap_tc, _) <- pdataReprTyCon (mkTyConApp wrap_tc [ty])
 530           return $ wrapNewTypeBody pwrap_tc [ty] expr
 531
 532
 533 buildFromArrPRepr :: TyCon -> TyCon -> TyCon -> SumRepr -> VM CoreExpr
 534 buildFromArrPRepr vect_tc prepr_tc pdata_tc r
 535   = do
 536       arg_ty <- mkPDataType =<< mkPReprType el_ty
 537       res_ty <- mkPDataType el_ty
 538       arg    <- newLocalVar (fsLit "xs") arg_ty
 539
 540       pdata_co <- mkBuiltinCo pdataTyCon
 541       let Just repr_co = tyConFamilyCoercion_maybe prepr_tc
 542           co           = mkAppCoercion pdata_co
 543                        $ mkTyConApp repr_co var_tys
 544
 545           scrut  = mkCoerce co (Var arg)
 546
 547           mk_result args = wrapFamInstBody pdata_tc var_tys
 548                          $ mkConApp pdata_con
 549                          $ map Type var_tys ++ args
 550
 551       (expr, _) <- fixV $ \ ~(_, args) ->
 552                      from_sum res_ty (mk_result args) scrut r
 553
 554       return $ Lam arg expr
 555
 556       -- (args, mk) <- from_sum res_ty scrut r
 557
 558       -- let result = wrapFamInstBody pdata_tc var_tys
 559       --           . mkConApp pdata_dc
 560       --           $ map Type var_tys ++ args
 561
 562       -- return $ Lam arg (mk result)
 563   where
 564     var_tys = mkTyVarTys $ tyConTyVars vect_tc
 565     el_ty   = mkTyConApp vect_tc var_tys
 566
 567     [pdata_con] = tyConDataCons pdata_tc
 568
 569     from_sum _ res _ EmptySum = return (res, [])
 570     from_sum res_ty res expr (UnarySum r) = from_con res_ty res expr r
 571     from_sum res_ty res expr (Sum { repr_psum_tc = psum_tc
 572                                   , repr_sel_ty  = sel_ty
 573                                   , repr_con_tys = tys
 574                                   , repr_cons    = cons })
 575       = do
 576           sel  <- newLocalVar (fsLit "sel") sel_ty
 577           ptys <- mapM mkPDataType tys
 578           vars <- newLocalVars (fsLit "xs") ptys
 579           (res', args) <- fold from_con res_ty res (map Var vars) cons
 580           let scrut = unwrapFamInstScrut psum_tc tys expr
 581               body  = mkWildCase scrut (exprType scrut) res_ty
 582                       [(DataAlt psum_con, sel : vars, res')]
 583           return (body, Var sel : args)
 584       where
 585         [psum_con] = tyConDataCons psum_tc
 586
 587
 588     from_con res_ty res expr (ConRepr _ r) = from_prod res_ty res expr r
 589
 590     from_prod _ res _ EmptyProd = return (res, [])
 591     from_prod res_ty res expr (UnaryProd r)
 592       = from_comp res_ty res expr r
 593     from_prod res_ty res expr (Prod { repr_ptup_tc  = ptup_tc
 594                                     , repr_comp_tys = tys
 595                                     , repr_comps    = comps })
 596       = do
 597           ptys <- mapM mkPDataType tys
 598           vars <- newLocalVars (fsLit "ys") ptys
 599           (res', args) <- fold from_comp res_ty res (map Var vars) comps
 600           let scrut = unwrapFamInstScrut ptup_tc tys expr
 601               body  = mkWildCase scrut (exprType scrut) res_ty
 602                       [(DataAlt ptup_con, vars, res')]
 603           return (body, args)
 604       where
 605         [ptup_con] = tyConDataCons ptup_tc
 606
 607     from_comp _ res expr (Keep _ _) = return (res, [expr])
 608     from_comp _ res expr (Wrap ty)
 609       = do
 610           wrap_tc  <- builtin wrapTyCon
 611           (pwrap_tc, _) <- pdataReprTyCon (mkTyConApp wrap_tc [ty])
 612           return (res, [unwrapNewTypeBody pwrap_tc [ty]
 613                         $ unwrapFamInstScrut pwrap_tc [ty] expr])
 614
 615     fold f res_ty res exprs rs = foldrM f' (res, []) (zip exprs rs)
 616       where
 617         f' (expr, r) (res, args) = do
 618                                      (res', args') <- f res_ty res expr r
 619                                      return (res', args' ++ args)
 620
 621 buildPRDict :: TyCon -> TyCon -> TyCon -> SumRepr -> VM CoreExpr
 622 buildPRDict vect_tc prepr_tc _ r
 623   = do
 624       dict <- sum_dict r
 625       pr_co <- mkBuiltinCo prTyCon
 626       let co = mkAppCoercion pr_co
 627              . mkSymCoercion
 628              $ mkTyConApp arg_co ty_args
 629       return (mkCoerce co dict)
 630   where
 631     ty_args = mkTyVarTys (tyConTyVars vect_tc)
 632     Just arg_co = tyConFamilyCoercion_maybe prepr_tc
 633
 634     sum_dict EmptySum = prDFunOfTyCon =<< builtin voidTyCon
 635     sum_dict (UnarySum r) = con_dict r
 636     sum_dict (Sum { repr_sum_tc  = sum_tc
 637                   , repr_con_tys = tys
 638                   , repr_cons    = cons
 639                   })
 640       = do
 641           dicts <- mapM con_dict cons
 642           dfun  <- prDFunOfTyCon sum_tc
 643           return $ dfun `mkTyApps` tys `mkApps` dicts
 644
 645     con_dict (ConRepr _ r) = prod_dict r
 646
 647     prod_dict EmptyProd = prDFunOfTyCon =<< builtin voidTyCon
 648     prod_dict (UnaryProd r) = comp_dict r
 649     prod_dict (Prod { repr_tup_tc   = tup_tc
 650                     , repr_comp_tys = tys
 651                     , repr_comps    = comps })
 652       = do
 653           dicts <- mapM comp_dict comps
 654           dfun <- prDFunOfTyCon tup_tc
 655           return $ dfun `mkTyApps` tys `mkApps` dicts
 656
 657     comp_dict (Keep _ pr) = return pr
 658     comp_dict (Wrap ty)   = wrapPR ty
 659
 660
 661 buildPDataTyCon :: TyCon -> TyCon -> SumRepr -> VM TyCon
 662 buildPDataTyCon orig_tc vect_tc repr = fixV $ \repr_tc ->
 663   do
 664     name' <- cloneName mkPDataTyConOcc orig_name
 665     rhs   <- buildPDataTyConRhs orig_name vect_tc repr_tc repr
 666     pdata <- builtin pdataTyCon
 667
 668     liftDs $ buildAlgTyCon name'
 669                            tyvars
 670                            []          -- no stupid theta
 671                            rhs
 672                            rec_flag    -- FIXME: is this ok?
 673                            False       -- FIXME: no generics
 674                            False       -- not GADT syntax
 675                            (Just $ mk_fam_inst pdata vect_tc)
 676   where
 677     orig_name = tyConName orig_tc
 678     tyvars = tyConTyVars vect_tc
 679     rec_flag = boolToRecFlag (isRecursiveTyCon vect_tc)
 680
 681
 682 buildPDataTyConRhs :: Name -> TyCon -> TyCon -> SumRepr -> VM AlgTyConRhs
 683 buildPDataTyConRhs orig_name vect_tc repr_tc repr
 684   = do
 685       data_con <- buildPDataDataCon orig_name vect_tc repr_tc repr
 686       return $ DataTyCon { data_cons = [data_con], is_enum = False }
 687
 688 buildPDataDataCon :: Name -> TyCon -> TyCon -> SumRepr -> VM DataCon
 689 buildPDataDataCon orig_name vect_tc repr_tc repr
 690   = do
 691       dc_name  <- cloneName mkPDataDataConOcc orig_name
 692       comp_tys <- sum_tys repr
 693
 694       liftDs $ buildDataCon dc_name
 695                             False                  -- not infix
 696                             (map (const HsNoBang) comp_tys)
 697                             []                     -- no field labels
 698                             tvs
 699                             []                     -- no existentials
 700                             []                     -- no eq spec
 701                             []                     -- no context
 702                             comp_tys
 703                             (mkFamilyTyConApp repr_tc (mkTyVarTys tvs))
 704                             repr_tc
 705   where
 706     tvs   = tyConTyVars vect_tc
 707
 708     sum_tys EmptySum = return []
 709     sum_tys (UnarySum r) = con_tys r
 710     sum_tys (Sum { repr_sel_ty = sel_ty
 711                  , repr_cons   = cons })
 712       = liftM (sel_ty :) (concatMapM con_tys cons)
 713
 714     con_tys (ConRepr _ r) = prod_tys r
 715
 716     prod_tys EmptyProd = return []
 717     prod_tys (UnaryProd r) = liftM singleton (comp_ty r)
 718     prod_tys (Prod { repr_comps = comps }) = mapM comp_ty comps
 719
 720     comp_ty r = mkPDataType (compOrigType r)
 721
 722
 723 buildTyConBindings :: TyCon -> TyCon -> TyCon -> TyCon -> SumRepr
 724                    -> VM Var
 725 buildTyConBindings orig_tc vect_tc prepr_tc pdata_tc repr
 726   = do
 727       vectDataConWorkers orig_tc vect_tc pdata_tc
 728       buildPADict vect_tc prepr_tc pdata_tc repr
 729
 730 vectDataConWorkers :: TyCon -> TyCon -> TyCon -> VM ()
 731 vectDataConWorkers orig_tc vect_tc arr_tc
 732   = do
 733       bs <- sequence
 734           . zipWith3 def_worker  (tyConDataCons orig_tc) rep_tys
 735           $ zipWith4 mk_data_con (tyConDataCons vect_tc)
 736                                  rep_tys
 737                                  (inits rep_tys)
 738                                  (tail $ tails rep_tys)
 739       mapM_ (uncurry hoistBinding) bs
 740   where
 741     tyvars   = tyConTyVars vect_tc
 742     var_tys  = mkTyVarTys tyvars
 743     ty_args  = map Type var_tys
 744     res_ty   = mkTyConApp vect_tc var_tys
 745
 746     cons     = tyConDataCons vect_tc
 747     arity    = length cons
 748     [arr_dc] = tyConDataCons arr_tc
 749
 750     rep_tys  = map dataConRepArgTys $ tyConDataCons vect_tc
 751
 752
 753     mk_data_con con tys pre post
 754       = liftM2 (,) (vect_data_con con)
 755                    (lift_data_con tys pre post (mkDataConTag con))
 756
 757     sel_replicate len tag
 758       | arity > 1 = do
 759                       rep <- builtin (selReplicate arity)
 760                       return [rep `mkApps` [len, tag]]
 761
 762       | otherwise = return []
 763
 764     vect_data_con con = return $ mkConApp con ty_args
 765     lift_data_con tys pre_tys post_tys tag
 766       = do
 767           len  <- builtin liftingContext
 768           args <- mapM (newLocalVar (fsLit "xs"))
 769                   =<< mapM mkPDataType tys
 770
 771           sel  <- sel_replicate (Var len) tag
 772
 773           pre   <- mapM emptyPD (concat pre_tys)
 774           post  <- mapM emptyPD (concat post_tys)
 775
 776           return . mkLams (len : args)
 777                  . wrapFamInstBody arr_tc var_tys
 778                  . mkConApp arr_dc
 779                  $ ty_args ++ sel ++ pre ++ map Var args ++ post
 780
 781     def_worker data_con arg_tys mk_body
 782       = do
 783           arity <- polyArity tyvars
 784           body <- closedV
 785                 . inBind orig_worker
 786                 . polyAbstract tyvars $ \args ->
 787                   liftM (mkLams (tyvars ++ args) . vectorised)
 788                 $ buildClosures tyvars [] arg_tys res_ty mk_body
 789
 790           raw_worker <- cloneId mkVectOcc orig_worker (exprType body)
 791           let vect_worker = raw_worker `setIdUnfolding`
 792                               mkInlineRule body (Just arity)
 793           defGlobalVar orig_worker vect_worker
 794           return (vect_worker, body)
 795       where
 796         orig_worker = dataConWorkId data_con
 797
 798 buildPADict :: TyCon -> TyCon -> TyCon -> SumRepr -> VM Var
 799 buildPADict vect_tc prepr_tc arr_tc repr
 800   = polyAbstract tvs $ \args ->
 801     do
 802       method_ids <- mapM (method args) paMethods
 803
 804       pa_tc  <- builtin paTyCon
 805       pa_con <- builtin paDataCon
 806       let dict = mkLams (tvs ++ args)
 807                $ mkConApp pa_con
 808                $ Type inst_ty : map (method_call args) method_ids
 809
 810           dfun_ty = mkForAllTys tvs
 811                   $ mkFunTys (map varType args) (mkTyConApp pa_tc [inst_ty])
 812
 813       raw_dfun <- newExportedVar dfun_name dfun_ty
 814       let dfun = raw_dfun `setIdUnfolding` mkDFunUnfolding pa_con method_ids
 815                           `setInlinePragma` dfunInlinePragma
 816
 817       hoistBinding dfun dict
 818       return dfun
 819   where
 820     tvs = tyConTyVars vect_tc
 821     arg_tys = mkTyVarTys tvs
 822     inst_ty = mkTyConApp vect_tc arg_tys
 823
 824     dfun_name = mkPADFunOcc (getOccName vect_tc)
 825
 826     method args (name, build)
 827       = localV
 828       $ do
 829           expr <- build vect_tc prepr_tc arr_tc repr
 830           let body = mkLams (tvs ++ args) expr
 831           raw_var <- newExportedVar (method_name name) (exprType body)
 832           let var = raw_var
 833                       `setIdUnfolding` mkInlineRule body (Just (length args))
 834                       `setInlinePragma` alwaysInlinePragma
 835           hoistBinding var body
 836           return var
 837
 838     method_call args id = mkApps (Var id) (map Type arg_tys ++ map Var args)
 839
 840     method_name name = mkVarOcc $ occNameString dfun_name ++ ('$' : name)
 841
 842
 843 paMethods :: [(String, TyCon -> TyCon -> TyCon -> SumRepr -> VM CoreExpr)]
 844 paMethods = [("dictPRepr",    buildPRDict),
 845              ("toPRepr",      buildToPRepr),
 846              ("fromPRepr",    buildFromPRepr),
 847              ("toArrPRepr",   buildToArrPRepr),
 848              ("fromArrPRepr", buildFromArrPRepr)]
 849
 850 -- | Split the given tycons into two sets depending on whether they have to be
 851 -- converted (first list) or not (second list). The first argument contains
 852 -- information about the conversion status of external tycons:
 853 --
 854 --   * tycons which have converted versions are mapped to True
 855 --   * tycons which are not changed by vectorisation are mapped to False
 856 --   * tycons which can't be converted are not elements of the map
 857 --
 858 classifyTyCons :: UniqFM Bool -> [TyConGroup] -> ([TyCon], [TyCon])
 859 classifyTyCons = classify [] []
 860   where
 861     classify conv keep _  [] = (conv, keep)
 862     classify conv keep cs ((tcs, ds) : rs)
 863       | can_convert && must_convert
 864         = classify (tcs ++ conv) keep (cs `addListToUFM` [(tc,True) | tc <- tcs]) rs
 865       | can_convert
 866         = classify conv (tcs ++ keep) (cs `addListToUFM` [(tc,False) | tc <- tcs]) rs
 867       | otherwise
 868         = classify conv keep cs rs
 869       where
 870         refs = ds `delListFromUniqSet` tcs
 871
 872         can_convert  = isNullUFM (refs `minusUFM` cs) && all convertable tcs
 873         must_convert = foldUFM (||) False (intersectUFM_C const cs refs)
 874
 875         convertable tc = isDataTyCon tc && all isVanillaDataCon (tyConDataCons tc)
 876
 877 -- | Compute mutually recursive groups of tycons in topological order
 878 --
 879 tyConGroups :: [TyCon] -> [TyConGroup]
 880 tyConGroups tcs = map mk_grp (stronglyConnCompFromEdgedVertices edges)
 881   where
 882     edges = [((tc, ds), tc, uniqSetToList ds) | tc <- tcs
 883                                 , let ds = tyConsOfTyCon tc]
 884
 885     mk_grp (AcyclicSCC (tc, ds)) = ([tc], ds)
 886     mk_grp (CyclicSCC els)       = (tcs, unionManyUniqSets dss)
 887       where
 888         (tcs, dss) = unzip els
 889
 890 tyConsOfTyCon :: TyCon -> UniqSet TyCon
 891 tyConsOfTyCon
 892   = tyConsOfTypes . concatMap dataConRepArgTys . tyConDataCons
 893
 894 tyConsOfType :: Type -> UniqSet TyCon
 895 tyConsOfType ty
 896   | Just ty' <- coreView ty    = tyConsOfType ty'
 897 tyConsOfType (TyVarTy _)       = emptyUniqSet
 898 tyConsOfType (TyConApp tc tys) = extend (tyConsOfTypes tys)
 899   where
 900     extend | isUnLiftedTyCon tc
 901            || isTupleTyCon   tc = id
 902
 903            | otherwise          = (`addOneToUniqSet` tc)
 904
 905 tyConsOfType (AppTy a b)       = tyConsOfType a `unionUniqSets` tyConsOfType b
 906 tyConsOfType (FunTy a b)       = (tyConsOfType a `unionUniqSets` tyConsOfType b)
 907                                  `addOneToUniqSet` funTyCon
 908 tyConsOfType (ForAllTy _ ty)   = tyConsOfType ty
 909 tyConsOfType other             = pprPanic "ClosureConv.tyConsOfType" $ ppr other
 910
 911 tyConsOfTypes :: [Type] -> UniqSet TyCon
 912 tyConsOfTypes = unionManyUniqSets . map tyConsOfType
 913
 914
 915 -- ----------------------------------------------------------------------------
 916 -- Conversions
 917
 918 fromVect :: Type -> CoreExpr -> VM CoreExpr
 919 fromVect ty expr | Just ty' <- coreView ty = fromVect ty' expr
 920 fromVect (FunTy arg_ty res_ty) expr
 921   = do
 922       arg     <- newLocalVar (fsLit "x") arg_ty
 923       varg    <- toVect arg_ty (Var arg)
 924       varg_ty <- vectType arg_ty
 925       vres_ty <- vectType res_ty
 926       apply   <- builtin applyVar
 927       body    <- fromVect res_ty
 928                $ Var apply `mkTyApps` [varg_ty, vres_ty] `mkApps` [expr, varg]
 929       return $ Lam arg body
 930 fromVect ty expr
 931   = identityConv ty >> return expr
 932
 933 toVect :: Type -> CoreExpr -> VM CoreExpr
 934 toVect ty expr = identityConv ty >> return expr
 935
 936 identityConv :: Type -> VM ()
 937 identityConv ty | Just ty' <- coreView ty = identityConv ty'
 938 identityConv (TyConApp tycon tys)
 939   = do
 940       mapM_ identityConv tys
 941       identityConvTyCon tycon
 942 identityConv _ = noV
 943
 944 identityConvTyCon :: TyCon -> VM ()
 945 identityConvTyCon tc
 946   | isBoxedTupleTyCon tc = return ()
 947   | isUnLiftedTyCon tc   = return ()
 948   | otherwise            = do
 949                              tc' <- maybeV (lookupTyCon tc)
 950                              if tc == tc' then return () else noV
 951