Complete PA dictionary generation for product types
[ghc-hetmet.git] / compiler / vectorise / VectType.hs
index 17a2b44..c977e4c 100644 (file)
@@ -26,7 +26,7 @@ import Var               ( Var )
 import Id                ( mkWildId )
 import Name              ( Name, getOccName )
 import NameEnv
-import TysWiredIn        ( unitTy, intTy, intDataCon )
+import TysWiredIn        ( unitTy, unitTyCon, intTy, intDataCon, unitDataConId )
 import TysPrim           ( intPrimTy )
 
 import Unique
@@ -208,29 +208,287 @@ 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 repr vect_tc prepr_tc _
+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
-      arg <- newLocalVar FSLIT("x") arg_ty
-      bndrss <- mapM (mapM (newLocalVar FSLIT("x")))
-                     (repr_tys repr)
+      arg  <- newLocalVar FSLIT("x") arg_ty
+      vars <- mapM (mapM (newLocalVar FSLIT("x"))) repr_tys
 
       return . Lam arg
              . wrapFamInstBody prepr_tc var_tys
-             . Case (Var arg) (mkWildId arg_ty) (repr_type repr)
-             . zipWith3 mk_alt data_cons bndrss
-             . mkToPRepr repr $ map (map Var) bndrss
+             . Case (Var arg) (mkWildId arg_ty) repr_type
+             . 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_alt data_con bndrs body = (DataAlt data_con, bndrs, body)
+    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 dc) tys exprs  = mkConApp dc (map Type tys ++ exprs)
+
+buildFromPRepr :: TyConRepr -> TyCon -> TyCon -> TyCon -> VM CoreExpr
+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
+      arg_ty <- mkPReprType res_ty
+      arg    <- newLocalVar FSLIT("x") arg_ty
+
+      liftM (Lam 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
+    var_tys   = mkTyVarTys $ tyConTyVars vect_tc
+    ty_args   = map Type var_tys
+    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
+          return (var, mkConApp dc ty_args)
+    un_prod dc _ _ [ty]
+      = do
+          var <- newLocalVar FSLIT("x") ty
+          return (var, mkConApp dc (ty_args ++ [Var var]))
+
+    un_prod dc (Just prod_dc) prod_ty tys
+      = do
+          vars  <- mapM (newLocalVar FSLIT("x")) tys
+          pv    <- newLocalVar FSLIT("p") prod_ty
+
+          let res  = mkConApp dc (ty_args ++ map Var vars)
+              expr = Case (Var pv) (mkWildId prod_ty) res_ty
+                        [(DataAlt prod_dc, vars, res)]
+
+          return (pv, expr)
+
+    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
@@ -265,37 +523,99 @@ buildToArrPRepr _ vect_tc prepr_tc arr_tc
 
     has_selector | [_] <- data_cons = False
                  | otherwise        = True
+-}
 
-
-buildFromPRepr :: TyConRepr -> TyCon -> TyCon -> TyCon -> VM CoreExpr
-buildFromPRepr _ vect_tc prepr_tc _
+buildFromArrPRepr :: TyConRepr -> TyCon -> TyCon -> TyCon -> VM CoreExpr
+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
-      arg_ty <- mkPReprType res_ty
-      arg <- newLocalVar FSLIT("x") arg_ty
-      alts <- mapM mk_alt data_cons
-      body <- mkFromPRepr (unwrapFamInstScrut prepr_tc var_tys (Var arg))
-                          res_ty alts
-      return $ Lam arg body
+      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
-    res_ty    = mkTyConApp vect_tc var_tys
-    data_cons = tyConDataCons vect_tc
+    var_tys = mkTyVarTys $ tyConTyVars vect_tc
+    el_ty   = mkTyConApp vect_tc var_tys
 
-    mk_alt dc = do
-                  bndrs <- mapM (newLocalVar FSLIT("x")) $ dataConRepArgTys dc
-                  return (bndrs, mkConApp dc (map Type var_tys ++ map Var bndrs))
+    [arr_dc] = tyConDataCons arr_tc
+buildFromArrPRepr _ _ _ _ = return (Var unitDataConId)
 
-buildFromArrPRepr :: TyConRepr -> TyCon -> TyCon -> TyCon -> VM CoreExpr
-buildFromArrPRepr _ vect_tc prepr_tc arr_tc
-  = mkFromArrPRepr undefined undefined undefined undefined undefined undefined
+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
@@ -327,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