Fix bug in vectorised DataCon worker generation

[ghc-hetmet.git] / compiler / vectorise / VectType.hs

diff --git a/compiler/vectorise/VectType.hs b/compiler/vectorise/VectType.hs
index 3e6fa2d..96745c5 100644 (file)
--- a/compiler/vectorise/VectType.hs
+++ b/compiler/vectorise/VectType.hs
@@ -1,6 +1,5 @@
 module VectType ( vectTyCon, vectType, vectTypeEnv,
-                   PAInstance, buildPADict,
-                   vectDataConWorkers )
+                   PAInstance, buildPADict )
 where
 
 #include "HsVersions.h"
@@ -37,7 +36,7 @@ import Digraph           ( SCC(..), stronglyConnComp )
 import Outputable
 
 import Control.Monad  ( liftM, liftM2, zipWithM, zipWithM_ )
-import Data.List      ( inits, tails )
+import Data.List      ( inits, tails, zipWith4 )
 
 -- ----------------------------------------------------------------------------
 -- Types
@@ -84,7 +83,7 @@ data PAInstance = PAInstance {
                   , painstArrTyCon  :: TyCon
                   }
 
-vectTypeEnv :: TypeEnv -> VM (TypeEnv, [FamInst])
+vectTypeEnv :: TypeEnv -> VM (TypeEnv, [FamInst], [(Var, CoreExpr)])
 vectTypeEnv env
   = do
       cs <- readGEnv $ mk_map . global_tycons
@@ -100,6 +99,7 @@ vectTypeEnv env
       parr_tcs <- zipWithM buildPArrayTyCon orig_tcs vect_tcs
       dfuns    <- mapM mkPADFun vect_tcs
       defTyConPAs (zip vect_tcs dfuns)
+      binds    <- sequence (zipWith4 buildTyConBindings orig_tcs vect_tcs parr_tcs dfuns)
       
       let all_new_tcs = new_tcs ++ parr_tcs
 
@@ -108,7 +108,7 @@ vectTypeEnv env
                         ++ [ADataCon dc | tc <- all_new_tcs
                                         , dc <- tyConDataCons tc])
 
-      return (new_env, map mkLocalFamInst parr_tcs)
+      return (new_env, map mkLocalFamInst parr_tcs, concat binds)
   where
     tycons = typeEnvTyCons env
     groups = tyConGroups tycons
@@ -303,31 +303,30 @@ tyConShape vect_tc
                                                e <- replicatePA len n
                                                return [e]
                }
-        
-vectDataConWorkers :: PAInstance -> VM [(Var, CoreExpr)]
-vectDataConWorkers (PAInstance { painstOrigTyCon = orig_tc
-                               , painstVectTyCon = vect_tc
-                               , painstArrTyCon  = arr_tc
-                               })
+
+buildTyConBindings :: TyCon -> TyCon -> TyCon -> Var -> VM [(Var, CoreExpr)]
+buildTyConBindings orig_tc vect_tc arr_tc dfun
   = do
       shape <- tyConShape vect_tc
-      sequence_ (zipWith3 (vectDataConWorker shape vect_tc arr_tc arr_dc)
-                          num_dcs
+      sequence_ (zipWith4 (vectDataConWorker shape vect_tc arr_tc arr_dc)
+                          orig_dcs
+                          vect_dcs
                           (inits repr_tys)
                           (tails repr_tys))
-      takeHoisted
+      dict <- buildPADict shape vect_tc arr_tc dfun
+      binds <- takeHoisted
+      return $ (dfun, dict) : binds
   where
     orig_dcs = tyConDataCons orig_tc
     vect_dcs = tyConDataCons vect_tc
     [arr_dc] = tyConDataCons arr_tc
 
-    num_dcs  = zip3 orig_dcs vect_dcs [0..]
     repr_tys = map dataConRepArgTys vect_dcs
 
 vectDataConWorker :: Shape -> TyCon -> TyCon -> DataCon
-                  -> (DataCon, DataCon, Int) -> [[Type]] -> [[Type]]
+                  -> DataCon -> DataCon -> [[Type]] -> [[Type]]
                   -> VM ()
-vectDataConWorker shape vect_tc arr_tc arr_dc (orig_dc, vect_dc, dc_num) pre (dc_tys : post)
+vectDataConWorker shape vect_tc arr_tc arr_dc orig_dc vect_dc pre (dc_tys : post)
   = do
       clo <- closedV
            . inBind orig_worker
@@ -351,7 +350,9 @@ vectDataConWorker shape vect_tc arr_tc arr_dc (orig_dc, vect_dc, dc_num) pre (dc
                 len     <- newLocalVar FSLIT("n") intPrimTy
                 arr_tys <- mapM mkPArrayType dc_tys
                 args    <- mapM (newLocalVar FSLIT("xs")) arr_tys
-                shapes  <- shapeReplicate shape (Var len) (mkIntLitInt dc_num)
+                shapes  <- shapeReplicate shape
+                                          (Var len)
+                                          (mkDataConTag vect_dc)
                 
                 empty_pre  <- mapM emptyPA (concat pre)
                 empty_post <- mapM emptyPA (concat post)
@@ -364,21 +365,17 @@ vectDataConWorker shape vect_tc arr_tc arr_dc (orig_dc, vect_dc, dc_num) pre (dc
                                           ++ map Var args
                                           ++ empty_post
 
-buildPADict :: PAInstance -> VM [(Var, CoreExpr)]
-buildPADict (PAInstance {
-               painstDFun      = dfun
-             , painstVectTyCon = vect_tc
-             , painstArrTyCon  = arr_tc })
-  = polyAbstract (tyConTyVars arr_tc) $ \abstract ->
+buildPADict :: Shape -> TyCon -> TyCon -> Var -> VM CoreExpr
+buildPADict shape vect_tc arr_tc dfun
+  = polyAbstract tvs $ \abstract ->
     do
-      shape <- tyConShape vect_tc
       meth_binds <- mapM (mk_method shape) paMethods
       let meth_exprs = map (Var . fst) meth_binds
 
       pa_dc <- builtin paDataCon
       let dict = mkConApp pa_dc (Type (mkTyConApp vect_tc arg_tys) : meth_exprs)
           body = Let (Rec meth_binds) dict
-      return [(dfun, mkInlineMe $ abstract body)]
+      return . mkInlineMe $ abstract body
   where
     tvs = tyConTyVars arr_tc
     arg_tys = mkTyVarTys tvs