Refactoring
[ghc-hetmet.git] / compiler / vectorise / VectType.hs
index 405d6ab..bef08f7 100644 (file)
@@ -32,6 +32,7 @@ import TysPrim           ( intPrimTy )
 import Unique
 import UniqFM
 import UniqSet
+import Util              ( singleton )
 import Digraph           ( SCC(..), stronglyConnComp )
 
 import Outputable
@@ -224,8 +225,30 @@ data Repr = ProdRepr {
             , sum_arr_data_con  :: DataCon
             }
 
-mkProduct :: [Type] -> VM Repr
-mkProduct tys
+          | IdRepr Type
+
+          | VoidRepr {
+              void_tycon        :: TyCon
+            , void_bottom       :: CoreExpr
+            }
+
+voidRepr :: VM Repr
+voidRepr
+  = do
+      tycon <- builtin voidTyCon
+      var   <- builtin voidVar
+      return $ VoidRepr {
+                 void_tycon  = tycon
+               , void_bottom = Var var
+               }
+
+unboxedProductRepr :: [Type] -> VM Repr
+unboxedProductRepr []   = voidRepr
+unboxedProductRepr [ty] = return $ IdRepr ty
+unboxedProductRepr tys  = boxedProductRepr tys
+
+boxedProductRepr :: [Type] -> VM Repr
+boxedProductRepr tys
   = do
       tycon <- builtin (prodTyCon arity)
       let [data_con] = tyConDataCons tycon
@@ -243,9 +266,10 @@ mkProduct tys
   where
     arity = length tys
 
-mkSum :: [Repr] -> VM Repr
-mkSum [repr] = return repr
-mkSum reprs
+sumRepr :: [Repr] -> VM Repr
+sumRepr []     = voidRepr
+sumRepr [repr] = boxRepr repr
+sumRepr reprs
   = do
       tycon <- builtin (sumTyCon arity)
       (arr_tycon, _) <- parrayReprTyCon
@@ -263,11 +287,22 @@ mkSum reprs
   where
     arity = length reprs
 
+splitSumRepr :: Repr -> [Repr]
+splitSumRepr (SumRepr { sum_components = reprs }) = reprs
+splitSumRepr repr                                 = [repr]
+
+boxRepr :: Repr -> VM Repr
+boxRepr (VoidRepr {}) = boxedProductRepr []
+boxRepr (IdRepr ty)   = boxedProductRepr [ty]
+boxRepr repr          = return repr
+
 reprType :: Repr -> Type
 reprType (ProdRepr { prod_tycon = tycon, prod_components = tys })
   = mkTyConApp tycon tys
 reprType (SumRepr { sum_tycon = tycon, sum_components = reprs })
   = mkTyConApp tycon (map reprType reprs)
+reprType (IdRepr ty) = ty
+reprType (VoidRepr { void_tycon = tycon }) = mkTyConApp tycon []
 
 arrReprType :: Repr -> VM Type
 arrReprType = mkPArrayType . reprType
@@ -277,7 +312,9 @@ arrShapeTys (SumRepr  {})
   = do
       int_arr <- builtin parrayIntPrimTyCon
       return [intPrimTy, mkTyConApp int_arr [], mkTyConApp int_arr []]
-arrShapeTys repr = return [intPrimTy]
+arrShapeTys (ProdRepr {}) = return [intPrimTy]
+arrShapeTys (IdRepr _)    = return []
+arrShapeTys (VoidRepr {}) = return [intPrimTy]
 
 arrShapeVars :: Repr -> VM [Var]
 arrShapeVars repr = mapM (newLocalVar FSLIT("sh")) =<< arrShapeTys repr
@@ -289,29 +326,47 @@ replicateShape (SumRepr {})  len tag
       rep <- builtin replicatePAIntPrimVar
       up  <- builtin upToPAIntPrimVar
       return [len, Var rep `mkApps` [len, tag], Var up `App` len]
-
-arrReprElemTys :: Repr -> [[Type]]
-arrReprElemTys (SumRepr { sum_components = prods })
-  = map arrProdElemTys prods
-arrReprElemTys prod@(ProdRepr {})
-  = [arrProdElemTys prod]
-
-arrProdElemTys (ProdRepr { prod_components = [] })
-  = [unitTy]
-arrProdElemTys (ProdRepr { prod_components = tys })
-  = tys
-
-arrReprTys :: Repr -> VM [[Type]]
-arrReprTys = mapM (mapM mkPArrayType) . arrReprElemTys
+replicateShape (IdRepr _) _ _ = return []
+replicateShape (VoidRepr {}) len _ = return [len]
+
+emptyArrRepr :: Repr -> VM [CoreExpr]
+emptyArrRepr (SumRepr { sum_components = prods })
+  = liftM concat $ mapM emptyArrRepr prods
+emptyArrRepr (ProdRepr { prod_components = [] })
+  = return [Var unitDataConId]
+emptyArrRepr (ProdRepr { prod_components = tys })
+  = mapM emptyPA tys
+emptyArrRepr (IdRepr ty)
+  = liftM singleton $ emptyPA ty
+emptyArrRepr (VoidRepr { void_tycon = tycon })
+  = liftM singleton $ emptyPA (mkTyConApp tycon [])
+
+arrReprTys :: Repr -> VM [Type]
+arrReprTys (SumRepr { sum_components = reprs })
+  = liftM concat $ mapM arrReprTys reprs
+arrReprTys (ProdRepr { prod_components = [] })
+  = return [unitTy]
+arrReprTys (ProdRepr { prod_components = tys })
+  = mapM mkPArrayType tys
+arrReprTys (IdRepr ty)
+  = liftM singleton $ mkPArrayType ty
+arrReprTys (VoidRepr { void_tycon = tycon })
+  = liftM singleton $ mkPArrayType (mkTyConApp tycon [])
+
+arrReprTys' :: Repr -> VM [[Type]]
+arrReprTys' (SumRepr { sum_components = reprs })
+  = mapM arrReprTys reprs
+arrReprTys' repr = liftM singleton $ arrReprTys repr
 
 arrReprVars :: Repr -> VM [[Var]]
 arrReprVars repr
-  = mapM (mapM (newLocalVar FSLIT("rs"))) =<< arrReprTys repr
+  = mapM (mapM (newLocalVar FSLIT("rs"))) =<< arrReprTys' repr
 
 mkRepr :: TyCon -> VM Repr
 mkRepr vect_tc
-  = mkSum
-  =<< mapM mkProduct (map dataConRepArgTys $ tyConDataCons vect_tc)
+  = sumRepr =<< mapM unboxedProductRepr rep_tys
+  where
+    rep_tys = map dataConRepArgTys $ tyConDataCons vect_tc
 
 buildPReprType :: TyCon -> VM Type
 buildPReprType = liftM reprType . mkRepr
@@ -358,6 +413,15 @@ buildToPRepr repr vect_tc prepr_tc _
           vars <- mapM (newLocalVar FSLIT("r")) tys
           return (vars, mkConApp data_con (map Type tys ++ map Var vars))
 
+    prod_alt (IdRepr ty)
+      = do
+          var <- newLocalVar FSLIT("y") ty
+          return ([var], Var var)
+
+    prod_alt (VoidRepr { void_bottom = bottom })
+      = return ([], bottom)
+
+
 buildFromPRepr :: Repr -> TyCon -> TyCon -> TyCon -> VM CoreExpr
 buildFromPRepr repr vect_tc prepr_tc _
   = do
@@ -397,6 +461,12 @@ buildFromPRepr repr vect_tc prepr_tc _
           return $ Case expr (mkWildId (reprType prod)) res_ty
                    [(DataAlt data_con, vars, con `mkVarApps` vars)]
 
+    from_prod (IdRepr _) con expr
+       = return $ con `App` expr
+
+    from_prod (VoidRepr {}) con expr
+       = return con
+
 buildToArrPRepr :: Repr -> TyCon -> TyCon -> TyCon -> VM CoreExpr
 buildToArrPRepr repr vect_tc prepr_tc arr_tc
   = do
@@ -435,7 +505,7 @@ buildToArrPRepr repr vect_tc prepr_tc arr_tc
                      , sum_arr_tycon    = tycon
                      , sum_arr_data_con = data_con })
       = do
-          exprs <- zipWithM (to_prod len_var) repr_vars prods
+          exprs <- zipWithM to_prod repr_vars prods
 
           return . wrapFamInstBody tycon tys
                  . mkConApp data_con
@@ -443,16 +513,28 @@ buildToArrPRepr repr vect_tc prepr_tc arr_tc
       where
         tys = map reprType prods
 
-    to_repr [len_var] [repr_vars] prod = to_prod len_var repr_vars prod
+    to_repr [len_var]
+            [repr_vars]
+            (ProdRepr { prod_components   = tys
+                      , prod_arr_tycon    = tycon
+                      , prod_arr_data_con = data_con })
+       = return . wrapFamInstBody tycon tys
+                . mkConApp data_con
+                $ map Type tys ++ map Var (len_var : repr_vars)
 
-    to_prod len_var
-            repr_vars
+    to_prod repr_vars@(r : _)
             (ProdRepr { prod_components   = tys
                       , prod_arr_tycon    = tycon
                       , prod_arr_data_con = data_con })
-      = return . wrapFamInstBody tycon tys
-               . mkConApp data_con
-               $ map Type tys ++ map Var (len_var : repr_vars)
+      = do
+          len <- lengthPA (Var r)
+          return . wrapFamInstBody tycon tys
+                 . mkConApp data_con
+                 $ map Type tys ++ len : map Var repr_vars
+
+    to_prod [var] (IdRepr ty)   = return (Var var)
+    to_prod [var] (VoidRepr {}) = return (Var var)
+
 
 buildFromArrPRepr :: Repr -> TyCon -> TyCon -> TyCon -> VM CoreExpr
 buildFromArrPRepr repr vect_tc prepr_tc arr_tc
@@ -531,7 +613,26 @@ buildFromArrPRepr repr vect_tc prepr_tc arr_tc
           return $ Case scrut (mkWildId scrut_ty) res_ty
                    [(DataAlt data_con, shape_vars ++ repr_vars, body)]
 
+    from_prod (IdRepr ty)
+              expr
+              shape_vars
+              [repr_var]
+              res_ty
+              body
+      = return $ Let (NonRec repr_var expr) body
+
+    from_prod (VoidRepr {})
+              expr
+              shape_vars
+              [repr_var]
+              res_ty
+              body
+      = return $ Let (NonRec repr_var expr) body
+
 buildPRDictRepr :: Repr -> VM CoreExpr
+buildPRDictRepr (VoidRepr { void_tycon = tycon })
+  = prDFunOfTyCon tycon
+buildPRDictRepr (IdRepr ty) = mkPR ty
 buildPRDictRepr (ProdRepr {
                    prod_components = tys
                  , prod_tycon      = tycon
@@ -601,7 +702,7 @@ buildPArrayDataCon orig_name vect_tc repr_tc
       shape_tys <- arrShapeTys repr
       repr_tys  <- arrReprTys  repr
 
-      let tys = shape_tys ++ concat repr_tys
+      let tys = shape_tys ++ repr_tys
 
       liftDs $ buildDataCon dc_name
                             False                  -- not infix
@@ -642,8 +743,8 @@ vectDataConWorkers repr orig_tc vect_tc arr_tc
           . zipWith3 def_worker  (tyConDataCons orig_tc) rep_tys
           $ zipWith4 mk_data_con (tyConDataCons vect_tc)
                                  rep_tys
-                                 (inits arr_tys)
-                                 (tail $ tails arr_tys)
+                                 (inits reprs)
+                                 (tail $ tails reprs)
       mapM_ (uncurry hoistBinding) bs
   where
     tyvars   = tyConTyVars vect_tc
@@ -653,18 +754,16 @@ vectDataConWorkers repr orig_tc vect_tc arr_tc
     res_ty   = mkTyConApp vect_tc var_tys
 
     rep_tys  = map dataConRepArgTys $ tyConDataCons vect_tc
-    arr_tys  = arrReprElemTys repr
+    reprs    = splitSumRepr repr
 
     [arr_dc] = tyConDataCons arr_tc
 
     mk_data_con con tys pre post
       = liftM2 (,) (vect_data_con con)
-                   (lift_data_con tys (concat pre)
-                                      (concat post)
-                                      (mkDataConTag con))
+                   (lift_data_con tys pre post (mkDataConTag con))
 
     vect_data_con con = return $ mkConApp con ty_args
-    lift_data_con tys pre_tys post_tys tag
+    lift_data_con tys pre_reprs post_reprs tag
       = do
           len  <- builtin liftingContext
           args <- mapM (newLocalVar FSLIT("xs"))
@@ -673,8 +772,8 @@ vectDataConWorkers repr orig_tc vect_tc arr_tc
           shape <- replicateShape repr (Var len) tag
           repr  <- mk_arr_repr (Var len) (map Var args)
 
-          pre   <- mapM emptyPA pre_tys
-          post  <- mapM emptyPA post_tys
+          pre   <- liftM concat $ mapM emptyArrRepr pre_reprs
+          post  <- liftM concat $ mapM emptyArrRepr post_reprs
 
           return . mkLams (len : args)
                  . wrapFamInstBody arr_tc var_tys