Complete PA dictionary generation for product types
[ghc-hetmet.git] / compiler / vectorise / VectType.hs
index 10c3bbf..c977e4c 100644 (file)
@@ -208,15 +208,140 @@ buildPReprTyCon orig_tc vect_tc
   where
     tyvars = tyConTyVars vect_tc
 
+data TyConRepr = ProdRepr {
+                   repr_prod_arg_tys      :: [Type]
+                 , repr_prod_tycon        :: TyCon
+                 , repr_prod_data_con     :: DataCon
+                 , repr_prod_arr_tycon    :: TyCon
+                 , repr_prod_arr_data_con :: DataCon
+                 , repr_type              :: Type
+                 }
+               | SumRepr {
+                   repr_tys            :: [[Type]]
+                 , repr_prod_tycons    :: [Maybe TyCon]
+                 , repr_prod_data_cons :: [Maybe DataCon]
+                 , repr_prod_tys       :: [Type]
+                 , repr_sum_tycon      :: TyCon
+                 , repr_type           :: Type
+                 }
+
+arrShapeTys :: TyConRepr -> VM [Type]
+arrShapeTys (ProdRepr {}) = return [intPrimTy]
+arrShapeTys (SumRepr  {})
+  = do
+      uarr <- builtin uarrTyCon
+      return [intPrimTy, mkTyConApp uarr [intTy]]
+
+arrReprTys :: TyConRepr -> VM [Type]
+arrReprTys (ProdRepr { repr_prod_arg_tys = tys })
+  = mapM mkPArrayType tys
+arrReprTys (SumRepr { repr_tys = tys })
+  = concat `liftM` mapM (mapM mkPArrayType) (map mk_prod tys)
+  where
+    mk_prod []  = [unitTy]
+    mk_prod tys = tys
+      
+
+mkTyConRepr :: TyCon -> VM TyConRepr
+mkTyConRepr vect_tc
+  | is_product
+  = let
+      [prod_arg_tys] = repr_tys
+      arity          = length prod_arg_tys
+    in
+    do
+      prod_tycon <- builtin (prodTyCon arity)
+      let [prod_data_con] = tyConDataCons prod_tycon
+
+      (arr_tycon, _) <- parrayReprTyCon
+                      . mkTyConApp prod_tycon
+                      $ replicate arity unitTy
+
+      let [arr_data_con] = tyConDataCons arr_tycon
+
+      return $ ProdRepr {
+                 repr_prod_arg_tys      = prod_arg_tys
+               , repr_prod_tycon        = prod_tycon
+               , repr_prod_data_con     = prod_data_con
+               , repr_prod_arr_tycon    = arr_tycon
+               , repr_prod_arr_data_con = arr_data_con
+               , repr_type              = mkTyConApp prod_tycon prod_arg_tys
+               }
+
+  | otherwise
+  = do
+      uarr <- builtin uarrTyCon
+      prod_tycons  <- mapM (mk_tycon prodTyCon) repr_tys
+      let prod_tys = zipWith mk_tc_app_maybe prod_tycons repr_tys
+      sum_tycon    <- builtin (sumTyCon $ length repr_tys)
+      arr_repr_tys <- mapM (mapM mkPArrayType . arr_repr_elem_tys) repr_tys
+
+      return $ SumRepr {
+                 repr_tys            = repr_tys
+               , repr_prod_tycons    = prod_tycons
+               , repr_prod_data_cons = map (fmap mk_single_datacon) prod_tycons
+               , repr_prod_tys       = prod_tys
+               , repr_sum_tycon      = sum_tycon
+               , repr_type           = mkTyConApp sum_tycon prod_tys
+               }
+  where
+    tyvars    = tyConTyVars vect_tc
+    data_cons = tyConDataCons vect_tc
+    repr_tys  = map dataConRepArgTys data_cons
+
+    is_product | [_] <- data_cons = True
+               | otherwise        = False
+
+    mk_shape uarr = intPrimTy : mk_sel uarr
+
+    mk_sel uarr | is_product = []
+                | otherwise  = [uarr_int, uarr_int]
+      where
+        uarr_int = mkTyConApp uarr [intTy]
+
+    mk_tycon get_tc tys
+      | n > 1     = builtin (Just . get_tc n)
+      | otherwise = return Nothing
+      where n = length tys
+
+    mk_single_datacon tc | [dc] <- tyConDataCons tc = dc
+
+    mk_tc_app_maybe Nothing   []   = unitTy
+    mk_tc_app_maybe Nothing   [ty] = ty
+    mk_tc_app_maybe (Just tc) tys  = mkTyConApp tc tys
+
+    arr_repr_elem_tys []  = [unitTy]
+    arr_repr_elem_tys tys = tys
+
 buildPReprType :: TyCon -> VM Type
 buildPReprType = liftM repr_type . mkTyConRepr
 
 buildToPRepr :: TyConRepr -> TyCon -> TyCon -> TyCon -> VM CoreExpr
-buildToPRepr (TyConRepr {
+buildToPRepr (ProdRepr {
+                repr_prod_arg_tys  = prod_arg_tys
+              , repr_prod_data_con = prod_data_con
+              , repr_type          = repr_type
+              })
+             vect_tc prepr_tc _
+  = do
+      arg  <- newLocalVar FSLIT("x") arg_ty
+      vars <- mapM (newLocalVar FSLIT("x")) prod_arg_tys
+
+      return . Lam arg
+             . wrapFamInstBody prepr_tc var_tys
+             $ Case (Var arg) (mkWildId arg_ty) repr_type
+               [(DataAlt data_con, vars,
+                 mkConApp prod_data_con (map Type prod_arg_tys ++ map Var vars))]
+  where
+    var_tys    = mkTyVarTys $ tyConTyVars vect_tc
+    arg_ty     = mkTyConApp vect_tc var_tys
+    [data_con] = tyConDataCons vect_tc
+
+buildToPRepr (SumRepr {
                 repr_tys            = repr_tys
               , repr_prod_data_cons = prod_data_cons
               , repr_prod_tys       = prod_tys
-              , repr_sum_data_cons  = sum_data_cons
+              , repr_sum_tycon      = sum_tycon
               , repr_type           = repr_type
               })
               vect_tc prepr_tc _
@@ -227,16 +352,14 @@ buildToPRepr (TyConRepr {
       return . Lam arg
              . wrapFamInstBody prepr_tc var_tys
              . Case (Var arg) (mkWildId arg_ty) repr_type
-             . mk_alts data_cons vars
+             . zipWith4 mk_alt data_cons vars sum_data_cons
              . zipWith3 mk_prod prod_data_cons repr_tys $ map (map Var) vars
   where
     var_tys   = mkTyVarTys $ tyConTyVars vect_tc
     arg_ty    = mkTyConApp vect_tc var_tys
     data_cons = tyConDataCons vect_tc
 
-    mk_alts _    _      []     = [(DEFAULT, [], Var unitDataConId)]
-    mk_alts [dc] [vars] [expr] = [(DataAlt dc, vars, expr)]
-    mk_alts dcs  vars   exprs  = zipWith4 mk_alt dcs vars sum_data_cons exprs 
+    sum_data_cons = tyConDataCons sum_tycon
 
     mk_alt dc vars sum_dc expr = (DataAlt dc, vars,
                                   mkConApp sum_dc (map Type prod_tys ++ [expr]))
@@ -246,11 +369,34 @@ buildToPRepr (TyConRepr {
     mk_prod (Just dc) tys exprs  = mkConApp dc (map Type tys ++ exprs)
 
 buildFromPRepr :: TyConRepr -> TyCon -> TyCon -> TyCon -> VM CoreExpr
-buildFromPRepr (TyConRepr {
+buildFromPRepr (ProdRepr {
+                  repr_prod_arg_tys  = prod_arg_tys
+                , repr_prod_data_con = prod_data_con
+                , repr_type          = repr_type
+                })
+             vect_tc prepr_tc _
+  = do
+      arg_ty <- mkPReprType res_ty
+      arg    <- newLocalVar FSLIT("x") arg_ty
+      vars   <- mapM (newLocalVar FSLIT("x")) prod_arg_tys
+
+      return . Lam arg
+             $ Case (unwrapFamInstScrut prepr_tc var_tys (Var arg))
+                    (mkWildId repr_type)
+                    res_ty
+               [(DataAlt prod_data_con, vars,
+                 mkConApp data_con (map Type var_tys ++ map Var vars))]
+  where
+    var_tys    = mkTyVarTys $ tyConTyVars vect_tc
+    ty_args    = map Type var_tys
+    res_ty     = mkTyConApp vect_tc var_tys
+    [data_con] = tyConDataCons vect_tc
+
+buildFromPRepr (SumRepr {
                 repr_tys            = repr_tys
               , repr_prod_data_cons = prod_data_cons
               , repr_prod_tys       = prod_tys
-              , repr_sum_data_cons  = sum_data_cons
+              , repr_sum_tycon      = sum_tycon
               , repr_type           = repr_type
               })
               vect_tc prepr_tc _
@@ -259,7 +405,10 @@ buildFromPRepr (TyConRepr {
       arg    <- newLocalVar FSLIT("x") arg_ty
 
       liftM (Lam arg
-             . un_sum (unwrapFamInstScrut prepr_tc var_tys (Var arg)))
+             . Case (unwrapFamInstScrut prepr_tc var_tys (Var arg))
+                    (mkWildId repr_type)
+                    res_ty
+             . zipWith mk_alt sum_data_cons)
             (sequence
              $ zipWith4 un_prod data_cons prod_data_cons prod_tys repr_tys)
   where
@@ -268,6 +417,8 @@ buildFromPRepr (TyConRepr {
     res_ty    = mkTyConApp vect_tc var_tys
     data_cons = tyConDataCons vect_tc
 
+    sum_data_cons = tyConDataCons sum_tycon
+
     un_prod dc _ _ []
       = do
           var <- newLocalVar FSLIT("u") unitTy
@@ -288,15 +439,56 @@ buildFromPRepr (TyConRepr {
 
           return (pv, expr)
 
-    un_sum scrut [(var, expr)] = Let (NonRec var scrut) expr
-    un_sum scrut alts
-      = Case scrut (mkWildId repr_type) res_ty
-      $ zipWith mk_alt sum_data_cons alts
-
     mk_alt sum_dc (var, expr) = (DataAlt sum_dc, [var], expr)
 
 
 buildToArrPRepr :: TyConRepr -> TyCon -> TyCon -> TyCon -> VM CoreExpr
+buildToArrPRepr repr@(ProdRepr {
+                   repr_prod_arg_tys      = prod_arg_tys
+                 , repr_prod_arr_tycon    = prod_arr_tycon
+                 , repr_prod_arr_data_con = prod_arr_data_con
+                 , repr_type              = repr_type
+                 })
+                vect_tc prepr_tc arr_tc
+  = do
+      arg_ty     <- mkPArrayType el_ty
+      shape_tys  <- arrShapeTys repr
+      arr_tys    <- arrReprTys repr
+      res_ty     <- mkPArrayType repr_type
+      rep_el_ty  <- mkPReprType el_ty
+
+      arg        <- newLocalVar FSLIT("xs") arg_ty
+      shape_vars <- mapM (newLocalVar FSLIT("sh")) shape_tys
+      rep_vars   <- mapM (newLocalVar FSLIT("ys")) arr_tys
+      
+      let vars = shape_vars ++ rep_vars
+
+      parray_co  <- mkBuiltinCo parrayTyCon
+
+      let res = wrapFamInstBody prod_arr_tycon prod_arg_tys
+              . mkConApp prod_arr_data_con
+              $ map Type prod_arg_tys ++ map Var vars
+
+          Just repr_co = tyConFamilyCoercion_maybe prepr_tc
+          co  = mkAppCoercion parray_co
+              . mkSymCoercion
+              $ mkTyConApp repr_co var_tys
+
+      return . Lam arg
+             . mkCoerce co
+             $ Case (unwrapFamInstScrut arr_tc var_tys (Var arg))
+                    (mkWildId (mkTyConApp arr_tc var_tys))
+                    res_ty
+               [(DataAlt arr_dc, vars, res)]
+  where
+    var_tys = mkTyVarTys $ tyConTyVars vect_tc
+    el_ty   = mkTyConApp vect_tc var_tys
+
+    [arr_dc] = tyConDataCons arr_tc
+
+
+buildToArrPRepr _ _ _ _ = return (Var unitDataConId)
+{-
 buildToArrPRepr _ vect_tc prepr_tc arr_tc
   = do
       arg_ty  <- mkPArrayType el_ty
@@ -331,29 +523,76 @@ buildToArrPRepr _ vect_tc prepr_tc arr_tc
 
     has_selector | [_] <- data_cons = False
                  | otherwise        = True
+-}
 
 buildFromArrPRepr :: TyConRepr -> TyCon -> TyCon -> TyCon -> VM CoreExpr
-buildFromArrPRepr _ vect_tc prepr_tc arr_tc
-  = mkFromArrPRepr undefined undefined undefined undefined undefined undefined
-
-buildPRDict :: TyConRepr -> TyCon -> TyCon -> TyCon -> VM CoreExpr
-buildPRDict (TyConRepr {
-                repr_tys         = repr_tys
-              , repr_prod_tycons = prod_tycons
-              , repr_prod_tys    = prod_tys
-              , repr_sum_tycon   = repr_sum_tycon
-              })
-            vect_tc prepr_tc _
+buildFromArrPRepr repr@(ProdRepr {
+                          repr_prod_arg_tys      = prod_arg_tys
+                        , repr_prod_arr_tycon    = prod_arr_tycon
+                        , repr_prod_arr_data_con = prod_arr_data_con
+                        , repr_type              = repr_type
+                        })
+                       vect_tc prepr_tc arr_tc
   = do
-      prs      <- mapM (mapM mkPR) repr_tys
-      prod_prs <- sequence $ zipWith3 mk_prod_pr prod_tycons repr_tys prs
-      sum_pr   <- mk_sum_pr prod_prs
-      prCoerce prepr_tc var_tys sum_pr
+      rep_el_ty  <- mkPReprType el_ty
+      arg_ty     <- mkPArrayType rep_el_ty
+      shape_tys  <- arrShapeTys repr
+      arr_tys    <- arrReprTys repr
+      res_ty     <- mkPArrayType el_ty
+
+      arg        <- newLocalVar FSLIT("xs") arg_ty
+      shape_vars <- mapM (newLocalVar FSLIT("sh")) shape_tys
+      rep_vars   <- mapM (newLocalVar FSLIT("ys")) arr_tys
+
+      let vars = shape_vars ++ rep_vars
+
+      parray_co  <- mkBuiltinCo parrayTyCon
+
+      let res = wrapFamInstBody arr_tc var_tys
+              . mkConApp arr_dc
+              $ map Type var_tys ++ map Var vars
+
+          Just repr_co = tyConFamilyCoercion_maybe prepr_tc
+          co  = mkAppCoercion parray_co
+              $ mkTyConApp repr_co var_tys
+
+          scrut = unwrapFamInstScrut prod_arr_tycon prod_arg_tys
+                $ mkCoerce co (Var arg)
+
+      return . Lam arg
+             $ Case (scrut)
+                    (mkWildId (mkTyConApp prod_arr_tycon prod_arg_tys))
+                    res_ty
+               [(DataAlt prod_arr_data_con, vars, res)]
   where
     var_tys = mkTyVarTys $ tyConTyVars vect_tc
+    el_ty   = mkTyConApp vect_tc var_tys
 
-    Just sum_tycon = repr_sum_tycon
+    [arr_dc] = tyConDataCons arr_tc
+buildFromArrPRepr _ _ _ _ = return (Var unitDataConId)
 
+buildPRDictRepr :: TyConRepr -> VM CoreExpr
+buildPRDictRepr (ProdRepr {
+                   repr_prod_arg_tys = prod_arg_tys
+                 , repr_prod_tycon   = prod_tycon
+                 })
+  = do
+      prs  <- mapM mkPR prod_arg_tys
+      dfun <- prDFunOfTyCon prod_tycon
+      return $ dfun `mkTyApps` prod_arg_tys `mkApps` prs
+
+buildPRDictRepr (SumRepr {
+                   repr_tys         = repr_tys
+                 , repr_prod_tycons = prod_tycons
+                 , repr_prod_tys    = prod_tys
+                 , repr_sum_tycon   = sum_tycon
+                 })
+  = do
+      prs      <- mapM (mapM mkPR) repr_tys
+      prod_prs <- sequence $ zipWith3 mk_prod_pr prod_tycons repr_tys prs
+      sum_dfun <- prDFunOfTyCon sum_tycon
+      return $ sum_dfun `mkTyApps` prod_tys `mkApps` prod_prs
+  where
     mk_prod_pr _         _   []   = prDFunOfTyCon unitTyCon
     mk_prod_pr _         _   [pr] = return pr
     mk_prod_pr (Just tc) tys prs
@@ -361,11 +600,21 @@ buildPRDict (TyConRepr {
           dfun <- prDFunOfTyCon tc
           return $ dfun `mkTyApps` tys `mkApps` prs
 
-    mk_sum_pr [pr] = return pr
-    mk_sum_pr prs
-      = do
-          dfun <- prDFunOfTyCon sum_tycon
-          return $ dfun `mkTyApps` prod_tys `mkApps` prs
+buildPRDict :: TyConRepr -> TyCon -> TyCon -> TyCon -> VM CoreExpr
+buildPRDict repr vect_tc prepr_tc _
+  = do
+      dict  <- buildPRDictRepr repr
+
+      pr_co <- mkBuiltinCo prTyCon
+      let co = mkAppCoercion pr_co
+             . mkSymCoercion
+             $ mkTyConApp arg_co var_tys
+
+      return $ mkCoerce co dict
+  where
+    var_tys = mkTyVarTys $ tyConTyVars vect_tc
+
+    Just arg_co = tyConFamilyCoercion_maybe prepr_tc
 
 buildPArrayTyCon :: TyCon -> TyCon -> VM TyCon
 buildPArrayTyCon orig_tc vect_tc = fixV $ \repr_tc ->
@@ -400,17 +649,20 @@ buildPArrayDataCon orig_name vect_tc repr_tc
       dc_name  <- cloneName mkPArrayDataConOcc orig_name
       repr     <- mkTyConRepr vect_tc
 
-      let all_tys = arr_shape_tys repr ++ concat (arr_repr_tys repr)
+      shape_tys <- arrShapeTys repr
+      repr_tys  <- arrReprTys  repr
+
+      let tys = shape_tys ++ repr_tys
 
       liftDs $ buildDataCon dc_name
                             False                  -- not infix
-                            (map (const NotMarkedStrict) all_tys)
+                            (map (const NotMarkedStrict) tys)
                             []                     -- no field labels
                             (tyConTyVars vect_tc)
                             []                     -- no existentials
                             []                     -- no eq spec
                             []                     -- no context
-                            all_tys
+                            tys
                             repr_tc
 
 mkPADFun :: TyCon -> VM Var