Vectorise.hs 17.5 KB
Newer Older
1

2 3 4
module Vectorise( vectorise )
where

5
import VectMonad
6
import VectUtils
7
import VectType
8
import VectCore
9

10
import HscTypes hiding      ( MonadThings(..) )
11

12
import Module               ( PackageId )
13
import CoreSyn
14
import CoreUtils
15
import CoreUnfold           ( mkInlineRule )
16
import MkCore               ( mkWildCase )
17
import CoreFVs
Ian Lynagh's avatar
Ian Lynagh committed
18
import CoreMonad            ( CoreM, getHscEnv )
19
import DataCon
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
20
import TyCon
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
21
import Type
22
import FamInstEnv           ( extendFamInstEnvList )
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
23 24
import Var
import VarEnv
25
import VarSet
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
26
import Id
27
import OccName
28
import BasicTypes           ( isLoopBreaker )
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
29

30
import Literal              ( Literal, mkMachInt )
31
import TysWiredIn
32
import TysPrim              ( intPrimTy )
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
33

34
import Outputable
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
35
import FastString
36 37
import Util                 ( zipLazy )
import Control.Monad
38
import Data.List            ( sortBy, unzip4 )
39

40 41 42 43 44 45 46
vectorise :: PackageId -> ModGuts -> CoreM ModGuts
vectorise backend guts = do
    hsc_env <- getHscEnv
    liftIO $ vectoriseIO backend hsc_env guts

vectoriseIO :: PackageId -> HscEnv -> ModGuts -> IO ModGuts
vectoriseIO backend hsc_env guts
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
47 48 49
  = do
      eps <- hscEPS hsc_env
      let info = hptVectInfo hsc_env `plusVectInfo` eps_vect_info eps
50
      Just (info', guts') <- initV backend hsc_env guts info (vectModule guts)
51
      return (guts' { mg_vect_info = info' })
52

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
53
vectModule :: ModGuts -> VM ModGuts
54 55
vectModule guts
  = do
56
      (types', fam_insts, tc_binds) <- vectTypeEnv (mg_types guts)
Ian Lynagh's avatar
Ian Lynagh committed
57

58 59
      let fam_inst_env' = extendFamInstEnvList (mg_fam_inst_env guts) fam_insts
      updGEnv (setFamInstEnv fam_inst_env')
Ian Lynagh's avatar
Ian Lynagh committed
60

61 62
      -- dicts   <- mapM buildPADict pa_insts
      -- workers <- mapM vectDataConWorkers pa_insts
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
63
      binds'  <- mapM vectTopBind (mg_binds guts)
64
      return $ guts { mg_types        = types'
65
                    , mg_binds        = Rec tc_binds : binds'
66 67 68
                    , mg_fam_inst_env = fam_inst_env'
                    , mg_fam_insts    = mg_fam_insts guts ++ fam_insts
                    }
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
69

70
vectTopBind :: CoreBind -> VM CoreBind
71 72
vectTopBind b@(NonRec var expr)
  = do
73 74
      (inline, expr') <- vectTopRhs var expr
      var' <- vectTopBinder var inline expr'
75
      hs    <- takeHoisted
76 77
      cexpr <- tryConvert var var' expr
      return . Rec $ (var, cexpr) : (var', expr') : hs
78 79 80 81 82
  `orElseV`
    return b

vectTopBind b@(Rec bs)
  = do
83 84 85 86 87 88 89
      (vars', _, exprs') <- fixV $ \ ~(_, inlines, rhss) ->
        do
          vars' <- sequence [vectTopBinder var inline rhs
                               | (var, ~(inline, rhs))
                                 <- zipLazy vars (zip inlines rhss)]
          (inlines', exprs') <- mapAndUnzipM (uncurry vectTopRhs) bs
          return (vars', inlines', exprs')
90
      hs     <- takeHoisted
91 92
      cexprs <- sequence $ zipWith3 tryConvert vars vars' exprs
      return . Rec $ zip vars cexprs ++ zip vars' exprs' ++ hs
93 94 95 96 97
  `orElseV`
    return b
  where
    (vars, exprs) = unzip bs

98 99 100 101
-- NOTE: vectTopBinder *MUST* be lazy in inline and expr because of how it is
-- used inside of fixV in vectTopBind
vectTopBinder :: Var -> Inline -> CoreExpr -> VM Var
vectTopBinder var inline expr
102
  = do
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
103
      vty  <- vectType (idType var)
104
      var' <- liftM (`setIdUnfolding` unfolding) $ cloneId mkVectOcc var vty
105 106
      defGlobalVar var var'
      return var'
107 108
  where
    unfolding = case inline of
109
                  Inline arity -> mkInlineRule expr (Just arity)
110
                  DontInline   -> noUnfolding
Ian Lynagh's avatar
Ian Lynagh committed
111

112
vectTopRhs :: Var -> CoreExpr -> VM (Inline, CoreExpr)
113
vectTopRhs var expr
114 115 116 117 118 119
  = closedV
  $ do
      (inline, vexpr) <- inBind var
                       $ vectPolyExpr (isLoopBreaker $ idOccInfo var)
                                      (freeVars expr)
      return (inline, vectorised vexpr)
120

121 122 123 124
tryConvert :: Var -> Var -> CoreExpr -> VM CoreExpr
tryConvert var vect_var rhs
  = fromVect (idType var) (Var vect_var) `orElseV` return rhs

125 126
-- ----------------------------------------------------------------------------
-- Bindings
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
127

128
vectBndr :: Var -> VM VVar
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
129 130
vectBndr v
  = do
131
      (vty, lty) <- vectAndLiftType (idType v)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
132 133 134 135 136
      let vv = v `Id.setIdType` vty
          lv = v `Id.setIdType` lty
      updLEnv (mapTo vv lv)
      return (vv, lv)
  where
137
    mapTo vv lv env = env { local_vars = extendVarEnv (local_vars env) v (vv, lv) }
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
138

139 140 141 142 143 144 145 146 147 148
vectBndrNew :: Var -> FastString -> VM VVar
vectBndrNew v fs
  = do
      vty <- vectType (idType v)
      vv  <- newLocalVVar fs vty
      updLEnv (upd vv)
      return vv
  where
    upd vv env = env { local_vars = extendVarEnv (local_vars env) v vv }

149
vectBndrIn :: Var -> VM a -> VM (VVar, a)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
150 151 152
vectBndrIn v p
  = localV
  $ do
153
      vv <- vectBndr v
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
154
      x <- p
155
      return (vv, x)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
156

157 158 159 160 161 162 163 164
vectBndrNewIn :: Var -> FastString -> VM a -> VM (VVar, a)
vectBndrNewIn v fs p
  = localV
  $ do
      vv <- vectBndrNew v fs
      x  <- p
      return (vv, x)

165
vectBndrsIn :: [Var] -> VM a -> VM ([VVar], a)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
166 167 168
vectBndrsIn vs p
  = localV
  $ do
169
      vvs <- mapM vectBndr vs
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
170
      x <- p
171
      return (vvs, x)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
172

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
173
-- ----------------------------------------------------------------------------
174 175
-- Expressions

176 177
vectVar :: Var -> VM VExpr
vectVar v
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
178 179 180
  = do
      r <- lookupVar v
      case r of
181 182 183
        Local (vv,lv) -> return (Var vv, Var lv)
        Global vv     -> do
                           let vexpr = Var vv
184
                           lexpr <- liftPD vexpr
185
                           return (vexpr, lexpr)
186

187 188
vectPolyVar :: Var -> [Type] -> VM VExpr
vectPolyVar v tys
189
  = do
190
      vtys <- mapM vectType tys
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
191
      r <- lookupVar v
192
      case r of
193 194 195 196
        Local (vv, lv) -> liftM2 (,) (polyApply (Var vv) vtys)
                                     (polyApply (Var lv) vtys)
        Global poly    -> do
                            vexpr <- polyApply (Var poly) vtys
197
                            lexpr <- liftPD vexpr
198
                            return (vexpr, lexpr)
199

200 201
vectLiteral :: Literal -> VM VExpr
vectLiteral lit
202
  = do
203
      lexpr <- liftPD (Lit lit)
204 205
      return (Lit lit, lexpr)

206 207 208 209 210 211 212 213 214 215 216 217 218
vectPolyExpr :: Bool -> CoreExprWithFVs -> VM (Inline, VExpr)
vectPolyExpr loop_breaker (_, AnnNote note expr)
  = do
      (inline, expr') <- vectPolyExpr loop_breaker expr
      return (inline, vNote note expr')
vectPolyExpr loop_breaker expr
  = do
      arity <- polyArity tvs
      polyAbstract tvs $ \args ->
        do
          (inline, mono') <- vectFnExpr False loop_breaker mono
          return (addInlineArity inline arity,
                  mapVect (mkLams $ tvs ++ args) mono')
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
219
  where
Ian Lynagh's avatar
Ian Lynagh committed
220 221
    (tvs, mono) = collectAnnTypeBinders expr

222 223
vectExpr :: CoreExprWithFVs -> VM VExpr
vectExpr (_, AnnType ty)
224
  = liftM vType (vectType ty)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
225

226
vectExpr (_, AnnVar v) = vectVar v
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
227

228
vectExpr (_, AnnLit lit) = vectLiteral lit
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
229

230 231
vectExpr (_, AnnNote note expr)
  = liftM (vNote note) (vectExpr expr)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
232

233
vectExpr e@(_, AnnApp _ arg)
234
  | isAnnTypeArg arg
235
  = vectTyAppExpr fn tys
236 237
  where
    (fn, tys) = collectAnnTypeArgs e
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
238

239 240 241 242 243
vectExpr (_, AnnApp (_, AnnVar v) (_, AnnLit lit))
  | Just con <- isDataConId_maybe v
  , is_special_con con
  = do
      let vexpr = App (Var v) (Lit lit)
244
      lexpr <- liftPD vexpr
245 246 247
      return (vexpr, lexpr)
  where
    is_special_con con = con `elem` [intDataCon, floatDataCon, doubleDataCon]
Ian Lynagh's avatar
Ian Lynagh committed
248

249

250
vectExpr (_, AnnApp fn arg)
251
  = do
252 253 254 255 256 257 258
      arg_ty' <- vectType arg_ty
      res_ty' <- vectType res_ty
      fn'     <- vectExpr fn
      arg'    <- vectExpr arg
      mkClosureApp arg_ty' res_ty' fn' arg'
  where
    (arg_ty, res_ty) = splitFunTy . exprType $ deAnnotate fn
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
259

260
vectExpr (_, AnnCase scrut bndr ty alts)
261 262 263
  | Just (tycon, ty_args) <- splitTyConApp_maybe scrut_ty
  , isAlgTyCon tycon
  = vectAlgCase tycon ty_args scrut bndr ty alts
264 265 266
  where
    scrut_ty = exprType (deAnnotate scrut)

267
vectExpr (_, AnnLet (AnnNonRec bndr rhs) body)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
268
  = do
269
      vrhs <- localV . inBind bndr . liftM snd $ vectPolyExpr False rhs
270
      (vbndr, vbody) <- vectBndrIn bndr (vectExpr body)
271
      return $ vLet (vNonRec vbndr vrhs) vbody
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
272

273
vectExpr (_, AnnLet (AnnRec bs) body)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
274
  = do
275 276
      (vbndrs, (vrhss, vbody)) <- vectBndrsIn bndrs
                                $ liftM2 (,)
277
                                  (zipWithM vect_rhs bndrs rhss)
278
                                  (vectExpr body)
279
      return $ vLet (vRec vbndrs vrhss) vbody
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
280
  where
281
    (bndrs, rhss) = unzip bs
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
282

283 284
    vect_rhs bndr rhs = localV
                      . inBind bndr
285 286
                      . liftM snd
                      $ vectPolyExpr (isLoopBreaker $ idOccInfo bndr) rhs
287

288
vectExpr e@(_, AnnLam bndr _)
289
  | isId bndr = liftM snd $ vectFnExpr True False e
290 291 292
{-
onlyIfV (isEmptyVarSet fvs) (vectScalarLam bs $ deAnnotate body)
                `orElseV` vectLam True fvs bs body
293 294
  where
    (bs,body) = collectAnnValBinders e
295
-}
296

297
vectExpr e = cantVectorise "Can't vectorise expression" (ppr $ deAnnotate e)
298

299 300 301 302 303
vectFnExpr :: Bool -> Bool -> CoreExprWithFVs -> VM (Inline, VExpr)
vectFnExpr inline loop_breaker e@(fvs, AnnLam bndr _)
  | isId bndr = onlyIfV (isEmptyVarSet fvs)
                        (mark DontInline . vectScalarLam bs $ deAnnotate body)
                `orElseV` mark inlineMe (vectLam inline loop_breaker fvs bs body)
304 305
  where
    (bs,body) = collectAnnValBinders e
306
vectFnExpr _ _ e = mark DontInline $ vectExpr e
307

308 309
mark :: Inline -> VM a -> VM (Inline, a)
mark b p = do { x <- p; return (b,x) }
310

311 312 313 314 315 316
vectScalarLam :: [Var] -> CoreExpr -> VM VExpr
vectScalarLam args body
  = do
      scalars <- globalScalars
      onlyIfV (all is_scalar_ty arg_tys
               && is_scalar_ty res_ty
317 318
               && is_scalar (extendVarSetList scalars args) body
               && uses scalars body)
319
        $ do
320
            fn_var <- hoistExpr (fsLit "fn") (mkLams args body) DontInline
321 322 323
            zipf <- zipScalars arg_tys res_ty
            clo <- scalarClosure arg_tys res_ty (Var fn_var)
                                                (zipf `App` Var fn_var)
324
            clo_var <- hoistExpr (fsLit "clo") clo DontInline
325
            lclo <- liftPD (Var clo_var)
326 327 328 329 330 331 332 333 334 335 336 337 338
            return (Var clo_var, lclo)
  where
    arg_tys = map idType args
    res_ty  = exprType body

    is_scalar_ty ty | Just (tycon, []) <- splitTyConApp_maybe ty
                    = tycon == intTyCon
                      || tycon == floatTyCon
                      || tycon == doubleTyCon

                    | otherwise = False

    is_scalar vs (Var v)     = v `elemVarSet` vs
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
339
    is_scalar _ e@(Lit _)    = is_scalar_ty $ exprType e
340 341 342
    is_scalar vs (App e1 e2) = is_scalar vs e1 && is_scalar vs e2
    is_scalar _ _            = False

343 344 345 346 347 348 349 350
    -- A scalar function has to actually compute something. Without the check,
    -- we would treat (\(x :: Int) -> x) as a scalar function and lift it to
    -- (map (\x -> x)) which is very bad. Normal lifting transforms it to
    -- (\n# x -> x) which is what we want.
    uses funs (Var v)     = v `elemVarSet` funs 
    uses funs (App e1 e2) = uses funs e1 || uses funs e2
    uses _ _              = False

351 352
vectLam :: Bool -> Bool -> VarSet -> [Var] -> CoreExprWithFVs -> VM VExpr
vectLam inline loop_breaker fvs bs body
353
  = do
354
      tyvars <- localTyVars
355 356 357
      (vs, vvs) <- readLEnv $ \env ->
                   unzip [(var, vv) | var <- varSetElems fvs
                                    , Just vv <- [lookupVarEnv (local_vars env) var]]
358

359 360 361
      arg_tys <- mapM (vectType . idType) bs
      res_ty  <- vectType (exprType $ deAnnotate body)

362
      buildClosures tyvars vvs arg_tys res_ty
363
        . hoistPolyVExpr tyvars (maybe_inline (length vs + length bs))
364
        $ do
365
            lc <- builtin liftingContext
366
            (vbndrs, vbody) <- vectBndrsIn (vs ++ bs)
367
                                           (vectExpr body)
368 369
            vbody' <- break_loop lc res_ty vbody
            return $ vLams lc vbndrs vbody'
370
  where
371 372 373 374 375 376 377 378 379 380 381 382 383 384
    maybe_inline n | inline    = Inline n
                   | otherwise = DontInline

    break_loop lc ty (ve, le)
      | loop_breaker
      = do
          empty <- emptyPD ty
          lty <- mkPDataType ty
          return (ve, mkWildCase (Var lc) intPrimTy lty
                        [(DEFAULT, [], le),
                         (LitAlt (mkMachInt 0), [], empty)])

      | otherwise = return (ve, le)
 
Ian Lynagh's avatar
Ian Lynagh committed
385

386 387
vectTyAppExpr :: CoreExprWithFVs -> [Type] -> VM VExpr
vectTyAppExpr (_, AnnVar v) tys = vectPolyVar v tys
388 389
vectTyAppExpr e tys = cantVectorise "Can't vectorise expression"
                        (ppr $ deAnnotate e `mkTyApps` tys)
390 391 392 393 394 395 396

-- We convert
--
--   case e :: t of v { ... }
--
-- to
--
397 398
--   V:    let v' = e in case v' of _ { ... }
--   L:    let v' = e in case v' `cast` ... of _ { ... }
399 400
--
-- When lifting, we have to do it this way because v must have the type
401 402
-- [:V(T):] but the scrutinee must be cast to the representation type. We also
-- have to handle the case where v is a wild var correctly.
Ian Lynagh's avatar
Ian Lynagh committed
403
--
404 405

-- FIXME: this is too lazy
Ian Lynagh's avatar
Ian Lynagh committed
406 407 408 409
vectAlgCase :: TyCon -> [Type] -> CoreExprWithFVs -> Var -> Type
            -> [(AltCon, [Var], CoreExprWithFVs)]
            -> VM VExpr
vectAlgCase _tycon _ty_args scrut bndr ty [(DEFAULT, [], body)]
410
  = do
411 412
      vscrut         <- vectExpr scrut
      (vty, lty)     <- vectAndLiftType ty
413 414 415
      (vbndr, vbody) <- vectBndrIn bndr (vectExpr body)
      return $ vCaseDEFAULT vscrut vbndr vty lty vbody

Ian Lynagh's avatar
Ian Lynagh committed
416
vectAlgCase _tycon _ty_args scrut bndr ty [(DataAlt _, [], body)]
417
  = do
418 419
      vscrut         <- vectExpr scrut
      (vty, lty)     <- vectAndLiftType ty
420 421 422
      (vbndr, vbody) <- vectBndrIn bndr (vectExpr body)
      return $ vCaseDEFAULT vscrut vbndr vty lty vbody

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
423
vectAlgCase _tycon _ty_args scrut bndr ty [(DataAlt dc, bndrs, body)]
424
  = do
425 426
      (vty, lty) <- vectAndLiftType ty
      vexpr      <- vectExpr scrut
427 428 429 430 431 432
      (vbndr, (vbndrs, (vect_body, lift_body)))
         <- vect_scrut_bndr
          . vectBndrsIn bndrs
          $ vectExpr body
      let (vect_bndrs, lift_bndrs) = unzip vbndrs
      (vscrut, lscrut, pdata_tc, _arg_tys) <- mkVScrut (vVar vbndr)
433
      vect_dc <- maybeV (lookupDataCon dc)
434 435 436 437 438 439
      let [pdata_dc] = tyConDataCons pdata_tc

      let vcase = mk_wild_case vscrut vty vect_dc  vect_bndrs vect_body
          lcase = mk_wild_case lscrut lty pdata_dc lift_bndrs lift_body

      return $ vLet (vNonRec vbndr vexpr) (vcase, lcase)
440
  where
Ian Lynagh's avatar
Ian Lynagh committed
441
    vect_scrut_bndr | isDeadBinder bndr = vectBndrNewIn bndr (fsLit "scrut")
442 443
                    | otherwise         = vectBndrIn bndr

444 445 446
    mk_wild_case expr ty dc bndrs body
      = mkWildCase expr (exprType expr) ty [(DataAlt dc, bndrs, body)]

Ian Lynagh's avatar
Ian Lynagh committed
447
vectAlgCase tycon _ty_args scrut bndr ty alts
448
  = do
449 450
      vect_tc     <- maybeV (lookupTyCon tycon)
      (vty, lty)  <- vectAndLiftType ty
451

452 453 454 455 456 457 458
      let arity = length (tyConDataCons vect_tc)
      sel_ty <- builtin (selTy arity)
      sel_bndr <- newLocalVar (fsLit "sel") sel_ty
      let sel = Var sel_bndr

      (vbndr, valts) <- vect_scrut_bndr
                      $ mapM (proc_alt arity sel vty lty) alts'
459 460 461
      let (vect_dcs, vect_bndrss, lift_bndrss, vbodies) = unzip4 valts

      vexpr <- vectExpr scrut
462 463
      (vect_scrut, lift_scrut, pdata_tc, _arg_tys) <- mkVScrut (vVar vbndr)
      let [pdata_dc] = tyConDataCons pdata_tc
464

465
      let (vect_bodies, lift_bodies) = unzip vbodies
466

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
467 468 469
      vdummy <- newDummyVar (exprType vect_scrut)
      ldummy <- newDummyVar (exprType lift_scrut)
      let vect_case = Case vect_scrut vdummy vty
470 471
                           (zipWith3 mk_vect_alt vect_dcs vect_bndrss vect_bodies)

472 473
      lc <- builtin liftingContext
      lbody <- combinePD vty (Var lc) sel lift_bodies
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
474
      let lift_case = Case lift_scrut ldummy lty
475
                           [(DataAlt pdata_dc, sel_bndr : concat lift_bndrss,
476 477 478 479 480
                             lbody)]

      return . vLet (vNonRec vbndr vexpr)
             $ (vect_case, lift_case)
  where
Ian Lynagh's avatar
Ian Lynagh committed
481
    vect_scrut_bndr | isDeadBinder bndr = vectBndrNewIn bndr (fsLit "scrut")
482 483 484 485 486 487 488 489
                    | otherwise         = vectBndrIn bndr

    alts' = sortBy (\(alt1, _, _) (alt2, _, _) -> cmp alt1 alt2) alts

    cmp (DataAlt dc1) (DataAlt dc2) = dataConTag dc1 `compare` dataConTag dc2
    cmp DEFAULT       DEFAULT       = EQ
    cmp DEFAULT       _             = LT
    cmp _             DEFAULT       = GT
Ian Lynagh's avatar
Ian Lynagh committed
490
    cmp _             _             = panic "vectAlgCase/cmp"
491

492
    proc_alt arity sel _ lty (DataAlt dc, bndrs, body)
493 494
      = do
          vect_dc <- maybeV (lookupDataCon dc)
495 496 497 498
          let ntag = dataConTagZ vect_dc
              tag  = mkDataConTag vect_dc
              fvs  = freeVarsOf body `delVarSetList` bndrs

499
          sel_tags  <- liftM (`App` sel) (builtin (selTags arity))
500 501 502 503 504 505 506
          lc        <- builtin liftingContext
          elems     <- builtin (selElements arity ntag)

          (vbndrs, vbody)
            <- vectBndrsIn bndrs
             . localV
             $ do
507
                 binds    <- mapM (pack_var (Var lc) sel_tags tag)
508 509 510 511
                           . filter isLocalId
                           $ varSetElems fvs
                 (ve, le) <- vectExpr body
                 return (ve, Case (elems `App` sel) lc lty
512 513 514 515 516 517
                             [(DEFAULT, [], (mkLets (concat binds) le))])
                 -- empty    <- emptyPD vty
                 -- return (ve, Case (elems `App` sel) lc lty
                 --             [(DEFAULT, [], Let (NonRec flags_var flags_expr)
                 --                             $ mkLets (concat binds) le),
                 --               (LitAlt (mkMachInt 0), [], empty)])
518
          let (vect_bndrs, lift_bndrs) = unzip vbndrs
519 520
          return (vect_dc, vect_bndrs, lift_bndrs, vbody)

521
    proc_alt _ _ _ _ _ = panic "vectAlgCase/proc_alt"
522 523 524

    mk_vect_alt vect_dc bndrs body = (DataAlt vect_dc, bndrs, body)

525
    pack_var len tags t v
526 527 528 529 530 531
      = do
          r <- lookupVar v
          case r of
            Local (vv, lv) ->
              do
                lv'  <- cloneVar lv
532
                expr <- packByTagPD (idType vv) (Var lv) len tags t
533 534 535 536 537
                updLEnv (\env -> env { local_vars = extendVarEnv
                                                (local_vars env) v (vv, lv') })
                return [(NonRec lv' expr)]

            _ -> return []
538