Complete PA dictionary generation for product types
[ghc-hetmet.git] / compiler / vectorise / VectType.hs
index e505cd3..c977e4c 100644 (file)
@@ -26,7 +26,7 @@ import Var               ( Var )
 import Id                ( mkWildId )
 import Name              ( Name, getOccName )
 import NameEnv
-import TysWiredIn        ( unitTy, intTy, intDataCon, unitDataConId )
+import TysWiredIn        ( unitTy, unitTyCon, intTy, intDataCon, unitDataConId )
 import TysPrim           ( intPrimTy )
 
 import Unique
@@ -208,16 +208,141 @@ 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 {
-                repr_tys         = repr_tys
-              , repr_prod_tycons = prod_tycons
-              , repr_prod_tys    = prod_tys
-              , repr_sum_tycon   = repr_sum_tycon
-              , repr_type        = repr_type
+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_tycon      = sum_tycon
+              , repr_type           = repr_type
               })
               vect_tc prepr_tc _
   = do
@@ -227,36 +352,52 @@ buildToPRepr (TyConRepr {
       return . Lam arg
              . wrapFamInstBody prepr_tc var_tys
              . Case (Var arg) (mkWildId arg_ty) repr_type
-             . mk_alts data_cons vars
-             . zipWith3 mk_prod prod_tycons repr_tys $ map (map Var) 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
 
-    Just sum_tycon = repr_sum_tycon
-    sum_data_cons  = tyConDataCons sum_tycon
-
-    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]))
 
     mk_prod _         _   []     = Var unitDataConId
     mk_prod _         _   [expr] = expr
-    mk_prod (Just tc) tys exprs  = mkConApp dc (map Type tys ++ exprs)
-      where
-        [dc] = tyConDataCons tc
+    mk_prod (Just dc) tys exprs  = mkConApp dc (map Type tys ++ exprs)
 
 buildFromPRepr :: TyConRepr -> TyCon -> TyCon -> TyCon -> VM CoreExpr
-buildFromPRepr (TyConRepr {
-                repr_tys         = repr_tys
-              , repr_prod_tycons = prod_tycons
-              , repr_prod_tys    = prod_tys
-              , repr_sum_tycon   = repr_sum_tycon
-              , repr_type        = repr_type
+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_tycon      = sum_tycon
+              , repr_type           = repr_type
               })
               vect_tc prepr_tc _
   = do
@@ -264,16 +405,19 @@ buildFromPRepr (TyConRepr {
       arg    <- newLocalVar FSLIT("x") arg_ty
 
       liftM (Lam arg
-             . un_sum (unwrapFamInstScrut prepr_tc var_tys (Var arg)))
-            (sequence $ zipWith4 un_prod data_cons prod_tycons prod_tys repr_tys)
+             . 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
     var_tys   = mkTyVarTys $ tyConTyVars vect_tc
     ty_args   = map Type var_tys
     res_ty    = mkTyConApp vect_tc var_tys
     data_cons = tyConDataCons vect_tc
 
-    Just sum_tc   = repr_sum_tycon
-    sum_data_cons = tyConDataCons sum_tc
+    sum_data_cons = tyConDataCons sum_tycon
 
     un_prod dc _ _ []
       = do
@@ -284,7 +428,7 @@ buildFromPRepr (TyConRepr {
           var <- newLocalVar FSLIT("x") ty
           return (var, mkConApp dc (ty_args ++ [Var var]))
 
-    un_prod dc (Just prod_tc) prod_ty tys
+    un_prod dc (Just prod_dc) prod_ty tys
       = do
           vars  <- mapM (newLocalVar FSLIT("x")) tys
           pv    <- newLocalVar FSLIT("p") prod_ty
@@ -294,19 +438,57 @@ buildFromPRepr (TyConRepr {
                         [(DataAlt prod_dc, vars, res)]
 
           return (pv, expr)
-      where
-        [prod_dc] = tyConDataCons prod_tc
-
-
-    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
@@ -341,18 +523,99 @@ 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
+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
+      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
+
+    [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
+      = do
+          dfun <- prDFunOfTyCon tc
+          return $ dfun `mkTyApps` tys `mkApps` prs
 
 buildPRDict :: TyConRepr -> TyCon -> TyCon -> TyCon -> VM CoreExpr
-buildPRDict _ vect_tc prepr_tc _
-  = prCoerce prepr_tc var_tys
-  =<< prDictOfType (mkTyConApp prepr_tc var_tys)
+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 ->
   do
@@ -384,22 +647,23 @@ buildPArrayDataCon :: Name -> TyCon -> TyCon -> VM DataCon
 buildPArrayDataCon orig_name vect_tc repr_tc
   = do
       dc_name  <- cloneName mkPArrayDataConOcc orig_name
-      shape    <- tyConShape vect_tc
-      repr_tys <- mapM mkPArrayType types
+      repr     <- mkTyConRepr vect_tc
+
+      shape_tys <- arrShapeTys repr
+      repr_tys  <- arrReprTys  repr
+
+      let tys = shape_tys ++ repr_tys
 
       liftDs $ buildDataCon dc_name
                             False                  -- not infix
-                            (shapeStrictness shape ++ map (const NotMarkedStrict) repr_tys)
+                            (map (const NotMarkedStrict) tys)
                             []                     -- no field labels
                             (tyConTyVars vect_tc)
                             []                     -- no existentials
                             []                     -- no eq spec
                             []                     -- no context
-                            (shapeReprTys shape ++ repr_tys)
+                            tys
                             repr_tc
-  where
-    types = [ty | dc <- tyConDataCons vect_tc
-                , ty <- dataConRepArgTys dc]
 
 mkPADFun :: TyCon -> VM Var
 mkPADFun vect_tc