compiler/vectorise/VectType.hs

   1 {-# OPTIONS -w #-}
   2 -- The above warning supression flag is a temporary kludge.
   3 -- While working on this module you are encouraged to remove it and fix
   4 -- any warnings in the module. See
   5 --     http://hackage.haskell.org/trac/ghc/wiki/Commentary/CodingStyle#Warnings
   6 -- for details
   7
   8 module VectType ( vectTyCon, vectType, vectTypeEnv,
   9                   mkRepr, arrShapeTys, arrShapeVars, arrSelector,
  10                   PAInstance, buildPADict,
  11                   fromVect )
  12 where
  13
  14 import VectMonad
  15 import VectUtils
  16 import VectCore
  17
  18 import HscTypes          ( TypeEnv, extendTypeEnvList, typeEnvTyCons )
  19 import CoreSyn
  20 import CoreUtils
  21 import BuildTyCl
  22 import DataCon
  23 import TyCon
  24 import Type
  25 import TypeRep
  26 import Coercion
  27 import FamInstEnv        ( FamInst, mkLocalFamInst )
  28 import InstEnv           ( Instance, mkLocalInstance, instanceDFunId )
  29 import OccName
  30 import MkId
  31 import BasicTypes        ( StrictnessMark(..), OverlapFlag(..), boolToRecFlag )
  32 import Var               ( Var )
  33 import Id                ( mkWildId )
  34 import Name              ( Name, getOccName )
  35 import NameEnv
  36 import TysWiredIn
  37 import TysPrim           ( intPrimTy )
  38
  39 import Unique
  40 import UniqFM
  41 import UniqSet
  42 import Util              ( singleton )
  43 import Digraph           ( SCC(..), stronglyConnComp )
  44
  45 import Outputable
  46 import FastString
  47
  48 import Control.Monad  ( liftM, liftM2, zipWithM, zipWithM_, mapAndUnzipM )
  49 import Data.List      ( inits, tails, zipWith4, zipWith5 )
  50
  51 -- ----------------------------------------------------------------------------
  52 -- Types
  53
  54 vectTyCon :: TyCon -> VM TyCon
  55 vectTyCon tc
  56   | isFunTyCon tc        = builtin closureTyCon
  57   | isBoxedTupleTyCon tc = return tc
  58   | isUnLiftedTyCon tc   = return tc
  59   | otherwise = do
  60                   r <- lookupTyCon tc
  61                   case r of
  62                     Just tc' -> return tc'
  63
  64                     -- FIXME: just for now
  65                     Nothing  -> pprTrace "ccTyCon:" (ppr tc) $ return tc
  66
  67 vectType :: Type -> VM Type
  68 vectType ty | Just ty' <- coreView ty = vectType ty'
  69 vectType (TyVarTy tv) = return $ TyVarTy tv
  70 vectType (AppTy ty1 ty2) = liftM2 AppTy (vectType ty1) (vectType ty2)
  71 vectType (TyConApp tc tys) = liftM2 TyConApp (vectTyCon tc) (mapM vectType tys)
  72 vectType (FunTy ty1 ty2)   = liftM2 TyConApp (builtin closureTyCon)
  73                                              (mapM vectAndBoxType [ty1,ty2])
  74 vectType ty@(ForAllTy _ _)
  75   = do
  76       mdicts   <- mapM paDictArgType tyvars
  77       mono_ty' <- vectType mono_ty
  78       return $ tyvars `mkForAllTys` ([dict | Just dict <- mdicts] `mkFunTys` mono_ty')
  79   where
  80     (tyvars, mono_ty) = splitForAllTys ty
  81
  82 vectType ty = pprPanic "vectType:" (ppr ty)
  83
  84 vectAndBoxType :: Type -> VM Type
  85 vectAndBoxType ty = vectType ty >>= boxType
  86
  87 -- ----------------------------------------------------------------------------
  88 -- Boxing
  89
  90 boxType :: Type -> VM Type
  91 boxType ty
  92   | Just (tycon, []) <- splitTyConApp_maybe ty
  93   , isUnLiftedTyCon tycon
  94   = do
  95       r <- lookupBoxedTyCon tycon
  96       case r of
  97         Just tycon' -> return $ mkTyConApp tycon' []
  98         Nothing     -> return ty
  99 boxType ty = return ty
 100
 101 -- ----------------------------------------------------------------------------
 102 -- Type definitions
 103
 104 type TyConGroup = ([TyCon], UniqSet TyCon)
 105
 106 data PAInstance = PAInstance {
 107                     painstDFun      :: Var
 108                   , painstOrigTyCon :: TyCon
 109                   , painstVectTyCon :: TyCon
 110                   , painstArrTyCon  :: TyCon
 111                   }
 112
 113 vectTypeEnv :: TypeEnv -> VM (TypeEnv, [FamInst], [(Var, CoreExpr)])
 114 vectTypeEnv env
 115   = do
 116       cs <- readGEnv $ mk_map . global_tycons
 117       let (conv_tcs, keep_tcs) = classifyTyCons cs groups
 118           keep_dcs             = concatMap tyConDataCons keep_tcs
 119       zipWithM_ defTyCon   keep_tcs keep_tcs
 120       zipWithM_ defDataCon keep_dcs keep_dcs
 121       new_tcs <- vectTyConDecls conv_tcs
 122
 123       let orig_tcs = keep_tcs ++ conv_tcs
 124           vect_tcs  = keep_tcs ++ new_tcs
 125
 126       repr_tcs <- zipWithM buildPReprTyCon   orig_tcs vect_tcs
 127       parr_tcs <- zipWithM buildPArrayTyCon orig_tcs vect_tcs
 128       dfuns    <- mapM mkPADFun vect_tcs
 129       defTyConPAs (zip vect_tcs dfuns)
 130       binds    <- sequence (zipWith5 buildTyConBindings orig_tcs
 131                                                         vect_tcs
 132                                                         repr_tcs
 133                                                         parr_tcs
 134                                                         dfuns)
 135
 136       let all_new_tcs = new_tcs ++ repr_tcs ++ parr_tcs
 137
 138       let new_env = extendTypeEnvList env
 139                        (map ATyCon all_new_tcs
 140                         ++ [ADataCon dc | tc <- all_new_tcs
 141                                         , dc <- tyConDataCons tc])
 142
 143       return (new_env, map mkLocalFamInst (repr_tcs ++ parr_tcs), concat binds)
 144   where
 145     tycons = typeEnvTyCons env
 146     groups = tyConGroups tycons
 147
 148     mk_map env = listToUFM_Directly [(u, getUnique n /= u) | (u,n) <- nameEnvUniqueElts env]
 149
 150     keep_tc tc = let dcs = tyConDataCons tc
 151                  in
 152                  defTyCon tc tc >> zipWithM_ defDataCon dcs dcs
 153
 154
 155 vectTyConDecls :: [TyCon] -> VM [TyCon]
 156 vectTyConDecls tcs = fixV $ \tcs' ->
 157   do
 158     mapM_ (uncurry defTyCon) (lazy_zip tcs tcs')
 159     mapM vectTyConDecl tcs
 160   where
 161     lazy_zip [] _ = []
 162     lazy_zip (x:xs) ~(y:ys) = (x,y) : lazy_zip xs ys
 163
 164 vectTyConDecl :: TyCon -> VM TyCon
 165 vectTyConDecl tc
 166   = do
 167       name' <- cloneName mkVectTyConOcc name
 168       rhs'  <- vectAlgTyConRhs (algTyConRhs tc)
 169
 170       liftDs $ buildAlgTyCon name'
 171                              tyvars
 172                              []           -- no stupid theta
 173                              rhs'
 174                              rec_flag     -- FIXME: is this ok?
 175                              False        -- FIXME: no generics
 176                              False        -- not GADT syntax
 177                              Nothing      -- not a family instance
 178   where
 179     name   = tyConName tc
 180     tyvars = tyConTyVars tc
 181     rec_flag = boolToRecFlag (isRecursiveTyCon tc)
 182
 183 vectAlgTyConRhs :: AlgTyConRhs -> VM AlgTyConRhs
 184 vectAlgTyConRhs (DataTyCon { data_cons = data_cons
 185                            , is_enum   = is_enum
 186                            })
 187   = do
 188       data_cons' <- mapM vectDataCon data_cons
 189       zipWithM_ defDataCon data_cons data_cons'
 190       return $ DataTyCon { data_cons = data_cons'
 191                          , is_enum   = is_enum
 192                          }
 193
 194 vectDataCon :: DataCon -> VM DataCon
 195 vectDataCon dc
 196   | not . null $ dataConExTyVars dc = pprPanic "vectDataCon: existentials" (ppr dc)
 197   | not . null $ dataConEqSpec   dc = pprPanic "vectDataCon: eq spec" (ppr dc)
 198   | otherwise
 199   = do
 200       name'    <- cloneName mkVectDataConOcc name
 201       tycon'   <- vectTyCon tycon
 202       arg_tys  <- mapM vectType rep_arg_tys
 203
 204       liftDs $ buildDataCon name'
 205                             False           -- not infix
 206                             (map (const NotMarkedStrict) arg_tys)
 207                             []              -- no labelled fields
 208                             univ_tvs
 209                             []              -- no existential tvs for now
 210                             []              -- no eq spec for now
 211                             []              -- no context
 212                             arg_tys
 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 buildPReprTyCon :: TyCon -> TyCon -> VM TyCon
 225 buildPReprTyCon orig_tc vect_tc
 226   = do
 227       name     <- cloneName mkPReprTyConOcc (tyConName orig_tc)
 228       rhs_ty   <- buildPReprType vect_tc
 229       prepr_tc <- builtin preprTyCon
 230       liftDs $ buildSynTyCon name
 231                              tyvars
 232                              (SynonymTyCon rhs_ty)
 233                              (Just $ mk_fam_inst prepr_tc vect_tc)
 234   where
 235     tyvars = tyConTyVars vect_tc
 236
 237
 238 data Repr = ProdRepr {
 239               prod_components   :: [Type]
 240             , prod_tycon        :: TyCon
 241             , prod_data_con     :: DataCon
 242             , prod_arr_tycon    :: TyCon
 243             , prod_arr_data_con :: DataCon
 244             }
 245
 246           | SumRepr {
 247               sum_components    :: [Repr]
 248             , sum_tycon         :: TyCon
 249             , sum_arr_tycon     :: TyCon
 250             , sum_arr_data_con  :: DataCon
 251             }
 252
 253           | IdRepr Type
 254
 255           | VoidRepr {
 256               void_tycon        :: TyCon
 257             , void_bottom       :: CoreExpr
 258             }
 259
 260           | EnumRepr {
 261               enum_tycon        :: TyCon
 262             , enum_data_con     :: DataCon
 263             , enum_arr_tycon    :: TyCon
 264             , enum_arr_data_con :: DataCon
 265             }
 266
 267 voidRepr :: VM Repr
 268 voidRepr
 269   = do
 270       tycon <- builtin voidTyCon
 271       var   <- builtin voidVar
 272       return $ VoidRepr {
 273                  void_tycon  = tycon
 274                , void_bottom = Var var
 275                }
 276
 277 enumRepr :: VM Repr
 278 enumRepr
 279   = do
 280       tycon <- builtin enumerationTyCon
 281       let [data_con] = tyConDataCons tycon
 282       (arr_tycon, _) <- parrayReprTyCon (mkTyConApp tycon [])
 283       let [arr_data_con] = tyConDataCons arr_tycon
 284
 285       return $ EnumRepr {
 286                  enum_tycon       = tycon
 287                , enum_data_con     = data_con
 288                , enum_arr_tycon    = arr_tycon
 289                , enum_arr_data_con = arr_data_con
 290                }
 291
 292 unboxedProductRepr :: [Type] -> VM Repr
 293 unboxedProductRepr []   = voidRepr
 294 unboxedProductRepr [ty] = return $ IdRepr ty
 295 unboxedProductRepr tys  = boxedProductRepr tys
 296
 297 boxedProductRepr :: [Type] -> VM Repr
 298 boxedProductRepr tys
 299   = do
 300       tycon <- builtin (prodTyCon arity)
 301       let [data_con] = tyConDataCons tycon
 302
 303       tys' <- mapM boxType tys
 304       (arr_tycon, _) <- parrayReprTyCon $ mkTyConApp tycon tys'
 305       let [arr_data_con] = tyConDataCons arr_tycon
 306
 307       return $ ProdRepr {
 308                  prod_components   = tys
 309                , prod_tycon        = tycon
 310                , prod_data_con     = data_con
 311                , prod_arr_tycon    = arr_tycon
 312                , prod_arr_data_con = arr_data_con
 313                }
 314   where
 315     arity = length tys
 316
 317 sumRepr :: [Repr] -> VM Repr
 318 sumRepr []     = voidRepr
 319 sumRepr [repr] = boxRepr repr
 320 sumRepr reprs
 321   = do
 322       tycon <- builtin (sumTyCon arity)
 323       (arr_tycon, _) <- parrayReprTyCon
 324                       . mkTyConApp tycon
 325                       $ map reprType reprs
 326
 327       let [arr_data_con] = tyConDataCons arr_tycon
 328
 329       return $ SumRepr {
 330                  sum_components   = reprs
 331                , sum_tycon        = tycon
 332                , sum_arr_tycon    = arr_tycon
 333                , sum_arr_data_con = arr_data_con
 334                }
 335   where
 336     arity = length reprs
 337
 338 splitSumRepr :: Repr -> [Repr]
 339 splitSumRepr (SumRepr { sum_components = reprs }) = reprs
 340 splitSumRepr repr                                 = [repr]
 341
 342 boxRepr :: Repr -> VM Repr
 343 boxRepr (VoidRepr {}) = boxedProductRepr []
 344 boxRepr (IdRepr ty)   = boxedProductRepr [ty]
 345 boxRepr repr          = return repr
 346
 347 reprType :: Repr -> Type
 348 reprType (ProdRepr { prod_tycon = tycon, prod_components = tys })
 349   = mkTyConApp tycon tys
 350 reprType (SumRepr { sum_tycon = tycon, sum_components = reprs })
 351   = mkTyConApp tycon (map reprType reprs)
 352 reprType (IdRepr ty) = ty
 353 reprType (VoidRepr { void_tycon = tycon }) = mkTyConApp tycon []
 354 reprType (EnumRepr { enum_tycon = tycon }) = mkTyConApp tycon []
 355
 356 arrReprType :: Repr -> VM Type
 357 arrReprType = mkPArrayType . reprType
 358
 359 arrShapeTys :: Repr -> VM [Type]
 360 arrShapeTys (SumRepr  {}) = sumShapeTys
 361 arrShapeTys (ProdRepr {}) = return [intPrimTy]
 362 arrShapeTys (IdRepr _)    = return []
 363 arrShapeTys (VoidRepr {}) = return [intPrimTy]
 364 arrShapeTys (EnumRepr {}) = sumShapeTys
 365
 366 sumShapeTys :: VM [Type]
 367 sumShapeTys = do
 368                 int_arr <- builtin intPrimArrayTy
 369                 return [intPrimTy, int_arr, int_arr]
 370
 371
 372 arrShapeVars :: Repr -> VM [Var]
 373 arrShapeVars repr = mapM (newLocalVar (fsLit "sh")) =<< arrShapeTys repr
 374
 375 replicateShape :: Repr -> CoreExpr -> CoreExpr -> VM [CoreExpr]
 376 replicateShape (ProdRepr {}) len _   = return [len]
 377 replicateShape (SumRepr {})  len tag = replicateSumShape len tag
 378 replicateShape (IdRepr _) _ _        = return []
 379 replicateShape (VoidRepr {}) len _   = return [len]
 380 replicateShape (EnumRepr {}) len tag = replicateSumShape len tag
 381
 382 replicateSumShape :: CoreExpr -> CoreExpr -> VM [CoreExpr]
 383 replicateSumShape len tag
 384   = do
 385       rep <- builtin replicatePAIntPrimVar
 386       up  <- builtin upToPAIntPrimVar
 387       return [len, Var rep `mkApps` [len, tag], Var up `App` len]
 388
 389 arrSelector :: Repr -> [CoreExpr] -> VM (CoreExpr, CoreExpr, CoreExpr)
 390 arrSelector (SumRepr {})  [len, sel, is] = return (len, sel, is)
 391 arrSelector (EnumRepr {}) [len, sel, is] = return (len, sel, is)
 392
 393 emptyArrRepr :: Repr -> VM [CoreExpr]
 394 emptyArrRepr (SumRepr { sum_components = prods })
 395   = liftM concat $ mapM emptyArrRepr prods
 396 emptyArrRepr (ProdRepr { prod_components = [] })
 397   = return [Var unitDataConId]
 398 emptyArrRepr (ProdRepr { prod_components = tys })
 399   = mapM emptyPA tys
 400 emptyArrRepr (IdRepr ty)
 401   = liftM singleton $ emptyPA ty
 402 emptyArrRepr (VoidRepr { void_tycon = tycon })
 403   = liftM singleton $ emptyPA (mkTyConApp tycon [])
 404 emptyArrRepr (EnumRepr { enum_tycon = tycon })
 405   = return []
 406
 407 arrReprTys :: Repr -> VM [Type]
 408 arrReprTys (SumRepr { sum_components = reprs })
 409   = liftM concat $ mapM arrReprTys reprs
 410 arrReprTys (ProdRepr { prod_components = [] })
 411   = return [unitTy]
 412 arrReprTys (ProdRepr { prod_components = tys })
 413   = mapM mkPArrayType tys
 414 arrReprTys (IdRepr ty)
 415   = liftM singleton $ mkPArrayType ty
 416 arrReprTys (VoidRepr { void_tycon = tycon })
 417   = liftM singleton $ mkPArrayType (mkTyConApp tycon [])
 418 arrReprTys (EnumRepr {})
 419   = return []
 420
 421 arrReprTys' :: Repr -> VM [[Type]]
 422 arrReprTys' (SumRepr { sum_components = reprs })
 423   = mapM arrReprTys reprs
 424 arrReprTys' repr = liftM singleton $ arrReprTys repr
 425
 426 arrReprVars :: Repr -> VM [[Var]]
 427 arrReprVars repr
 428   = mapM (mapM (newLocalVar (fsLit "rs"))) =<< arrReprTys' repr
 429
 430 mkRepr :: TyCon -> VM Repr
 431 mkRepr vect_tc
 432   | [tys] <- rep_tys = boxedProductRepr tys
 433   -- | all null rep_tys = enumRepr
 434   | otherwise        = sumRepr =<< mapM unboxedProductRepr rep_tys
 435   where
 436     rep_tys = map dataConRepArgTys $ tyConDataCons vect_tc
 437
 438 buildPReprType :: TyCon -> VM Type
 439 buildPReprType = liftM reprType . mkRepr
 440
 441 buildToPRepr :: Repr -> TyCon -> TyCon -> TyCon -> VM CoreExpr
 442 buildToPRepr repr vect_tc prepr_tc _
 443   = do
 444       arg    <- newLocalVar (fsLit "x") arg_ty
 445       result <- to_repr repr (Var arg)
 446
 447       return . Lam arg
 448              . wrapFamInstBody prepr_tc var_tys
 449              $ result
 450   where
 451     var_tys = mkTyVarTys $ tyConTyVars vect_tc
 452     arg_ty  = mkTyConApp vect_tc var_tys
 453     res_ty  = reprType repr
 454
 455     cons    = tyConDataCons vect_tc
 456     [con]   = cons
 457
 458     to_repr (SumRepr { sum_components = prods
 459                      , sum_tycon      = tycon })
 460             expr
 461       = do
 462           (vars, bodies) <- mapAndUnzipM to_unboxed prods
 463           return . Case expr (mkWildId (exprType expr)) res_ty
 464                  $ zipWith4 mk_alt cons vars (tyConDataCons tycon) bodies
 465       where
 466         mk_alt con vars sum_con body
 467           = (DataAlt con, vars, mkConApp sum_con (ty_args ++ [body]))
 468
 469         ty_args = map (Type . reprType) prods
 470
 471     to_repr (EnumRepr { enum_data_con = data_con }) expr
 472       = return . Case expr (mkWildId (exprType expr)) res_ty
 473                $ map mk_alt cons
 474       where
 475         mk_alt con = (DataAlt con, [], mkConApp data_con [mkDataConTag con])
 476
 477     to_repr prod expr
 478       = do
 479           (vars, body) <- to_unboxed prod
 480           return $ Case expr (mkWildId (exprType expr)) res_ty
 481                    [(DataAlt con, vars, body)]
 482
 483     to_unboxed (ProdRepr { prod_components = tys
 484                          , prod_data_con   = data_con })
 485       = do
 486           vars <- mapM (newLocalVar (fsLit "r")) tys
 487           return (vars, mkConApp data_con (map Type tys ++ map Var vars))
 488
 489     to_unboxed (IdRepr ty)
 490       = do
 491           var <- newLocalVar (fsLit "y") ty
 492           return ([var], Var var)
 493
 494     to_unboxed (VoidRepr { void_bottom = bottom })
 495       = return ([], bottom)
 496
 497
 498 buildFromPRepr :: Repr -> TyCon -> TyCon -> TyCon -> VM CoreExpr
 499 buildFromPRepr repr vect_tc prepr_tc _
 500   = do
 501       arg_ty <- mkPReprType res_ty
 502       arg    <- newLocalVar (fsLit "x") arg_ty
 503
 504       liftM (Lam arg)
 505            . from_repr repr
 506            $ unwrapFamInstScrut prepr_tc var_tys (Var arg)
 507   where
 508     var_tys = mkTyVarTys $ tyConTyVars vect_tc
 509     res_ty  = mkTyConApp vect_tc var_tys
 510
 511     cons    = map (`mkConApp` map Type var_tys) (tyConDataCons vect_tc)
 512     [con]   = cons
 513
 514     from_repr repr@(SumRepr { sum_components = prods
 515                             , sum_tycon      = tycon })
 516               expr
 517       = do
 518           vars   <- mapM (newLocalVar (fsLit "x")) (map reprType prods)
 519           bodies <- sequence . zipWith3 from_unboxed prods cons
 520                              $ map Var vars
 521           return . Case expr (mkWildId (reprType repr)) res_ty
 522                  $ zipWith3 sum_alt (tyConDataCons tycon) vars bodies
 523       where
 524         sum_alt data_con var body = (DataAlt data_con, [var], body)
 525
 526     from_repr repr@(EnumRepr { enum_data_con = data_con }) expr
 527       = do
 528           var <- newLocalVar (fsLit "n") intPrimTy
 529
 530           let res = Case (Var var) (mkWildId intPrimTy) res_ty
 531                   $ (DEFAULT, [], error_expr)
 532                   : zipWith mk_alt (tyConDataCons vect_tc) cons
 533
 534           return $ Case expr (mkWildId (reprType repr)) res_ty
 535                    [(DataAlt data_con, [var], res)]
 536       where
 537         mk_alt data_con con = (LitAlt (mkDataConTagLit data_con), [], con)
 538
 539         error_expr = mkRuntimeErrorApp rUNTIME_ERROR_ID res_ty
 540                    . showSDoc
 541                    $ sep [text "Invalid NDP representation of", ppr vect_tc]
 542
 543     from_repr repr expr = from_unboxed repr con expr
 544
 545     from_unboxed prod@(ProdRepr { prod_components = tys
 546                                 , prod_data_con   = data_con })
 547               con
 548               expr
 549       = do
 550           vars <- mapM (newLocalVar (fsLit "y")) tys
 551           return $ Case expr (mkWildId (reprType prod)) res_ty
 552                    [(DataAlt data_con, vars, con `mkVarApps` vars)]
 553
 554     from_unboxed (IdRepr _) con expr
 555        = return $ con `App` expr
 556
 557     from_unboxed (VoidRepr {}) con expr
 558        = return con
 559
 560 buildToArrPRepr :: Repr -> TyCon -> TyCon -> TyCon -> VM CoreExpr
 561 buildToArrPRepr repr vect_tc prepr_tc arr_tc
 562   = do
 563       arg_ty     <- mkPArrayType el_ty
 564       arg        <- newLocalVar (fsLit "xs") arg_ty
 565
 566       res_ty     <- mkPArrayType (reprType repr)
 567
 568       shape_vars <- arrShapeVars repr
 569       repr_vars  <- arrReprVars  repr
 570
 571       parray_co  <- mkBuiltinCo parrayTyCon
 572
 573       let Just repr_co = tyConFamilyCoercion_maybe prepr_tc
 574           co           = mkAppCoercion parray_co
 575                        . mkSymCoercion
 576                        $ mkTyConApp repr_co var_tys
 577
 578           scrut   = unwrapFamInstScrut arr_tc var_tys (Var arg)
 579
 580       result <- to_repr shape_vars repr_vars repr
 581
 582       return . Lam arg
 583              . mkCoerce co
 584              $ Case scrut (mkWildId (mkTyConApp arr_tc var_tys)) res_ty
 585                [(DataAlt arr_dc, shape_vars ++ concat repr_vars, result)]
 586   where
 587     var_tys = mkTyVarTys $ tyConTyVars vect_tc
 588     el_ty   = mkTyConApp vect_tc var_tys
 589
 590     [arr_dc] = tyConDataCons arr_tc
 591
 592     to_repr shape_vars@(len_var : _)
 593             repr_vars
 594             (SumRepr { sum_components   = prods
 595                      , sum_arr_tycon    = tycon
 596                      , sum_arr_data_con = data_con })
 597       = do
 598           exprs <- zipWithM to_prod repr_vars prods
 599
 600           return . wrapFamInstBody tycon tys
 601                  . mkConApp data_con
 602                  $ map Type tys ++ map Var shape_vars ++ exprs
 603       where
 604         tys = map reprType prods
 605
 606     to_repr [len_var]
 607             [repr_vars]
 608             (ProdRepr { prod_components   = tys
 609                       , prod_arr_tycon    = tycon
 610                       , prod_arr_data_con = data_con })
 611        = return . wrapFamInstBody tycon tys
 612                 . mkConApp data_con
 613                 $ map Type tys ++ map Var (len_var : repr_vars)
 614
 615     to_repr shape_vars
 616             _
 617             (EnumRepr { enum_arr_tycon    = tycon
 618                       , enum_arr_data_con = data_con })
 619       = return . wrapFamInstBody tycon []
 620                . mkConApp data_con
 621                 $ map Var shape_vars
 622
 623     to_prod repr_vars@(r : _)
 624             (ProdRepr { prod_components   = tys@(ty : _)
 625                       , prod_arr_tycon    = tycon
 626                       , prod_arr_data_con = data_con })
 627       = do
 628           len <- lengthPA ty (Var r)
 629           return . wrapFamInstBody tycon tys
 630                  . mkConApp data_con
 631                  $ map Type tys ++ len : map Var repr_vars
 632
 633     to_prod [var] (IdRepr ty)   = return (Var var)
 634     to_prod [var] (VoidRepr {}) = return (Var var)
 635
 636
 637 buildFromArrPRepr :: Repr -> TyCon -> TyCon -> TyCon -> VM CoreExpr
 638 buildFromArrPRepr repr vect_tc prepr_tc arr_tc
 639   = do
 640       arg_ty     <- mkPArrayType =<< mkPReprType el_ty
 641       arg        <- newLocalVar (fsLit "xs") arg_ty
 642
 643       res_ty     <- mkPArrayType el_ty
 644
 645       shape_vars <- arrShapeVars repr
 646       repr_vars  <- arrReprVars  repr
 647
 648       parray_co  <- mkBuiltinCo parrayTyCon
 649
 650       let Just repr_co = tyConFamilyCoercion_maybe prepr_tc
 651           co           = mkAppCoercion parray_co
 652                        $ mkTyConApp repr_co var_tys
 653
 654           scrut  = mkCoerce co (Var arg)
 655
 656           result = wrapFamInstBody arr_tc var_tys
 657                  . mkConApp arr_dc
 658                  $ map Type var_tys ++ map Var (shape_vars ++ concat repr_vars)
 659
 660       liftM (Lam arg)
 661             (from_repr repr scrut shape_vars repr_vars res_ty result)
 662   where
 663     var_tys = mkTyVarTys $ tyConTyVars vect_tc
 664     el_ty   = mkTyConApp vect_tc var_tys
 665
 666     [arr_dc] = tyConDataCons arr_tc
 667
 668     from_repr (SumRepr { sum_components   = prods
 669                        , sum_arr_tycon    = tycon
 670                        , sum_arr_data_con = data_con })
 671               expr
 672               shape_vars
 673               repr_vars
 674               res_ty
 675               body
 676       = do
 677           vars <- mapM (newLocalVar (fsLit "xs")) =<< mapM arrReprType prods
 678           result <- go prods repr_vars vars body
 679
 680           let scrut = unwrapFamInstScrut tycon ty_args expr
 681           return . Case scrut (mkWildId scrut_ty) res_ty
 682                  $ [(DataAlt data_con, shape_vars ++ vars, result)]
 683       where
 684         ty_args  = map reprType prods
 685         scrut_ty = mkTyConApp tycon ty_args
 686
 687         go [] [] [] body = return body
 688         go (prod : prods) (repr_vars : rss) (var : vars) body
 689           = do
 690               shape_vars <- mapM (newLocalVar (fsLit "s")) =<< arrShapeTys prod
 691
 692               from_prod prod (Var var) shape_vars repr_vars res_ty
 693                 =<< go prods rss vars body
 694
 695     from_repr repr expr shape_vars [repr_vars] res_ty body
 696       = from_prod repr expr shape_vars repr_vars res_ty body
 697
 698     from_prod prod@(ProdRepr { prod_components = tys
 699                              , prod_arr_tycon  = tycon
 700                              , prod_arr_data_con = data_con })
 701               expr
 702               shape_vars
 703               repr_vars
 704               res_ty
 705               body
 706       = do
 707           let scrut    = unwrapFamInstScrut tycon tys expr
 708               scrut_ty = mkTyConApp tycon tys
 709           ty <- arrReprType prod
 710
 711           return $ Case scrut (mkWildId scrut_ty) res_ty
 712                    [(DataAlt data_con, shape_vars ++ repr_vars, body)]
 713
 714     from_prod (EnumRepr { enum_arr_tycon = tycon
 715                         , enum_arr_data_con = data_con })
 716               expr
 717               shape_vars
 718               _
 719               res_ty
 720               body
 721       = let scrut    = unwrapFamInstScrut tycon [] expr
 722             scrut_ty = mkTyConApp tycon []
 723         in
 724         return $ Case scrut (mkWildId scrut_ty) res_ty
 725                  [(DataAlt data_con, shape_vars, body)]
 726
 727     from_prod (IdRepr ty)
 728               expr
 729               shape_vars
 730               [repr_var]
 731               res_ty
 732               body
 733       = return $ Let (NonRec repr_var expr) body
 734
 735     from_prod (VoidRepr {})
 736               expr
 737               shape_vars
 738               [repr_var]
 739               res_ty
 740               body
 741       = return $ Let (NonRec repr_var expr) body
 742
 743 buildPRDictRepr :: Repr -> VM CoreExpr
 744 buildPRDictRepr (VoidRepr { void_tycon = tycon })
 745   = prDFunOfTyCon tycon
 746 buildPRDictRepr (IdRepr ty) = mkPR ty
 747 buildPRDictRepr (ProdRepr {
 748                    prod_components = tys
 749                  , prod_tycon      = tycon
 750                  })
 751   = do
 752       prs  <- mapM mkPR tys
 753       dfun <- prDFunOfTyCon tycon
 754       return $ dfun `mkTyApps` tys `mkApps` prs
 755
 756 buildPRDictRepr (SumRepr {
 757                    sum_components = prods
 758                  , sum_tycon      = tycon })
 759   = do
 760       prs  <- mapM buildPRDictRepr prods
 761       dfun <- prDFunOfTyCon tycon
 762       return $ dfun `mkTyApps` map reprType prods `mkApps` prs
 763
 764 buildPRDictRepr (EnumRepr { enum_tycon = tycon })
 765   = prDFunOfTyCon tycon
 766
 767 buildPRDict :: Repr -> TyCon -> TyCon -> TyCon -> VM CoreExpr
 768 buildPRDict repr vect_tc prepr_tc _
 769   = do
 770       dict  <- buildPRDictRepr repr
 771
 772       pr_co <- mkBuiltinCo prTyCon
 773       let co = mkAppCoercion pr_co
 774              . mkSymCoercion
 775              $ mkTyConApp arg_co var_tys
 776
 777       return $ mkCoerce co dict
 778   where
 779     var_tys = mkTyVarTys $ tyConTyVars vect_tc
 780
 781     Just arg_co = tyConFamilyCoercion_maybe prepr_tc
 782
 783 buildPArrayTyCon :: TyCon -> TyCon -> VM TyCon
 784 buildPArrayTyCon orig_tc vect_tc = fixV $ \repr_tc ->
 785   do
 786     name'  <- cloneName mkPArrayTyConOcc orig_name
 787     rhs    <- buildPArrayTyConRhs orig_name vect_tc repr_tc
 788     parray <- builtin parrayTyCon
 789
 790     liftDs $ buildAlgTyCon name'
 791                            tyvars
 792                            []          -- no stupid theta
 793                            rhs
 794                            rec_flag    -- FIXME: is this ok?
 795                            False       -- FIXME: no generics
 796                            False       -- not GADT syntax
 797                            (Just $ mk_fam_inst parray vect_tc)
 798   where
 799     orig_name = tyConName orig_tc
 800     tyvars = tyConTyVars vect_tc
 801     rec_flag = boolToRecFlag (isRecursiveTyCon vect_tc)
 802
 803
 804 buildPArrayTyConRhs :: Name -> TyCon -> TyCon -> VM AlgTyConRhs
 805 buildPArrayTyConRhs orig_name vect_tc repr_tc
 806   = do
 807       data_con <- buildPArrayDataCon orig_name vect_tc repr_tc
 808       return $ DataTyCon { data_cons = [data_con], is_enum = False }
 809
 810 buildPArrayDataCon :: Name -> TyCon -> TyCon -> VM DataCon
 811 buildPArrayDataCon orig_name vect_tc repr_tc
 812   = do
 813       dc_name  <- cloneName mkPArrayDataConOcc orig_name
 814       repr     <- mkRepr vect_tc
 815
 816       shape_tys <- arrShapeTys repr
 817       repr_tys  <- arrReprTys  repr
 818
 819       let tys = shape_tys ++ repr_tys
 820
 821       liftDs $ buildDataCon dc_name
 822                             False                  -- not infix
 823                             (map (const NotMarkedStrict) tys)
 824                             []                     -- no field labels
 825                             (tyConTyVars vect_tc)
 826                             []                     -- no existentials
 827                             []                     -- no eq spec
 828                             []                     -- no context
 829                             tys
 830                             repr_tc
 831
 832 mkPADFun :: TyCon -> VM Var
 833 mkPADFun vect_tc
 834   = newExportedVar (mkPADFunOcc $ getOccName vect_tc) =<< paDFunType vect_tc
 835
 836 buildTyConBindings :: TyCon -> TyCon -> TyCon -> TyCon -> Var
 837                    -> VM [(Var, CoreExpr)]
 838 buildTyConBindings orig_tc vect_tc prepr_tc arr_tc dfun
 839   = do
 840       repr  <- mkRepr vect_tc
 841       vectDataConWorkers repr orig_tc vect_tc arr_tc
 842       dict <- buildPADict repr vect_tc prepr_tc arr_tc dfun
 843       binds <- takeHoisted
 844       return $ (dfun, dict) : binds
 845   where
 846     orig_dcs = tyConDataCons orig_tc
 847     vect_dcs = tyConDataCons vect_tc
 848     [arr_dc] = tyConDataCons arr_tc
 849
 850     repr_tys = map dataConRepArgTys vect_dcs
 851
 852 vectDataConWorkers :: Repr -> TyCon -> TyCon -> TyCon
 853                    -> VM ()
 854 vectDataConWorkers repr orig_tc vect_tc arr_tc
 855   = do
 856       bs <- sequence
 857           . zipWith3 def_worker  (tyConDataCons orig_tc) rep_tys
 858           $ zipWith4 mk_data_con (tyConDataCons vect_tc)
 859                                  rep_tys
 860                                  (inits reprs)
 861                                  (tail $ tails reprs)
 862       mapM_ (uncurry hoistBinding) bs
 863   where
 864     tyvars   = tyConTyVars vect_tc
 865     var_tys  = mkTyVarTys tyvars
 866     ty_args  = map Type var_tys
 867
 868     res_ty   = mkTyConApp vect_tc var_tys
 869
 870     rep_tys  = map dataConRepArgTys $ tyConDataCons vect_tc
 871     reprs    = splitSumRepr repr
 872
 873     [arr_dc] = tyConDataCons arr_tc
 874
 875     mk_data_con con tys pre post
 876       = liftM2 (,) (vect_data_con con)
 877                    (lift_data_con tys pre post (mkDataConTag con))
 878
 879     vect_data_con con = return $ mkConApp con ty_args
 880     lift_data_con tys pre_reprs post_reprs tag
 881       = do
 882           len  <- builtin liftingContext
 883           args <- mapM (newLocalVar (fsLit "xs"))
 884                   =<< mapM mkPArrayType tys
 885
 886           shape <- replicateShape repr (Var len) tag
 887           repr  <- mk_arr_repr (Var len) (map Var args)
 888
 889           pre   <- liftM concat $ mapM emptyArrRepr pre_reprs
 890           post  <- liftM concat $ mapM emptyArrRepr post_reprs
 891
 892           return . mkLams (len : args)
 893                  . wrapFamInstBody arr_tc var_tys
 894                  . mkConApp arr_dc
 895                  $ ty_args ++ shape ++ pre ++ repr ++ post
 896
 897     mk_arr_repr len []
 898       = do
 899           units <- replicatePA len (Var unitDataConId)
 900           return [units]
 901
 902     mk_arr_repr len arrs = return arrs
 903
 904     def_worker data_con arg_tys mk_body
 905       = do
 906           body <- closedV
 907                 . inBind orig_worker
 908                 . polyAbstract tyvars $ \abstract ->
 909                   liftM (abstract . vectorised)
 910                 $ buildClosures tyvars [] arg_tys res_ty mk_body
 911
 912           vect_worker <- cloneId mkVectOcc orig_worker (exprType body)
 913           defGlobalVar orig_worker vect_worker
 914           return (vect_worker, body)
 915       where
 916         orig_worker = dataConWorkId data_con
 917
 918 buildPADict :: Repr -> TyCon -> TyCon -> TyCon -> Var -> VM CoreExpr
 919 buildPADict repr vect_tc prepr_tc arr_tc dfun
 920   = polyAbstract tvs $ \abstract ->
 921     do
 922       meth_binds <- mapM (mk_method repr) paMethods
 923       let meth_exprs = map (Var . fst) meth_binds
 924
 925       pa_dc <- builtin paDataCon
 926       let dict = mkConApp pa_dc (Type (mkTyConApp vect_tc arg_tys) : meth_exprs)
 927           body = Let (Rec meth_binds) dict
 928       return . mkInlineMe $ abstract body
 929   where
 930     tvs = tyConTyVars arr_tc
 931     arg_tys = mkTyVarTys tvs
 932
 933     mk_method repr (name, build)
 934       = localV
 935       $ do
 936           body <- build repr vect_tc prepr_tc arr_tc
 937           var  <- newLocalVar name (exprType body)
 938           return (var, mkInlineMe body)
 939
 940 paMethods = [(fsLit "toPRepr",      buildToPRepr),
 941              (fsLit "fromPRepr",    buildFromPRepr),
 942              (fsLit "toArrPRepr",   buildToArrPRepr),
 943              (fsLit "fromArrPRepr", buildFromArrPRepr),
 944              (fsLit "dictPRepr",    buildPRDict)]
 945
 946 -- | Split the given tycons into two sets depending on whether they have to be
 947 -- converted (first list) or not (second list). The first argument contains
 948 -- information about the conversion status of external tycons:
 949 --
 950 --   * tycons which have converted versions are mapped to True
 951 --   * tycons which are not changed by vectorisation are mapped to False
 952 --   * tycons which can't be converted are not elements of the map
 953 --
 954 classifyTyCons :: UniqFM Bool -> [TyConGroup] -> ([TyCon], [TyCon])
 955 classifyTyCons = classify [] []
 956   where
 957     classify conv keep cs [] = (conv, keep)
 958     classify conv keep cs ((tcs, ds) : rs)
 959       | can_convert && must_convert
 960         = classify (tcs ++ conv) keep (cs `addListToUFM` [(tc,True) | tc <- tcs]) rs
 961       | can_convert
 962         = classify conv (tcs ++ keep) (cs `addListToUFM` [(tc,False) | tc <- tcs]) rs
 963       | otherwise
 964         = classify conv keep cs rs
 965       where
 966         refs = ds `delListFromUniqSet` tcs
 967
 968         can_convert  = isNullUFM (refs `minusUFM` cs) && all convertable tcs
 969         must_convert = foldUFM (||) False (intersectUFM_C const cs refs)
 970
 971         convertable tc = isDataTyCon tc && all isVanillaDataCon (tyConDataCons tc)
 972
 973 -- | Compute mutually recursive groups of tycons in topological order
 974 --
 975 tyConGroups :: [TyCon] -> [TyConGroup]
 976 tyConGroups tcs = map mk_grp (stronglyConnComp edges)
 977   where
 978     edges = [((tc, ds), tc, uniqSetToList ds) | tc <- tcs
 979                                 , let ds = tyConsOfTyCon tc]
 980
 981     mk_grp (AcyclicSCC (tc, ds)) = ([tc], ds)
 982     mk_grp (CyclicSCC els)       = (tcs, unionManyUniqSets dss)
 983       where
 984         (tcs, dss) = unzip els
 985
 986 tyConsOfTyCon :: TyCon -> UniqSet TyCon
 987 tyConsOfTyCon
 988   = tyConsOfTypes . concatMap dataConRepArgTys . tyConDataCons
 989
 990 tyConsOfType :: Type -> UniqSet TyCon
 991 tyConsOfType ty
 992   | Just ty' <- coreView ty    = tyConsOfType ty'
 993 tyConsOfType (TyVarTy v)       = emptyUniqSet
 994 tyConsOfType (TyConApp tc tys) = extend (tyConsOfTypes tys)
 995   where
 996     extend | isUnLiftedTyCon tc
 997            || isTupleTyCon   tc = id
 998
 999            | otherwise          = (`addOneToUniqSet` tc)
1000
1001 tyConsOfType (AppTy a b)       = tyConsOfType a `unionUniqSets` tyConsOfType b
1002 tyConsOfType (FunTy a b)       = (tyConsOfType a `unionUniqSets` tyConsOfType b)
1003                                  `addOneToUniqSet` funTyCon
1004 tyConsOfType (ForAllTy _ ty)   = tyConsOfType ty
1005 tyConsOfType other             = pprPanic "ClosureConv.tyConsOfType" $ ppr other
1006
1007 tyConsOfTypes :: [Type] -> UniqSet TyCon
1008 tyConsOfTypes = unionManyUniqSets . map tyConsOfType
1009
1010
1011 -- ----------------------------------------------------------------------------
1012 -- Conversions
1013
1014 fromVect :: Type -> CoreExpr -> VM CoreExpr
1015 fromVect ty expr | Just ty' <- coreView ty = fromVect ty' expr
1016 fromVect (FunTy arg_ty res_ty) expr
1017   = do
1018       arg     <- newLocalVar (fsLit "x") arg_ty
1019       varg    <- toVect arg_ty (Var arg)
1020       varg_ty <- vectType arg_ty
1021       vres_ty <- vectType res_ty
1022       apply   <- builtin applyClosureVar
1023       body    <- fromVect res_ty
1024                $ Var apply `mkTyApps` [varg_ty, vres_ty] `mkApps` [expr, varg]
1025       return $ Lam arg body
1026 fromVect ty expr
1027   = identityConv ty >> return expr
1028
1029 toVect :: Type -> CoreExpr -> VM CoreExpr
1030 toVect ty expr = identityConv ty >> return expr
1031
1032 identityConv :: Type -> VM ()
1033 identityConv ty | Just ty' <- coreView ty = identityConv ty'
1034 identityConv (TyConApp tycon tys)
1035   = do
1036       mapM_ identityConv tys
1037       identityConvTyCon tycon
1038 identityConv ty = noV
1039
1040 identityConvTyCon :: TyCon -> VM ()
1041 identityConvTyCon tc
1042   | isBoxedTupleTyCon tc = return ()
1043   | isUnLiftedTyCon tc   = return ()
1044   | otherwise            = do
1045                              tc' <- maybeV (lookupTyCon tc)
1046                              if tc == tc' then return () else noV
1047