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