compiler/vectorise/VectType.hs

   1 {-# OPTIONS -fno-warn-missing-signatures #-}
   2
   3 module VectType ( vectTyCon, vectAndLiftType, vectType, vectTypeEnv,
   4                   -- arrSumArity, pdataCompTys, pdataCompVars,
   5                   buildPADict,
   6                   fromVect )
   7 where
   8
   9 import VectUtils
  10 import Vectorise.Env
  11 import Vectorise.Convert
  12 import Vectorise.Vect
  13 import Vectorise.Monad
  14 import Vectorise.Builtins
  15 import Vectorise.Type.Type
  16 import Vectorise.Type.TyConDecl
  17 import Vectorise.Type.Classify
  18 import Vectorise.Utils.Closure
  19 import Vectorise.Utils.Hoisting
  20
  21 import HscTypes          ( TypeEnv, extendTypeEnvList, typeEnvTyCons )
  22 import BasicTypes
  23 import CoreSyn
  24 import CoreUtils
  25 import CoreUnfold
  26 import MkCore            ( mkWildCase )
  27 import BuildTyCl
  28 import DataCon
  29 import TyCon
  30 import Type
  31 import Coercion
  32 import FamInstEnv        ( FamInst, mkLocalFamInst )
  33 import OccName
  34 import Id
  35 import MkId
  36 import Var
  37 import Name              ( Name, getOccName )
  38 import NameEnv
  39
  40 import Unique
  41 import UniqFM
  42 import Util
  43
  44 import Outputable
  45 import FastString
  46
  47 import MonadUtils     ( zipWith3M, foldrM, concatMapM )
  48 import Control.Monad  ( liftM, liftM2, zipWithM, zipWithM_, mapAndUnzipM )
  49 import Data.List
  50
  51 debug           = False
  52 dtrace s x      = if debug then pprTrace "VectType" s x else x
  53
  54
  55 -- | Vectorise a type environment.
  56 --   The type environment contains all the type things defined in a module.
  57 vectTypeEnv
  58         :: TypeEnv
  59         -> VM ( TypeEnv                 -- Vectorised type environment.
  60               , [FamInst]               -- New type family instances.
  61               , [(Var, CoreExpr)])      -- New top level bindings.
  62
  63 vectTypeEnv env
  64  = dtrace (ppr env)
  65  $ do
  66       cs <- readGEnv $ mk_map . global_tycons
  67
  68       -- Split the list of TyCons into the ones we have to vectorise vs the
  69       -- ones we can pass through unchanged. We also pass through algebraic
  70       -- types that use non Haskell98 features, as we don't handle those.
  71       let (conv_tcs, keep_tcs) = classifyTyCons cs groups
  72           keep_dcs             = concatMap tyConDataCons keep_tcs
  73
  74       zipWithM_ defTyCon   keep_tcs keep_tcs
  75       zipWithM_ defDataCon keep_dcs keep_dcs
  76
  77       new_tcs <- vectTyConDecls conv_tcs
  78
  79       let orig_tcs = keep_tcs ++ conv_tcs
  80
  81       -- We don't need to make new representation types for dictionary
  82       -- constructors. The constructors are always fully applied, and we don't
  83       -- need to lift them to arrays as a dictionary of a particular type
  84       -- always has the same value.
  85       let vect_tcs = filter (not . isClassTyCon)
  86                    $ keep_tcs ++ new_tcs
  87
  88       (_, binds, inst_tcs) <- fixV $ \ ~(dfuns', _, _) ->
  89         do
  90           defTyConPAs (zipLazy vect_tcs dfuns')
  91           reprs     <- mapM tyConRepr vect_tcs
  92           repr_tcs  <- zipWith3M buildPReprTyCon orig_tcs vect_tcs reprs
  93           pdata_tcs <- zipWith3M buildPDataTyCon orig_tcs vect_tcs reprs
  94
  95           dfuns     <- sequence
  96                     $  zipWith5 buildTyConBindings
  97                                orig_tcs
  98                                vect_tcs
  99                                repr_tcs
 100                                pdata_tcs
 101                                reprs
 102
 103           binds     <- takeHoisted
 104           return (dfuns, binds, repr_tcs ++ pdata_tcs)
 105
 106       let all_new_tcs = new_tcs ++ inst_tcs
 107
 108       let new_env = extendTypeEnvList env
 109                        (map ATyCon all_new_tcs
 110                         ++ [ADataCon dc | tc <- all_new_tcs
 111                                         , dc <- tyConDataCons tc])
 112
 113       return (new_env, map mkLocalFamInst inst_tcs, binds)
 114   where
 115     tycons = typeEnvTyCons env
 116     groups = tyConGroups tycons
 117
 118     mk_map env = listToUFM_Directly [(u, getUnique n /= u) | (u,n) <- nameEnvUniqueElts env]
 119
 120
 121 mk_fam_inst :: TyCon -> TyCon -> (TyCon, [Type])
 122 mk_fam_inst fam_tc arg_tc
 123   = (fam_tc, [mkTyConApp arg_tc . mkTyVarTys $ tyConTyVars arg_tc])
 124
 125
 126 buildPReprTyCon :: TyCon -> TyCon -> SumRepr -> VM TyCon
 127 buildPReprTyCon orig_tc vect_tc repr
 128   = do
 129       name     <- cloneName mkPReprTyConOcc (tyConName orig_tc)
 130       -- rhs_ty   <- buildPReprType vect_tc
 131       rhs_ty   <- sumReprType repr
 132       prepr_tc <- builtin preprTyCon
 133       liftDs $ buildSynTyCon name
 134                              tyvars
 135                              (SynonymTyCon rhs_ty)
 136                              (typeKind rhs_ty)
 137                              (Just $ mk_fam_inst prepr_tc vect_tc)
 138   where
 139     tyvars = tyConTyVars vect_tc
 140
 141 data CompRepr = Keep Type
 142                      CoreExpr     -- PR dictionary for the type
 143               | Wrap Type
 144
 145 data ProdRepr = EmptyProd
 146               | UnaryProd CompRepr
 147               | Prod { repr_tup_tc   :: TyCon  -- representation tuple tycon
 148                      , repr_ptup_tc  :: TyCon  -- PData representation tycon
 149                      , repr_comp_tys :: [Type] -- representation types of
 150                      , repr_comps    :: [CompRepr]          -- components
 151                      }
 152 data ConRepr  = ConRepr DataCon ProdRepr
 153
 154 data SumRepr  = EmptySum
 155               | UnarySum ConRepr
 156               | Sum  { repr_sum_tc   :: TyCon  -- representation sum tycon
 157                      , repr_psum_tc  :: TyCon  -- PData representation tycon
 158                      , repr_sel_ty   :: Type   -- type of selector
 159                      , repr_con_tys :: [Type]  -- representation types of
 160                      , repr_cons     :: [ConRepr]           -- components
 161                      }
 162
 163 tyConRepr :: TyCon -> VM SumRepr
 164 tyConRepr tc = sum_repr (tyConDataCons tc)
 165   where
 166     sum_repr []    = return EmptySum
 167     sum_repr [con] = liftM UnarySum (con_repr con)
 168     sum_repr cons  = do
 169                        rs     <- mapM con_repr cons
 170                        sum_tc <- builtin (sumTyCon arity)
 171                        tys    <- mapM conReprType rs
 172                        (psum_tc, _) <- pdataReprTyCon (mkTyConApp sum_tc tys)
 173                        sel_ty <- builtin (selTy arity)
 174                        return $ Sum { repr_sum_tc  = sum_tc
 175                                     , repr_psum_tc = psum_tc
 176                                     , repr_sel_ty  = sel_ty
 177                                     , repr_con_tys = tys
 178                                     , repr_cons    = rs
 179                                     }
 180       where
 181         arity = length cons
 182
 183     con_repr con = liftM (ConRepr con) (prod_repr (dataConRepArgTys con))
 184
 185     prod_repr []   = return EmptyProd
 186     prod_repr [ty] = liftM UnaryProd (comp_repr ty)
 187     prod_repr tys  = do
 188                        rs <- mapM comp_repr tys
 189                        tup_tc <- builtin (prodTyCon arity)
 190                        tys'    <- mapM compReprType rs
 191                        (ptup_tc, _) <- pdataReprTyCon (mkTyConApp tup_tc tys')
 192                        return $ Prod { repr_tup_tc   = tup_tc
 193                                      , repr_ptup_tc  = ptup_tc
 194                                      , repr_comp_tys = tys'
 195                                      , repr_comps    = rs
 196                                      }
 197       where
 198         arity = length tys
 199
 200     comp_repr ty = liftM (Keep ty) (prDictOfType ty)
 201                    `orElseV` return (Wrap ty)
 202
 203 sumReprType :: SumRepr -> VM Type
 204 sumReprType EmptySum = voidType
 205 sumReprType (UnarySum r) = conReprType r
 206 sumReprType (Sum { repr_sum_tc  = sum_tc, repr_con_tys = tys })
 207   = return $ mkTyConApp sum_tc tys
 208
 209 conReprType :: ConRepr -> VM Type
 210 conReprType (ConRepr _ r) = prodReprType r
 211
 212 prodReprType :: ProdRepr -> VM Type
 213 prodReprType EmptyProd = voidType
 214 prodReprType (UnaryProd r) = compReprType r
 215 prodReprType (Prod { repr_tup_tc = tup_tc, repr_comp_tys = tys })
 216   = return $ mkTyConApp tup_tc tys
 217
 218 compReprType :: CompRepr -> VM Type
 219 compReprType (Keep ty _) = return ty
 220 compReprType (Wrap ty) = do
 221                              wrap_tc <- builtin wrapTyCon
 222                              return $ mkTyConApp wrap_tc [ty]
 223
 224 compOrigType :: CompRepr -> Type
 225 compOrigType (Keep ty _) = ty
 226 compOrigType (Wrap ty) = ty
 227
 228 buildToPRepr :: TyCon -> TyCon -> TyCon -> SumRepr -> VM CoreExpr
 229 buildToPRepr vect_tc repr_tc _ repr
 230   = do
 231       let arg_ty = mkTyConApp vect_tc ty_args
 232       res_ty <- mkPReprType arg_ty
 233       arg    <- newLocalVar (fsLit "x") arg_ty
 234       result <- to_sum (Var arg) arg_ty res_ty repr
 235       return $ Lam arg result
 236   where
 237     ty_args = mkTyVarTys (tyConTyVars vect_tc)
 238
 239     wrap_repr_inst = wrapFamInstBody repr_tc ty_args
 240
 241     to_sum _ _ _ EmptySum
 242       = do
 243           void <- builtin voidVar
 244           return $ wrap_repr_inst $ Var void
 245
 246     to_sum arg arg_ty res_ty (UnarySum r)
 247       = do
 248           (pat, vars, body) <- con_alt r
 249           return $ mkWildCase arg arg_ty res_ty
 250                    [(pat, vars, wrap_repr_inst body)]
 251
 252     to_sum arg arg_ty res_ty (Sum { repr_sum_tc  = sum_tc
 253                                   , repr_con_tys = tys
 254                                   , repr_cons    =  cons })
 255       = do
 256           alts <- mapM con_alt cons
 257           let alts' = [(pat, vars, wrap_repr_inst
 258                                    $ mkConApp sum_con (map Type tys ++ [body]))
 259                         | ((pat, vars, body), sum_con)
 260                             <- zip alts (tyConDataCons sum_tc)]
 261           return $ mkWildCase arg arg_ty res_ty alts'
 262
 263     con_alt (ConRepr con r)
 264       = do
 265           (vars, body) <- to_prod r
 266           return (DataAlt con, vars, body)
 267
 268     to_prod EmptyProd
 269       = do
 270           void <- builtin voidVar
 271           return ([], Var void)
 272
 273     to_prod (UnaryProd comp)
 274       = do
 275           var  <- newLocalVar (fsLit "x") (compOrigType comp)
 276           body <- to_comp (Var var) comp
 277           return ([var], body)
 278
 279     to_prod(Prod { repr_tup_tc   = tup_tc
 280                  , repr_comp_tys = tys
 281                  , repr_comps    = comps })
 282       = do
 283           vars  <- newLocalVars (fsLit "x") (map compOrigType comps)
 284           exprs <- zipWithM to_comp (map Var vars) comps
 285           return (vars, mkConApp tup_con (map Type tys ++ exprs))
 286       where
 287         [tup_con] = tyConDataCons tup_tc
 288
 289     to_comp expr (Keep _ _) = return expr
 290     to_comp expr (Wrap ty)  = do
 291                                 wrap_tc <- builtin wrapTyCon
 292                                 return $ wrapNewTypeBody wrap_tc [ty] expr
 293
 294
 295 buildFromPRepr :: TyCon -> TyCon -> TyCon -> SumRepr -> VM CoreExpr
 296 buildFromPRepr vect_tc repr_tc _ repr
 297   = do
 298       arg_ty <- mkPReprType res_ty
 299       arg <- newLocalVar (fsLit "x") arg_ty
 300
 301       result <- from_sum (unwrapFamInstScrut repr_tc ty_args (Var arg))
 302                          repr
 303       return $ Lam arg result
 304   where
 305     ty_args = mkTyVarTys (tyConTyVars vect_tc)
 306     res_ty  = mkTyConApp vect_tc ty_args
 307
 308     from_sum _ EmptySum
 309       = do
 310           dummy <- builtin fromVoidVar
 311           return $ Var dummy `App` Type res_ty
 312
 313     from_sum expr (UnarySum r) = from_con expr r
 314     from_sum expr (Sum { repr_sum_tc  = sum_tc
 315                        , repr_con_tys = tys
 316                        , repr_cons    = cons })
 317       = do
 318           vars  <- newLocalVars (fsLit "x") tys
 319           es    <- zipWithM from_con (map Var vars) cons
 320           return $ mkWildCase expr (exprType expr) res_ty
 321                    [(DataAlt con, [var], e)
 322                       | (con, var, e) <- zip3 (tyConDataCons sum_tc) vars es]
 323
 324     from_con expr (ConRepr con r)
 325       = from_prod expr (mkConApp con $ map Type ty_args) r
 326
 327     from_prod _ con EmptyProd = return con
 328     from_prod expr con (UnaryProd r)
 329       = do
 330           e <- from_comp expr r
 331           return $ con `App` e
 332
 333     from_prod expr con (Prod { repr_tup_tc   = tup_tc
 334                              , repr_comp_tys = tys
 335                              , repr_comps    = comps
 336                              })
 337       = do
 338           vars <- newLocalVars (fsLit "y") tys
 339           es   <- zipWithM from_comp (map Var vars) comps
 340           return $ mkWildCase expr (exprType expr) res_ty
 341                    [(DataAlt tup_con, vars, con `mkApps` es)]
 342       where
 343         [tup_con] = tyConDataCons tup_tc
 344
 345     from_comp expr (Keep _ _) = return expr
 346     from_comp expr (Wrap ty)
 347       = do
 348           wrap <- builtin wrapTyCon
 349           return $ unwrapNewTypeBody wrap [ty] expr
 350
 351
 352 buildToArrPRepr :: TyCon -> TyCon -> TyCon -> SumRepr -> VM CoreExpr
 353 buildToArrPRepr vect_tc prepr_tc pdata_tc r
 354   = do
 355       arg_ty <- mkPDataType el_ty
 356       res_ty <- mkPDataType =<< mkPReprType el_ty
 357       arg    <- newLocalVar (fsLit "xs") arg_ty
 358
 359       pdata_co <- mkBuiltinCo pdataTyCon
 360       let Just repr_co = tyConFamilyCoercion_maybe prepr_tc
 361           co           = mkAppCoercion pdata_co
 362                        . mkSymCoercion
 363                        $ mkTyConApp repr_co ty_args
 364
 365           scrut   = unwrapFamInstScrut pdata_tc ty_args (Var arg)
 366
 367       (vars, result) <- to_sum r
 368
 369       return . Lam arg
 370              $ mkWildCase scrut (mkTyConApp pdata_tc ty_args) res_ty
 371                [(DataAlt pdata_dc, vars, mkCoerce co result)]
 372   where
 373     ty_args = mkTyVarTys $ tyConTyVars vect_tc
 374     el_ty   = mkTyConApp vect_tc ty_args
 375
 376     [pdata_dc] = tyConDataCons pdata_tc
 377
 378
 379     to_sum EmptySum = do
 380                         pvoid <- builtin pvoidVar
 381                         return ([], Var pvoid)
 382     to_sum (UnarySum r) = to_con r
 383     to_sum (Sum { repr_psum_tc = psum_tc
 384                 , repr_sel_ty  = sel_ty
 385                 , repr_con_tys = tys
 386                 , repr_cons    = cons
 387                 })
 388       = do
 389           (vars, exprs) <- mapAndUnzipM to_con cons
 390           sel <- newLocalVar (fsLit "sel") sel_ty
 391           return (sel : concat vars, mk_result (Var sel) exprs)
 392       where
 393         [psum_con] = tyConDataCons psum_tc
 394         mk_result sel exprs = wrapFamInstBody psum_tc tys
 395                             $ mkConApp psum_con
 396                             $ map Type tys ++ (sel : exprs)
 397
 398     to_con (ConRepr _ r) = to_prod r
 399
 400     to_prod EmptyProd = do
 401                           pvoid <- builtin pvoidVar
 402                           return ([], Var pvoid)
 403     to_prod (UnaryProd r)
 404       = do
 405           pty  <- mkPDataType (compOrigType r)
 406           var  <- newLocalVar (fsLit "x") pty
 407           expr <- to_comp (Var var) r
 408           return ([var], expr)
 409
 410     to_prod (Prod { repr_ptup_tc  = ptup_tc
 411                   , repr_comp_tys = tys
 412                   , repr_comps    = comps })
 413       = do
 414           ptys <- mapM (mkPDataType . compOrigType) comps
 415           vars <- newLocalVars (fsLit "x") ptys
 416           es   <- zipWithM to_comp (map Var vars) comps
 417           return (vars, mk_result es)
 418       where
 419         [ptup_con] = tyConDataCons ptup_tc
 420         mk_result exprs = wrapFamInstBody ptup_tc tys
 421                         $ mkConApp ptup_con
 422                         $ map Type tys ++ exprs
 423
 424     to_comp expr (Keep _ _) = return expr
 425
 426     -- FIXME: this is bound to be wrong!
 427     to_comp expr (Wrap ty)
 428       = do
 429           wrap_tc  <- builtin wrapTyCon
 430           (pwrap_tc, _) <- pdataReprTyCon (mkTyConApp wrap_tc [ty])
 431           return $ wrapNewTypeBody pwrap_tc [ty] expr
 432
 433
 434 buildFromArrPRepr :: TyCon -> TyCon -> TyCon -> SumRepr -> VM CoreExpr
 435 buildFromArrPRepr vect_tc prepr_tc pdata_tc r
 436   = do
 437       arg_ty <- mkPDataType =<< mkPReprType el_ty
 438       res_ty <- mkPDataType el_ty
 439       arg    <- newLocalVar (fsLit "xs") arg_ty
 440
 441       pdata_co <- mkBuiltinCo pdataTyCon
 442       let Just repr_co = tyConFamilyCoercion_maybe prepr_tc
 443           co           = mkAppCoercion pdata_co
 444                        $ mkTyConApp repr_co var_tys
 445
 446           scrut  = mkCoerce co (Var arg)
 447
 448           mk_result args = wrapFamInstBody pdata_tc var_tys
 449                          $ mkConApp pdata_con
 450                          $ map Type var_tys ++ args
 451
 452       (expr, _) <- fixV $ \ ~(_, args) ->
 453                      from_sum res_ty (mk_result args) scrut r
 454
 455       return $ Lam arg expr
 456
 457       -- (args, mk) <- from_sum res_ty scrut r
 458
 459       -- let result = wrapFamInstBody pdata_tc var_tys
 460       --           . mkConApp pdata_dc
 461       --           $ map Type var_tys ++ args
 462
 463       -- return $ Lam arg (mk result)
 464   where
 465     var_tys = mkTyVarTys $ tyConTyVars vect_tc
 466     el_ty   = mkTyConApp vect_tc var_tys
 467
 468     [pdata_con] = tyConDataCons pdata_tc
 469
 470     from_sum _ res _ EmptySum = return (res, [])
 471     from_sum res_ty res expr (UnarySum r) = from_con res_ty res expr r
 472     from_sum res_ty res expr (Sum { repr_psum_tc = psum_tc
 473                                   , repr_sel_ty  = sel_ty
 474                                   , repr_con_tys = tys
 475                                   , repr_cons    = cons })
 476       = do
 477           sel  <- newLocalVar (fsLit "sel") sel_ty
 478           ptys <- mapM mkPDataType tys
 479           vars <- newLocalVars (fsLit "xs") ptys
 480           (res', args) <- fold from_con res_ty res (map Var vars) cons
 481           let scrut = unwrapFamInstScrut psum_tc tys expr
 482               body  = mkWildCase scrut (exprType scrut) res_ty
 483                       [(DataAlt psum_con, sel : vars, res')]
 484           return (body, Var sel : args)
 485       where
 486         [psum_con] = tyConDataCons psum_tc
 487
 488
 489     from_con res_ty res expr (ConRepr _ r) = from_prod res_ty res expr r
 490
 491     from_prod _ res _ EmptyProd = return (res, [])
 492     from_prod res_ty res expr (UnaryProd r)
 493       = from_comp res_ty res expr r
 494     from_prod res_ty res expr (Prod { repr_ptup_tc  = ptup_tc
 495                                     , repr_comp_tys = tys
 496                                     , repr_comps    = comps })
 497       = do
 498           ptys <- mapM mkPDataType tys
 499           vars <- newLocalVars (fsLit "ys") ptys
 500           (res', args) <- fold from_comp res_ty res (map Var vars) comps
 501           let scrut = unwrapFamInstScrut ptup_tc tys expr
 502               body  = mkWildCase scrut (exprType scrut) res_ty
 503                       [(DataAlt ptup_con, vars, res')]
 504           return (body, args)
 505       where
 506         [ptup_con] = tyConDataCons ptup_tc
 507
 508     from_comp _ res expr (Keep _ _) = return (res, [expr])
 509     from_comp _ res expr (Wrap ty)
 510       = do
 511           wrap_tc  <- builtin wrapTyCon
 512           (pwrap_tc, _) <- pdataReprTyCon (mkTyConApp wrap_tc [ty])
 513           return (res, [unwrapNewTypeBody pwrap_tc [ty]
 514                         $ unwrapFamInstScrut pwrap_tc [ty] expr])
 515
 516     fold f res_ty res exprs rs = foldrM f' (res, []) (zip exprs rs)
 517       where
 518         f' (expr, r) (res, args) = do
 519                                      (res', args') <- f res_ty res expr r
 520                                      return (res', args' ++ args)
 521
 522 buildPRDict :: TyCon -> TyCon -> TyCon -> SumRepr -> VM CoreExpr
 523 buildPRDict vect_tc prepr_tc _ r
 524   = do
 525       dict <- sum_dict r
 526       pr_co <- mkBuiltinCo prTyCon
 527       let co = mkAppCoercion pr_co
 528              . mkSymCoercion
 529              $ mkTyConApp arg_co ty_args
 530       return (mkCoerce co dict)
 531   where
 532     ty_args = mkTyVarTys (tyConTyVars vect_tc)
 533     Just arg_co = tyConFamilyCoercion_maybe prepr_tc
 534
 535     sum_dict EmptySum = prDFunOfTyCon =<< builtin voidTyCon
 536     sum_dict (UnarySum r) = con_dict r
 537     sum_dict (Sum { repr_sum_tc  = sum_tc
 538                   , repr_con_tys = tys
 539                   , repr_cons    = cons
 540                   })
 541       = do
 542           dicts <- mapM con_dict cons
 543           dfun  <- prDFunOfTyCon sum_tc
 544           return $ dfun `mkTyApps` tys `mkApps` dicts
 545
 546     con_dict (ConRepr _ r) = prod_dict r
 547
 548     prod_dict EmptyProd = prDFunOfTyCon =<< builtin voidTyCon
 549     prod_dict (UnaryProd r) = comp_dict r
 550     prod_dict (Prod { repr_tup_tc   = tup_tc
 551                     , repr_comp_tys = tys
 552                     , repr_comps    = comps })
 553       = do
 554           dicts <- mapM comp_dict comps
 555           dfun <- prDFunOfTyCon tup_tc
 556           return $ dfun `mkTyApps` tys `mkApps` dicts
 557
 558     comp_dict (Keep _ pr) = return pr
 559     comp_dict (Wrap ty)   = wrapPR ty
 560
 561
 562 buildPDataTyCon :: TyCon -> TyCon -> SumRepr -> VM TyCon
 563 buildPDataTyCon orig_tc vect_tc repr = fixV $ \repr_tc ->
 564   do
 565     name' <- cloneName mkPDataTyConOcc orig_name
 566     rhs   <- buildPDataTyConRhs orig_name vect_tc repr_tc repr
 567     pdata <- builtin pdataTyCon
 568
 569     liftDs $ buildAlgTyCon name'
 570                            tyvars
 571                            []          -- no stupid theta
 572                            rhs
 573                            rec_flag    -- FIXME: is this ok?
 574                            False       -- FIXME: no generics
 575                            False       -- not GADT syntax
 576                            (Just $ mk_fam_inst pdata vect_tc)
 577   where
 578     orig_name = tyConName orig_tc
 579     tyvars = tyConTyVars vect_tc
 580     rec_flag = boolToRecFlag (isRecursiveTyCon vect_tc)
 581
 582
 583 buildPDataTyConRhs :: Name -> TyCon -> TyCon -> SumRepr -> VM AlgTyConRhs
 584 buildPDataTyConRhs orig_name vect_tc repr_tc repr
 585   = do
 586       data_con <- buildPDataDataCon orig_name vect_tc repr_tc repr
 587       return $ DataTyCon { data_cons = [data_con], is_enum = False }
 588
 589 buildPDataDataCon :: Name -> TyCon -> TyCon -> SumRepr -> VM DataCon
 590 buildPDataDataCon orig_name vect_tc repr_tc repr
 591   = do
 592       dc_name  <- cloneName mkPDataDataConOcc orig_name
 593       comp_tys <- sum_tys repr
 594
 595       liftDs $ buildDataCon dc_name
 596                             False                  -- not infix
 597                             (map (const HsNoBang) comp_tys)
 598                             []                     -- no field labels
 599                             tvs
 600                             []                     -- no existentials
 601                             []                     -- no eq spec
 602                             []                     -- no context
 603                             comp_tys
 604                             (mkFamilyTyConApp repr_tc (mkTyVarTys tvs))
 605                             repr_tc
 606   where
 607     tvs   = tyConTyVars vect_tc
 608
 609     sum_tys EmptySum = return []
 610     sum_tys (UnarySum r) = con_tys r
 611     sum_tys (Sum { repr_sel_ty = sel_ty
 612                  , repr_cons   = cons })
 613       = liftM (sel_ty :) (concatMapM con_tys cons)
 614
 615     con_tys (ConRepr _ r) = prod_tys r
 616
 617     prod_tys EmptyProd = return []
 618     prod_tys (UnaryProd r) = liftM singleton (comp_ty r)
 619     prod_tys (Prod { repr_comps = comps }) = mapM comp_ty comps
 620
 621     comp_ty r = mkPDataType (compOrigType r)
 622
 623
 624 buildTyConBindings :: TyCon -> TyCon -> TyCon -> TyCon -> SumRepr
 625                    -> VM Var
 626 buildTyConBindings orig_tc vect_tc prepr_tc pdata_tc repr
 627   = do
 628       vectDataConWorkers orig_tc vect_tc pdata_tc
 629       buildPADict vect_tc prepr_tc pdata_tc repr
 630
 631 vectDataConWorkers :: TyCon -> TyCon -> TyCon -> VM ()
 632 vectDataConWorkers orig_tc vect_tc arr_tc
 633   = do
 634       bs <- sequence
 635           . zipWith3 def_worker  (tyConDataCons orig_tc) rep_tys
 636           $ zipWith4 mk_data_con (tyConDataCons vect_tc)
 637                                  rep_tys
 638                                  (inits rep_tys)
 639                                  (tail $ tails rep_tys)
 640       mapM_ (uncurry hoistBinding) bs
 641   where
 642     tyvars   = tyConTyVars vect_tc
 643     var_tys  = mkTyVarTys tyvars
 644     ty_args  = map Type var_tys
 645     res_ty   = mkTyConApp vect_tc var_tys
 646
 647     cons     = tyConDataCons vect_tc
 648     arity    = length cons
 649     [arr_dc] = tyConDataCons arr_tc
 650
 651     rep_tys  = map dataConRepArgTys $ tyConDataCons vect_tc
 652
 653
 654     mk_data_con con tys pre post
 655       = liftM2 (,) (vect_data_con con)
 656                    (lift_data_con tys pre post (mkDataConTag con))
 657
 658     sel_replicate len tag
 659       | arity > 1 = do
 660                       rep <- builtin (selReplicate arity)
 661                       return [rep `mkApps` [len, tag]]
 662
 663       | otherwise = return []
 664
 665     vect_data_con con = return $ mkConApp con ty_args
 666     lift_data_con tys pre_tys post_tys tag
 667       = do
 668           len  <- builtin liftingContext
 669           args <- mapM (newLocalVar (fsLit "xs"))
 670                   =<< mapM mkPDataType tys
 671
 672           sel  <- sel_replicate (Var len) tag
 673
 674           pre   <- mapM emptyPD (concat pre_tys)
 675           post  <- mapM emptyPD (concat post_tys)
 676
 677           return . mkLams (len : args)
 678                  . wrapFamInstBody arr_tc var_tys
 679                  . mkConApp arr_dc
 680                  $ ty_args ++ sel ++ pre ++ map Var args ++ post
 681
 682     def_worker data_con arg_tys mk_body
 683       = do
 684           arity <- polyArity tyvars
 685           body <- closedV
 686                 . inBind orig_worker
 687                 . polyAbstract tyvars $ \args ->
 688                   liftM (mkLams (tyvars ++ args) . vectorised)
 689                 $ buildClosures tyvars [] arg_tys res_ty mk_body
 690
 691           raw_worker <- cloneId mkVectOcc orig_worker (exprType body)
 692           let vect_worker = raw_worker `setIdUnfolding`
 693                               mkInlineRule body (Just arity)
 694           defGlobalVar orig_worker vect_worker
 695           return (vect_worker, body)
 696       where
 697         orig_worker = dataConWorkId data_con
 698
 699 buildPADict :: TyCon -> TyCon -> TyCon -> SumRepr -> VM Var
 700 buildPADict vect_tc prepr_tc arr_tc repr
 701   = polyAbstract tvs $ \args ->
 702     do
 703       method_ids <- mapM (method args) paMethods
 704
 705       pa_tc  <- builtin paTyCon
 706       pa_dc  <- builtin paDataCon
 707       let dict = mkLams (tvs ++ args)
 708                $ mkConApp pa_dc
 709                $ Type inst_ty : map (method_call args) method_ids
 710
 711           dfun_ty = mkForAllTys tvs
 712                   $ mkFunTys (map varType args) (mkTyConApp pa_tc [inst_ty])
 713
 714       raw_dfun <- newExportedVar dfun_name dfun_ty
 715       let dfun = raw_dfun `setIdUnfolding` mkDFunUnfolding dfun_ty (map Var method_ids)
 716                           `setInlinePragma` dfunInlinePragma
 717
 718       hoistBinding dfun dict
 719       return dfun
 720   where
 721     tvs = tyConTyVars vect_tc
 722     arg_tys = mkTyVarTys tvs
 723     inst_ty = mkTyConApp vect_tc arg_tys
 724
 725     dfun_name = mkPADFunOcc (getOccName vect_tc)
 726
 727     method args (name, build)
 728       = localV
 729       $ do
 730           expr <- build vect_tc prepr_tc arr_tc repr
 731           let body = mkLams (tvs ++ args) expr
 732           raw_var <- newExportedVar (method_name name) (exprType body)
 733           let var = raw_var
 734                       `setIdUnfolding` mkInlineRule body (Just (length args))
 735                       `setInlinePragma` alwaysInlinePragma
 736           hoistBinding var body
 737           return var
 738
 739     method_call args id = mkApps (Var id) (map Type arg_tys ++ map Var args)
 740
 741     method_name name = mkVarOcc $ occNameString dfun_name ++ ('$' : name)
 742
 743
 744 paMethods :: [(String, TyCon -> TyCon -> TyCon -> SumRepr -> VM CoreExpr)]
 745 paMethods = [("dictPRepr",    buildPRDict),
 746              ("toPRepr",      buildToPRepr),
 747              ("fromPRepr",    buildFromPRepr),
 748              ("toArrPRepr",   buildToArrPRepr),
 749              ("fromArrPRepr", buildFromArrPRepr)]
 750
 751