compiler/vectorise/Vectorise/Utils/Poly.hs

   1
   2 module Vectorise.Utils.Poly (
   3         polyAbstract,
   4         polyApply,
   5         polyVApply,
   6         polyArity
   7 )
   8 where
   9 import Vectorise.Vect
  10 import Vectorise.Monad
  11 import Vectorise.Utils.PADict
  12 import CoreSyn
  13 import Type
  14 import FastString
  15 import Control.Monad
  16
  17
  18 -- Poly Functions -------------------------------------------------------------
  19 polyAbstract :: [TyVar] -> ([Var] -> VM a) -> VM a
  20 polyAbstract tvs p
  21   = localV
  22   $ do
  23       mdicts <- mapM mk_dict_var tvs
  24       zipWithM_ (\tv -> maybe (defLocalTyVar tv)
  25                               (defLocalTyVarWithPA tv . Var)) tvs mdicts
  26       p (mk_args mdicts)
  27   where
  28     mk_dict_var tv = do
  29                        r <- paDictArgType tv
  30                        case r of
  31                          Just ty -> liftM Just (newLocalVar (fsLit "dPA") ty)
  32                          Nothing -> return Nothing
  33
  34     mk_args mdicts = [dict | Just dict <- mdicts]
  35
  36
  37 polyArity :: [TyVar] -> VM Int
  38 polyArity tvs = do
  39                   tys <- mapM paDictArgType tvs
  40                   return $ length [() | Just _ <- tys]
  41
  42
  43 polyApply :: CoreExpr -> [Type] -> VM CoreExpr
  44 polyApply expr tys
  45  = do dicts <- mapM paDictOfType tys
  46       return $ expr `mkTyApps` tys `mkApps` dicts
  47
  48
  49 polyVApply :: VExpr -> [Type] -> VM VExpr
  50 polyVApply expr tys
  51  = do dicts <- mapM paDictOfType tys
  52       return     $ mapVect (\e -> e `mkTyApps` tys `mkApps` dicts) expr