diff options
Diffstat (limited to 'compiler/vectorise/Vectorise/Utils.hs')
| -rw-r--r-- | compiler/vectorise/Vectorise/Utils.hs | 19 |
1 files changed, 16 insertions, 3 deletions
diff --git a/compiler/vectorise/Vectorise/Utils.hs b/compiler/vectorise/Vectorise/Utils.hs index 255a6c5857..c5f1cb7cb1 100644 --- a/compiler/vectorise/Vectorise/Utils.hs +++ b/compiler/vectorise/Vectorise/Utils.hs @@ -7,6 +7,7 @@ module Vectorise.Utils ( -- * Annotated Exprs collectAnnTypeArgs, + collectAnnDictArgs, collectAnnTypeBinders, collectAnnValBinders, isAnnTypeArg, @@ -31,6 +32,7 @@ import Vectorise.Monad import Vectorise.Builtins import CoreSyn import CoreUtils +import Id import Type import Control.Monad @@ -43,17 +45,28 @@ collectAnnTypeArgs expr = go expr [] go (_, AnnApp f (_, AnnType ty)) tys = go f (ty : tys) go e tys = (e, tys) +collectAnnDictArgs :: AnnExpr Var ann -> (AnnExpr Var ann, [AnnExpr Var ann]) +collectAnnDictArgs expr = go expr [] + where + go e@(_, AnnApp f arg) dicts + | isPredTy . exprType . deAnnotate $ arg = go f (arg : dicts) + | otherwise = (e, dicts) + go e dicts = (e, dicts) + collectAnnTypeBinders :: AnnExpr Var ann -> ([Var], AnnExpr Var ann) collectAnnTypeBinders expr = go [] expr where - go bs (_, AnnLam b e) | isTyVar b = go (b:bs) e + go bs (_, AnnLam b e) | isTyVar b = go (b : bs) e go bs e = (reverse bs, e) +-- |Collect all consecutive value binders that are not dictionaries. +-- collectAnnValBinders :: AnnExpr Var ann -> ([Var], AnnExpr Var ann) collectAnnValBinders expr = go [] expr where - go bs (_, AnnLam b e) | isId b = go (b:bs) e - go bs e = (reverse bs, e) + go bs (_, AnnLam b e) | isId b + && (not . isPredTy . idType $ b) = go (b : bs) e + go bs e = (reverse bs, e) isAnnTypeArg :: AnnExpr b ann -> Bool isAnnTypeArg (_, AnnType _) = True |