compiler/vectorise/Vectorise/Type/PADict.hs

   1
   2 module Vectorise.Type.PADict
   3         (buildPADict)
   4 where
   5 import VectUtils
   6 import Vectorise.Monad
   7 import Vectorise.Builtins
   8 import Vectorise.Type.Repr
   9 import Vectorise.Type.PRepr
  10 import Vectorise.Type.PRDict
  11 import Vectorise.Utils.Hoisting
  12
  13 import BasicTypes
  14 import CoreSyn
  15 import CoreUtils
  16 import CoreUnfold
  17 import TyCon
  18 import Type
  19 import OccName
  20 import Id
  21 import Var
  22 import Name
  23
  24
  25
  26 buildPADict :: TyCon -> TyCon -> TyCon -> SumRepr -> VM Var
  27 buildPADict vect_tc prepr_tc arr_tc repr
  28   = polyAbstract tvs $ \args ->
  29     do
  30       method_ids <- mapM (method args) paMethods
  31
  32       pa_tc  <- builtin paTyCon
  33       pa_dc  <- builtin paDataCon
  34       let dict = mkLams (tvs ++ args)
  35                $ mkConApp pa_dc
  36                $ Type inst_ty : map (method_call args) method_ids
  37
  38           dfun_ty = mkForAllTys tvs
  39                   $ mkFunTys (map varType args) (mkTyConApp pa_tc [inst_ty])
  40
  41       raw_dfun <- newExportedVar dfun_name dfun_ty
  42       let dfun = raw_dfun `setIdUnfolding` mkDFunUnfolding dfun_ty (map Var method_ids)
  43                           `setInlinePragma` dfunInlinePragma
  44
  45       hoistBinding dfun dict
  46       return dfun
  47   where
  48     tvs = tyConTyVars vect_tc
  49     arg_tys = mkTyVarTys tvs
  50     inst_ty = mkTyConApp vect_tc arg_tys
  51
  52     dfun_name = mkPADFunOcc (getOccName vect_tc)
  53
  54     method args (name, build)
  55       = localV
  56       $ do
  57           expr <- build vect_tc prepr_tc arr_tc repr
  58           let body = mkLams (tvs ++ args) expr
  59           raw_var <- newExportedVar (method_name name) (exprType body)
  60           let var = raw_var
  61                       `setIdUnfolding` mkInlineRule body (Just (length args))
  62                       `setInlinePragma` alwaysInlinePragma
  63           hoistBinding var body
  64           return var
  65
  66     method_call args id = mkApps (Var id) (map Type arg_tys ++ map Var args)
  67
  68     method_name name = mkVarOcc $ occNameString dfun_name ++ ('$' : name)
  69
  70
  71 paMethods :: [(String, TyCon -> TyCon -> TyCon -> SumRepr -> VM CoreExpr)]
  72 paMethods = [("dictPRepr",    buildPRDict),
  73              ("toPRepr",      buildToPRepr),
  74              ("fromPRepr",    buildFromPRepr),
  75              ("toArrPRepr",   buildToArrPRepr),
  76              ("fromArrPRepr", buildFromArrPRepr)]
  77