Commit fdc19e64 authored by Ian Lynagh's avatar Ian Lynagh
Browse files

Merge branch 'master' of darcs.haskell.org:/srv/darcs//ghc

parents 85ae01ba 209e3750
...@@ -361,18 +361,18 @@ vectTopRhs recFs var expr ...@@ -361,18 +361,18 @@ vectTopRhs recFs var expr
rhs _globalScalar _isDFun (Just (_, expr')) -- Case (1) rhs _globalScalar _isDFun (Just (_, expr')) -- Case (1)
= return (inlineMe, False, expr') = return (inlineMe, False, expr')
rhs True False Nothing -- Case (2) rhs True False Nothing -- Case (2)
= do { expr' <- vectScalarFun recFs expr = do { expr' <- vectScalarFun expr
; return (inlineMe, True, vectorised expr') ; return (inlineMe, True, vectorised expr')
} }
rhs True True Nothing -- Case (3) rhs True True Nothing -- Case (3)
= do { expr' <- vectScalarDFun var recFs = do { expr' <- vectScalarDFun var
; return (DontInline, True, expr') ; return (DontInline, True, expr')
} }
rhs False False Nothing -- Case (4) — not a dfun rhs False False Nothing -- Case (4) — not a dfun
= do { let exprFvs = freeVars expr = do { let exprFvs = freeVars expr
; (inline, isScalar, vexpr) ; (inline, isScalar, vexpr)
<- inBind var $ <- inBind var $
vectPolyExpr (isStrongLoopBreaker $ idOccInfo var) recFs exprFvs vectPolyExpr (isStrongLoopBreaker $ idOccInfo var) recFs exprFvs Nothing
; return (inline, isScalar, vectorised vexpr) ; return (inline, isScalar, vectorised vexpr)
} }
rhs False True Nothing -- Case (4) — is a dfun rhs False True Nothing -- Case (4) — is a dfun
......
...@@ -51,271 +51,120 @@ import TcRnMonad (doptM) ...@@ -51,271 +51,120 @@ import TcRnMonad (doptM)
import DynFlags (DynFlag(Opt_AvoidVect)) import DynFlags (DynFlag(Opt_AvoidVect))
-- For prototyping, the VITree is a separate data structure with the same shape as the corresponding expression -- Main entry point to vectorise expressions -----------------------------------
-- tree. This will become part of the annotation
data VectInfo = VIParr
| VISimple
| VIComplex
| VIEncaps
deriving (Eq, Show)
data VITree = VITNode VectInfo [VITree]
deriving (Show)
viTrace :: CoreExprWithFVs -> VectInfo -> [VITree] -> VM ()
viTrace ce vi vTs =
-- return ()
traceVt ("vitrace " ++ (show vi) ++ "[" ++ (concat $ map (\(VITNode vi _) -> show vi ++ " ") vTs) ++"]") (ppr $ deAnnotate ce)
viOr :: [VITree] -> Bool
viOr = or . (map (\(VITNode vi _) -> vi == VIParr))
-- TODO: free scalar vars don't actually need to be passed through, since encapsulations makes sure, that there are
-- no free variables in encapsulated lambda expressions
vectInfo:: CoreExprWithFVs -> VM VITree
vectInfo ce@(_, AnnVar v)
= do { vi <- vectInfoType $ exprType $ deAnnotate ce
; viTrace ce vi []
; traceVt "vectInfo AnnVar" ((ppr v) <+> (ppr $ exprType $ deAnnotate ce))
; return $ VITNode vi []
}
vectInfo ce@(_, AnnLit _)
= do { vi <- vectInfoType $ exprType $ deAnnotate ce
; viTrace ce vi []
; traceVt "vectInfo AnnLit" (ppr $ exprType $ deAnnotate ce)
; return $ VITNode vi []
}
vectInfo ce@(_, AnnApp e1 e2)
= do { vt1 <- vectInfo e1
; vt2 <- vectInfo e2
; vi <- if viOr [vt1, vt2]
then return VIParr
else vectInfoType $ exprType $ deAnnotate ce
; viTrace ce vi [vt1, vt2]
; return $ VITNode vi [vt1, vt2]
}
vectInfo ce@(_, AnnLam _var body)
= do { vt@(VITNode vi _) <- vectInfo body
; viTrace ce vi [vt]
; if (vi == VIParr)
then return $ VITNode vi [vt]
else return $ VITNode VIComplex [vt]
}
vectInfo ce@(_, AnnLet (AnnNonRec _var expr) body)
= do { vtE <- vectInfo expr
; vtB <- vectInfo body
; vi <- if viOr [vtE, vtB]
then return VIParr
else vectInfoType $ exprType $ deAnnotate ce
; viTrace ce vi [vtE, vtB]
; return $ VITNode vi [vtE, vtB]
}
vectInfo ce@(_, AnnLet (AnnRec bnds) body)
= do { let (_, exprs) = unzip bnds
; vtBnds <- mapM (\e -> vectInfo e) exprs
; if (viOr vtBnds)
then do { vtBnds' <- mapM (\e -> vectInfo e) exprs
; vtB <- vectInfo body
; return (VITNode VIParr (vtB: vtBnds'))
}
else do { vtB@(VITNode vib _) <- vectInfo body
; ni <- if (vib == VIParr)
then return VIParr
else vectInfoType $ exprType $ deAnnotate ce
; viTrace ce ni (vtB : vtBnds)
; return $ VITNode ni (vtB : vtBnds)
}
}
vectInfo ce@(_, AnnCase expr _var _ty alts)
= do { vtExpr <- vectInfo expr
; vtAlts <- mapM (\(_, _, e) -> vectInfo e) alts
; ni <- if viOr (vtExpr : vtAlts)
then return VIParr
else vectInfoType $ exprType $ deAnnotate ce
; viTrace ce ni (vtExpr : vtAlts)
; return $ VITNode ni (vtExpr: vtAlts)
}
vectInfo (_, AnnCast expr _)
= do { vt@(VITNode vi _) <- vectInfo expr
; return $ VITNode vi [vt]
}
vectInfo (_, AnnTick _ expr )
= do { vt@(VITNode vi _) <- vectInfo expr
; return $ VITNode vi [vt]
}
vectInfo (_, AnnType {})
= return $ VITNode VISimple []
vectInfo (_, AnnCoercion {})
= return $ VITNode VISimple []
vectInfoType:: Type -> VM VectInfo
vectInfoType ty
| maybeParrTy ty = return VIParr
| otherwise
= do { sType <- isSimpleType ty
; if sType
then return VISimple
else return VIComplex
}
-- Checks whether the type might be a parallel array type. In particular, if the outermost
-- constructor is a type family, we conservatively assume that it may be a parallel array type.
maybeParrTy :: Type -> Bool
maybeParrTy ty
| Just ty' <- coreView ty = maybeParrTy ty'
| Just (tyCon, ts) <- splitTyConApp_maybe ty = isPArrTyCon tyCon || isSynFamilyTyCon tyCon
|| or (map maybeParrTy ts)
maybeParrTy _ = False
isSimpleType:: Type -> VM Bool -- |Vectorise a polymorphic expression.
isSimpleType ty --
| Just (c, _cs) <- splitTyConApp_maybe ty = return $ (tyConName c) `elem` [boolTyConName, intTyConName, word8TyConName, doubleTyConName, floatTyConName] -- If not yet available, precompute vectorisation avoidance information before vectorising. If
{- -- the vectorisation avoidance optimisation is enabled, also use the vectorisation avoidance
= do { globals <- globalScalarTyCons -- information to encapsulated subexpression that do not need to be vectorised.
; traceVt ("isSimpleType " ++ (show (elemNameSet (tyConName c) globals ))) (ppr c) --
; return (elemNameSet (tyConName c) globals ) vectPolyExpr :: Bool -> [Var] -> CoreExprWithFVs -> Maybe VITree
} -> VM (Inline, Bool, VExpr)
-} -- precompute vectorisation avoidance information (and possibly encapsulated subexpressions)
| Nothing <- splitTyConApp_maybe ty vectPolyExpr loop_breaker recFns expr Nothing
= return False = do
isSimpleType ty { vectAvoidance <- liftDs $ doptM Opt_AvoidVect
= pprPanic "Vectorise.Exp.isSimpleType not handled" (ppr ty) ; vi <- vectAvoidInfo expr
; (expr', vi') <-
varsSimple :: VarSet -> VM Bool if vectAvoidance
varsSimple vs then do
= do { varTypes <- mapM isSimpleType $ map varType $ varSetElems vs { (expr', vi') <- encapsulateScalars vi expr
; return $ and varTypes ; traceVt "vectPolyExpr encapsulated:" (ppr $ deAnnotate expr')
} ; return (expr', vi')
}
else return (expr, vi)
-- | Vectorise a polymorphic expression. ; vectPolyExpr loop_breaker recFns expr' (Just vi')
vectPolyExpr:: Bool -> [Var] -> CoreExprWithFVs }
-> VM (Inline, Bool, VExpr)
vectPolyExpr loop_breaker recFns (_, AnnTick tickish expr) -- traverse through ticks
= do { (inline, isScalarFn, expr') <- vectPolyExpr loop_breaker recFns expr vectPolyExpr loop_breaker recFns (_, AnnTick tickish expr) (Just (VITNode _ [vit]))
; return (inline, isScalarFn, vTick tickish expr') = do
} { (inline, isScalarFn, expr') <- vectPolyExpr loop_breaker recFns expr (Just vit)
; return (inline, isScalarFn, vTick tickish expr')
}
vectPolyExpr loop_breaker recFns expr -- collect and vectorise type abstractions; then, descent into the body
= do { vectAvoidance <- liftDs $ doptM Opt_AvoidVect vectPolyExpr loop_breaker recFns expr (Just vit)
; vi <- vectInfo expr = do
; ((tvs, mono), vi') <- { let (tvs, mono) = collectAnnTypeBinders expr
if vectAvoidance vit' = stripLevels (length tvs) vit
then do { (extExpr, vi') <- encapsulateScalar vi expr ; arity <- polyArity tvs
; traceVt "vectPolyExpr extended:" (ppr $ deAnnotate extExpr) ; polyAbstract tvs $ \args ->
; return $ (collectAnnTypeBinders extExpr , vi') do
} { (inline, isScalarFn, mono') <- vectFnExpr False loop_breaker recFns mono vit'
else return $ (collectAnnTypeBinders expr, vi) ; return (addInlineArity inline arity, isScalarFn, mapVect (mkLams $ tvs ++ args) mono')
; arity <- polyArity tvs }
; polyAbstract tvs $ \args -> }
do {(inline, isScalarFn, mono') <- vectFnExpr False loop_breaker recFns mono vi' where
; return (addInlineArity inline arity, isScalarFn, mapVect (mkLams $ tvs ++ args) mono') stripLevels 0 vit = vit
} stripLevels n (VITNode _ [vit]) = stripLevels (n - 1) vit
} stripLevels _ vit = pprPanic "vectPolyExpr: stripLevels:" (text (show vit))
-- todo: clean this
vectPolyExprVT:: Bool -> [Var] -> CoreExprWithFVs -> VITree
-> VM (Inline, Bool, VExpr)
-- vectPolyExprVT _loop_breaker _recFns e vi | not (checkTree vi (deAnnotate e))
-- = pprPanic "vectPolyExprVT" (ppr $ deAnnotate e)
vectPolyExprVT loop_breaker recFns (_, AnnTick tickish expr) (VITNode _ [vit])
= do { (inline, isScalarFn, expr') <- vectPolyExprVT loop_breaker recFns expr vit
; return (inline, isScalarFn, vTick tickish expr')
}
vectPolyExprVT loop_breaker recFns expr vi
= do { -- checkTreeAnnM vi expr ;
let (tvs, mono) = collectAnnTypeBinders expr
; arity <- polyArity tvs
; polyAbstract tvs $ \args ->
do { (inline, isScalarFn, mono') <- vectFnExpr False loop_breaker recFns mono vi
; return (addInlineArity inline arity, isScalarFn, mapVect (mkLams $ tvs ++ args) mono')
}
}
-- | encapsulate every purely sequential subexpression with a simple return type -- Encapsulate every purely sequential subexpression of a (potentially) parallel expression into a
-- of a (potentially) parallel expression into a lambda abstraction over all its -- into a lambda abstraction over all its free variables followed by the corresponding application
-- free variables followed by the corresponding application to those variables. -- to those variables. We can, then, avoid the vectorisation of the ensapsulated subexpressions.
-- Condition: --
-- all free variables and the result type must be of `simple' type -- Preconditions:
-- the expression is 'complex enough', which is, for now, every expression --
-- which is not constant and contains at least one operation. -- * All free variables and the result type must be /simple/ types.
-- * The expression is sufficientlt complex (top warrant special treatment). For now, that is
-- every expression that is not constant and contains at least one operation.
-- --
encapsulateScalar :: VITree -> CoreExprWithFVs -> VM (CoreExprWithFVs, VITree) encapsulateScalars :: VITree -> CoreExprWithFVs -> VM (CoreExprWithFVs, VITree)
encapsulateScalar vit ce@(_, AnnType _ty) encapsulateScalars vit ce@(_, AnnType _ty)
= return (ce, vit) = return (ce, vit)
encapsulateScalar vit ce@(_, AnnVar _v) encapsulateScalars vit ce@(_, AnnVar _v)
= return (ce, vit) = return (ce, vit)
encapsulateScalar vit ce@(_, AnnLit _) encapsulateScalars vit ce@(_, AnnLit _)
= return (ce, vit) = return (ce, vit)
encapsulateScalars (VITNode vi [vit]) (fvs, AnnTick tck expr)
encapsulateScalar (VITNode vi [vit]) (fvs, AnnTick tck expr) = do { (extExpr, vit') <- encapsulateScalars vit expr
= do { (extExpr, vit') <- encapsulateScalar vit expr
; return ((fvs, AnnTick tck extExpr), VITNode vi [vit']) ; return ((fvs, AnnTick tck extExpr), VITNode vi [vit'])
} }
encapsulateScalar _ (_fvs, AnnTick _tck _expr) encapsulateScalars _ (_fvs, AnnTick _tck _expr)
= panic "encapsulateScalar AnnTick doesn't match up" = panic "encapsulateScalar AnnTick doesn't match up"
encapsulateScalar (VITNode vi [vit]) ce@(fvs, AnnLam bndr expr) encapsulateScalars (VITNode vi [vit]) ce@(fvs, AnnLam bndr expr)
= do { varsS <- varsSimple fvs = do { varsS <- varsSimple fvs
; case (vi, varsS) of ; case (vi, varsS) of
(VISimple, True) -> do { let (e', vit') = encaps vit ce (VISimple, True) -> do { let (e', vit') = liftSimple vit ce
; return (e', vit') ; return (e', vit')
} }
_ -> do { (extExpr, vit') <- encapsulateScalar vit expr _ -> do { (extExpr, vit') <- encapsulateScalars vit expr
; return ((fvs, AnnLam bndr extExpr), VITNode vi [vit']) ; return ((fvs, AnnLam bndr extExpr), VITNode vi [vit'])
} }
} }
encapsulateScalar _ (_fvs, AnnLam _bndr _expr) encapsulateScalars _ (_fvs, AnnLam _bndr _expr)
= panic "encapsulateScalar AnnLam doesn't match up" = panic "encapsulateScalars AnnLam doesn't match up"
encapsulateScalars vt@(VITNode vi [vit1, vit2]) ce@(fvs, AnnApp ce1 ce2)
encapsulateScalar vt@(VITNode vi [vit1, vit2]) ce@(fvs, AnnApp ce1 ce2)
= do { varsS <- varsSimple fvs = do { varsS <- varsSimple fvs
; case (vi, varsS) of ; case (vi, varsS) of
(VISimple, True) -> do { let (e', vt') = encaps vt ce (VISimple, True) -> do { let (e', vt') = liftSimple vt ce
-- ; checkTreeAnnM vt' e' -- ; checkTreeAnnM vt' e'
-- ; traceVt "Passed checkTree test!!" (ppr $ deAnnotate e') -- ; traceVt "Passed checkTree test!!" (ppr $ deAnnotate e')
; return (e', vt') ; return (e', vt')
} }
_ -> do { (etaCe1, vit1') <- encapsulateScalar vit1 ce1 _ -> do { (etaCe1, vit1') <- encapsulateScalars vit1 ce1
; (etaCe2, vit2') <- encapsulateScalar vit2 ce2 ; (etaCe2, vit2') <- encapsulateScalars vit2 ce2
; return ((fvs, AnnApp etaCe1 etaCe2), VITNode vi [vit1', vit2']) ; return ((fvs, AnnApp etaCe1 etaCe2), VITNode vi [vit1', vit2'])
} }
} }
encapsulateScalar _ (_fvs, AnnApp _ce1 _ce2)
= panic "encapsulateScalar AnnApp doesn't match up" encapsulateScalars _ (_fvs, AnnApp _ce1 _ce2)
= panic "encapsulateScalars AnnApp doesn't match up"
encapsulateScalar vt@(VITNode vi (scrutVit : altVits)) ce@(fvs, AnnCase scrut bndr ty alts) encapsulateScalars vt@(VITNode vi (scrutVit : altVits)) ce@(fvs, AnnCase scrut bndr ty alts)
= do { varsS <- varsSimple fvs = do { varsS <- varsSimple fvs
; case (vi, varsS) of ; case (vi, varsS) of
(VISimple, True) -> return $ encaps vt ce (VISimple, True) -> return $ liftSimple vt ce
_ -> do { (extScrut, scrutVit') <- encapsulateScalar scrutVit scrut _ -> do { (extScrut, scrutVit') <- encapsulateScalars scrutVit scrut
; extAltsVits <- zipWithM expAlt altVits alts ; extAltsVits <- zipWithM expAlt altVits alts
; let (extAlts, altVits') = unzip extAltsVits ; let (extAlts, altVits') = unzip extAltsVits
; return ((fvs, AnnCase extScrut bndr ty extAlts), VITNode vi (scrutVit': altVits')) ; return ((fvs, AnnCase extScrut bndr ty extAlts), VITNode vi (scrutVit': altVits'))
...@@ -323,110 +172,100 @@ encapsulateScalar vt@(VITNode vi (scrutVit : altVits)) ce@(fvs, AnnCase scrut bn ...@@ -323,110 +172,100 @@ encapsulateScalar vt@(VITNode vi (scrutVit : altVits)) ce@(fvs, AnnCase scrut bn
} }
where where
expAlt vt (con, bndrs, expr) expAlt vt (con, bndrs, expr)
= do { (extExpr, vt') <- encapsulateScalar vt expr = do { (extExpr, vt') <- encapsulateScalars vt expr
; return ((con, bndrs, extExpr), vt') ; return ((con, bndrs, extExpr), vt')
} }
encapsulateScalar _ (_fvs, AnnCase _scrut _bndr _ty _alts) encapsulateScalars _ (_fvs, AnnCase _scrut _bndr _ty _alts)
= panic "encapsulateScalar AnnCase doesn't match up" = panic "encapsulateScalars AnnCase doesn't match up"
encapsulateScalar vt@(VITNode vi [vt1, vt2]) ce@(fvs, AnnLet (AnnNonRec bndr expr1) expr2) encapsulateScalars vt@(VITNode vi [vt1, vt2]) ce@(fvs, AnnLet (AnnNonRec bndr expr1) expr2)
= do { varsS <- varsSimple fvs = do { varsS <- varsSimple fvs
; case (vi, varsS) of ; case (vi, varsS) of
(VISimple, True) -> return $ encaps vt ce (VISimple, True) -> return $ liftSimple vt ce
_ -> do { (extExpr1, vt1') <- encapsulateScalar vt1 expr1 _ -> do { (extExpr1, vt1') <- encapsulateScalars vt1 expr1
; (extExpr2, vt2') <- encapsulateScalar vt2 expr2 ; (extExpr2, vt2') <- encapsulateScalars vt2 expr2
; return ((fvs, AnnLet (AnnNonRec bndr extExpr1) extExpr2), VITNode vi [vt1', vt2']) ; return ((fvs, AnnLet (AnnNonRec bndr extExpr1) extExpr2), VITNode vi [vt1', vt2'])
} }
} }
encapsulateScalar _ (_fvs, AnnLet (AnnNonRec _bndr _expr1) _expr2) encapsulateScalars _ (_fvs, AnnLet (AnnNonRec _bndr _expr1) _expr2)
= panic "encapsulateScalar AnnLet nonrec doesn't match up" = panic "encapsulateScalars AnnLet nonrec doesn't match up"
encapsulateScalar vt@(VITNode vi (vtB : vtBnds)) ce@(fvs, AnnLet (AnnRec bndngs) expr) encapsulateScalars vt@(VITNode vi (vtB : vtBnds)) ce@(fvs, AnnLet (AnnRec bndngs) expr)
= do { varsS <- varsSimple fvs = do { varsS <- varsSimple fvs
; case (vi, varsS) of ; case (vi, varsS) of
(VISimple, True) -> return $ encaps vt ce (VISimple, True) -> return $ liftSimple vt ce
_ -> do { extBndsVts <- zipWithM expBndg vtBnds bndngs _ -> do { extBndsVts <- zipWithM expBndg vtBnds bndngs
; let (extBnds, vtBnds') = unzip extBndsVts ; let (extBnds, vtBnds') = unzip extBndsVts
; (extExpr, vtB') <- encapsulateScalar vtB expr ; (extExpr, vtB') <- encapsulateScalars vtB expr
; let vt' = VITNode vi (vtB':vtBnds') ; let vt' = VITNode vi (vtB':vtBnds')
; return ((fvs, AnnLet (AnnRec extBnds) extExpr), vt') ; return ((fvs, AnnLet (AnnRec extBnds) extExpr), vt')
} }
} }
where where
expBndg vit (bndr, expr) expBndg vit (bndr, expr)
= do { (extExpr, vit') <- encapsulateScalar vit expr = do { (extExpr, vit') <- encapsulateScalars vit expr
; return ((bndr, extExpr), vit') ; return ((bndr, extExpr), vit')
} }
encapsulateScalar _ (_fvs, AnnLet (AnnRec _) _expr2) encapsulateScalars _ (_fvs, AnnLet (AnnRec _) _expr2)
= panic "encapsulateScalar AnnLet rec doesn't match up" = panic "encapsulateScalars AnnLet rec doesn't match up"
encapsulateScalars (VITNode vi [vit]) (fvs, AnnCast expr coercion)
= do { (extExpr, vit') <- encapsulateScalars vit expr
encapsulateScalar (VITNode vi [vit]) (fvs, AnnCast expr coercion)
= do { (extExpr, vit') <- encapsulateScalar vit expr
; return ((fvs, AnnCast extExpr coercion), VITNode vi [vit']) ; return ((fvs, AnnCast extExpr coercion), VITNode vi [vit'])
} }
encapsulateScalar _ (_fvs, AnnCast _expr _coercion) encapsulateScalars _ (_fvs, AnnCast _expr _coercion)
= panic "encapsulateScalar AnnCast rec doesn't match up" = panic "encapsulateScalars AnnCast rec doesn't match up"
encapsulateScalar _ _
= panic "encapsulateScalar case not handled"
encapsulateScalars _ _
= panic "encapsulateScalars case not handled"
-- Lambda-lift the given expression and apply it to the abstracted free variables.
--
-- CoreExprWithFVs, -- = AnnExpr Id VarSet -- If the expression is a case expression scrutinising anything but a primitive type, then lift
-- AnnExpr bndr VarSet = (annot, AnnExpr' bndr VarSet) -- each alternative individually.
-- AnnLam :: bndr -> (AnnExpr bndr VarSet) -> AnnExpr' bndr VarSet --
-- AnnLam bndr (AnnExpr bndr annot) liftSimple :: VITree -> CoreExprWithFVs -> (CoreExprWithFVs, VITree)
encaps :: VITree -> CoreExprWithFVs -> (CoreExprWithFVs, VITree) liftSimple (VITNode vi (scrutVit : altVits)) (fvs, AnnCase expr bndr t alts)
encaps (VITNode vi (scrutVit : altVits)) (fvs, AnnCase expr bndr t alts)
| Just (c,_) <- splitTyConApp_maybe (exprType $ deAnnotate $ expr), | Just (c,_) <- splitTyConApp_maybe (exprType $ deAnnotate $ expr),
(not $ elem c [boolTyCon, intTyCon, doubleTyCon, floatTyCon]) -- TODO: globalScalarTyCons (not $ elem c [boolTyCon, intTyCon, doubleTyCon, floatTyCon]) -- FIXME: shouldn't be hardcoded
= ((fvs, AnnCase expr bndr t alts'), VITNode vi (scrutVit : altVits')) = ((fvs, AnnCase expr bndr t alts'), VITNode vi (scrutVit : altVits'))
where
where (alts', altVits') = unzip $ map (\(ac,bndrs, (alt, avi)) -> ((ac,bndrs,alt), avi)) $
(alts', altVits') = unzip $ map (\(ac,bndrs, (alt, avi)) -> ((ac,bndrs,alt), avi)) $ zipWith (\(ac, bndrs, aex) -> \altVi -> (ac, bndrs, liftSimple altVi aex)) alts altVits
zipWith (\(ac, bndrs, aex) -> \altVi -> (ac, bndrs, encaps altVi aex)) alts altVits
liftSimple viTree ae@(fvs, _annEx)
encaps viTree ae@(fvs, _annEx)
= (mkAnnApps (mkAnnLams ae vars) vars, viTree') = (mkAnnApps (mkAnnLams ae vars) vars, viTree')
where where
mkViTreeLams (VITNode _ vits) [] = VITNode VIEncaps vits mkViTreeLams (VITNode _ vits) [] = VITNode VIEncaps vits
mkViTreeLams vi (_:vs) = VITNode VIEncaps [mkViTreeLams vi vs] mkViTreeLams vi (_:vs) = VITNode VIEncaps [mkViTreeLams vi vs]
mkViTreeApps vi [] = vi mkViTreeApps vi [] = vi
mkViTreeApps vi (_:vs) = VITNode VISimple [mkViTreeApps vi vs, VITNode VISimple []] mkViTreeApps vi (_:vs) = VITNode VISimple [mkViTreeApps vi vs, VITNode VISimple []]
vars = varSetElems fvs
viTree' = mkViTreeApps (mkViTreeLams viTree vars) vars
mkAnnLam :: bndr -> AnnExpr bndr VarSet -> AnnExpr' bndr VarSet
mkAnnLam bndr ce = AnnLam bndr ce
vars = varSetElems fvs mkAnnLams:: CoreExprWithFVs -> [Var] -> CoreExprWithFVs
viTree' = mkViTreeApps (mkViTreeLams viTree vars) vars mkAnnLams (fv, aex') [] = (fv, aex') -- fv should be empty. check!
mkAnnLams (fv, aex') (v:vs) = mkAnnLams (delVarSet fv v, (mkAnnLam v ((delVarSet fv v), aex'))) vs
mkAnnLam :: bndr -> AnnExpr bndr VarSet -> AnnExpr' bndr VarSet mkAnnApp :: (AnnExpr bndr VarSet) -> Var -> (AnnExpr' bndr VarSet)
mkAnnLam bndr ce = AnnLam bndr ce mkAnnApp aex v = AnnApp aex (unitVarSet v, (AnnVar v))
mkAnnLams:: CoreExprWithFVs -> [Var] -> CoreExprWithFVs mkAnnApps:: CoreExprWithFVs -> [Var] -> CoreExprWithFVs
mkAnnLams (fv, aex') [] = (fv, aex') -- fv should be empty. check! mkAnnApps (fv, aex') [] = (fv, aex')
mkAnnLams (fv, aex') (v:vs) = mkAnnLams (delVarSet fv v, (mkAnnLam v ((delVarSet fv v), aex'))) vs mkAnnApps ae (v:vs) =
let
mkAnnApp :: (AnnExpr bndr VarSet) -> Var -> (AnnExpr' bndr VarSet) (fv, aex') = mkAnnApps ae vs
mkAnnApp aex v = AnnApp aex (unitVarSet v, (AnnVar v)) in (extendVarSet fv v, mkAnnApp (fv, aex') v)
mkAnnApps:: CoreExprWithFVs -> [Var] -> CoreExprWithFVs
mkAnnApps (fv, aex') [] = (fv, aex')
mkAnnApps ae (v:vs) =