vectoriser: make PA PDatas instances for prod types work
authorBen Lippmeier <benl@ouroborus.net>
Tue, 15 Nov 2011 03:41:15 +0000 (14:41 +1100)
committerBen Lippmeier <benl@ouroborus.net>
Tue, 15 Nov 2011 03:41:15 +0000 (14:41 +1100)
compiler/vectorise/Vectorise/Builtins/Base.hs
compiler/vectorise/Vectorise/Builtins/Initialise.hs
compiler/vectorise/Vectorise/Generic/Description.hs
compiler/vectorise/Vectorise/Generic/PAMethods.hs
compiler/vectorise/Vectorise/Generic/PData.hs
compiler/vectorise/Vectorise/Type/Env.hs
compiler/vectorise/Vectorise/Utils/Base.hs

index 2d0931b..c5ef19f 100644 (file)
@@ -73,10 +73,7 @@ data Builtins
         { parrayTyCon          :: TyCon                     -- ^ PArray
         , parray_PrimTyCons    :: NameEnv TyCon             -- ^ PArray_Int# etc.
         , pdataTyCon           :: TyCon                     -- ^ PData
-
-        , pdatasTyCon          :: Maybe TyCon    
-        -- ^ PDatas. Not all lifted backends use 'PDatas', so it might not be defined.
-
+        , pdatasTyCon          :: TyCon                     -- ^ PDatas
         , prClass              :: Class                     -- ^ PR
         , prTyCon              :: TyCon                     -- ^ PR
         , preprTyCon           :: TyCon                     -- ^ PRepr
@@ -99,7 +96,9 @@ data Builtins
         , fromVoidVar          :: Var                       -- ^ fromVoid
         , sumTyCons            :: Array Int TyCon           -- ^ Sum2 .. Sum3
         , wrapTyCon            :: TyCon                     -- ^ Wrap
+        , wrapsTyCon           :: TyCon                     -- ^ Wraps
         , pvoidVar             :: Var                       -- ^ pvoid
+        , pvoidsVar            :: Var                       -- ^ pvoids
         , closureTyCon         :: TyCon                     -- ^ :->
         , closureVar           :: Var                       -- ^ closure
         , liftedClosureVar     :: Var                       -- ^ liftedClosure
index 329a7f5..a3cfe2e 100644 (file)
@@ -38,7 +38,7 @@ initBuiltins
           -- 'PData': type family mapping array element types to array representation types
           -- Not all backends use `PDatas`.
       ; pdataTyCon  <- externalTyCon (fsLit "PData")
-      ; pdatasTyCon <- externalTyCon_maybe (fsLit "PDatas")
+      ; pdatasTyCon <- externalTyCon (fsLit "PDatas")
 
           -- 'PR': class of basic array operators operating on 'PData' types
       ; prClass     <- externalClass (fsLit "PR")
@@ -90,6 +90,7 @@ initBuiltins
       ; let sumTyCons    = listArray (2, mAX_DPH_SUM) sum_tcs
       ; wrapTyCon        <- externalTyCon (fsLit "Wrap")
       ; pvoidVar         <- externalVar   (fsLit "pvoid")
+      ; pvoidsVar        <- externalVar   (fsLit "pvoids")
 
           -- Types and functions for closure conversion
       ; closureTyCon     <- externalTyCon (fsLit ":->")
@@ -141,6 +142,7 @@ initBuiltins
                , sumTyCons            = sumTyCons
                , wrapTyCon            = wrapTyCon
                , pvoidVar             = pvoidVar
+               , pvoidsVar            = pvoidsVar
                , closureTyCon         = closureTyCon
                , closureVar           = closureVar
                , liftedClosureVar     = liftedClosureVar
@@ -220,16 +222,6 @@ externalTyCon :: FastString -> DsM TyCon
 externalTyCon fs = dsLookupDPHRdrEnv (mkTcOccFS fs) >>= dsLookupTyCon
 
 
--- |Lookup a 'TyCon' in 'Data.Array.Parallel.Prim', given its name.
---  Return 'Nothing' if there isn't one.
-externalTyCon_maybe :: FastString -> DsM (Maybe TyCon)
-externalTyCon_maybe fs
- = do   mName   <- dsLookupDPHRdrEnv_maybe (mkTcOccFS fs)
-        case mName of
-         Nothing        -> return Nothing
-         Just name      -> liftM Just $ dsLookupTyCon name
-
-
 -- |Lookup some `Type` in 'Data.Array.Parallel.Prim', given its name.
 externalType :: FastString -> DsM Type
 externalType fs
index 99af58c..8a60d57 100644 (file)
@@ -44,8 +44,7 @@ data SumRepr
                , repr_psum_tc   :: TyCon
 
                -- | PDatas version of the sum tycon    (eg PDatasSum2)
-               --   Not all lifted backends use `PDatas`.
-               , repr_psums_tc  :: Maybe TyCon
+               , repr_psums_tc  :: TyCon
 
                -- | Type of the selector (eg Sel2)
                , repr_sel_ty    :: Type
@@ -84,7 +83,7 @@ data ProdRepr
 
                  -- | PDatas version of the product tycon  (eg PDatasTuple2s)
                  --   Not all lifted backends use `PDatas`.
-               , repr_ptups_tc :: Maybe TyCon
+               , repr_ptups_tc :: TyCon
 
                  -- | Types of each field.
                , repr_comp_tys :: [Type]
@@ -125,8 +124,8 @@ tyConRepr tc
            
            -- Get the 'PData' and 'PDatas' tycons for the sum.
            let sumapp   = mkTyConApp sum_tc tys
-           psum_tc      <- liftM fst         $ pdataReprTyCon sumapp
-           psums_tc     <- liftM (liftM fst) $ pdatasReprTyCon_maybe sumapp
+           psum_tc      <- liftM fst $ pdataReprTyCon  sumapp
+           psums_tc     <- liftM fst $ pdatasReprTyCon sumapp
            
            sel_ty       <- builtin (selTy arity)
            return $ Sum 
@@ -157,8 +156,8 @@ tyConRepr tc
 
            -- Get the 'PData' and 'PDatas' tycons for the product.
            let prodapp  = mkTyConApp tup_tc tys'
-           ptup_tc      <- liftM fst         $ pdataReprTyCon prodapp
-           ptups_tc     <- liftM (liftM fst) $ pdatasReprTyCon_maybe prodapp
+           ptup_tc      <- liftM fst $ pdataReprTyCon  prodapp
+           ptups_tc     <- liftM fst $ pdatasReprTyCon prodapp
            
            return $ Prod 
                   { repr_tup_tc   = tup_tc
index a5e999b..6330ddd 100644 (file)
@@ -1,3 +1,9 @@
+
+-- | Generate methods for the PA class.
+--
+--   TODO: there is a large amount of redundancy here between the 
+--   a, PData a, and PDatas a forms. See if we can factor some of this out.
+--
 module Vectorise.Generic.PAMethods
   ( buildPReprTyCon
   , buildPAScAndMethods 
@@ -20,7 +26,6 @@ import MkId
 import FastString
 import MonadUtils
 import Control.Monad
-import Data.Maybe
 
 
 buildPReprTyCon :: TyCon -> TyCon -> SumRepr -> VM TyCon
@@ -57,13 +62,10 @@ mk_fam_inst fam_tc arg_tc
 --      toArrPRepr    :: PData a          -> PData (PRepr a)
 --      fromArrPRepr  :: PData (PRepr a)  -> PData a
 --
---      toArrPReprs   :: PDatas a         -> PDatas (PRepr a)    (optional)
---      fromArrPReprs :: PDatas (PRepr a) -> PDatas a            (optional)
+--      toArrPReprs   :: PDatas a         -> PDatas (PRepr a)
+--      fromArrPReprs :: PDatas (PRepr a) -> PDatas a
 --   @
 --
---  Not all lifted backends use the 'toArrPReprs' and 'fromArrPReprs' methods, 
---  so we only generate these if the 'PDatas' type family is defined.
---
 type PAInstanceBuilder
         =  TyCon        -- ^ Vectorised TyCon 
         -> TyCon        -- ^ Representation TyCon
@@ -75,17 +77,13 @@ type PAInstanceBuilder
 
 buildPAScAndMethods :: VM [(String, PAInstanceBuilder)]
 buildPAScAndMethods
- = do   hasPDatas <- liftM isJust $ builtin pdatasTyCon
-        return 
-         $    [ ("PR",            buildPRDict)
-              , ("toPRepr",       buildToPRepr)
-              , ("fromPRepr",     buildFromPRepr)
-              , ("toArrPRepr",    buildToArrPRepr)
-              , ("fromArrPRepr",  buildFromArrPRepr)]
-         ++ (if hasPDatas then
-              [ ("toArrPReprs",   buildToArrPReprs)
-              , ("fromArrPReprs", buildFromArrPReprs)]
-              else [])             
+ = return [ ("PR",            buildPRDict)
+          , ("toPRepr",       buildToPRepr)
+          , ("fromPRepr",     buildFromPRepr)
+          , ("toArrPRepr",    buildToArrPRepr)
+          , ("fromArrPRepr",  buildFromArrPRepr)
+          , ("toArrPReprs",   buildToArrPReprs)
+          , ("fromArrPReprs", buildFromArrPReprs)]
 
 
 buildPRDict :: PAInstanceBuilder
@@ -246,60 +244,46 @@ buildToArrPRepr vect_tc prepr_tc pdata_tc _ r
              $ mkWildCase scrut (mkTyConApp pdata_tc ty_args) res_ty
                [(DataAlt pdata_dc, vars, mkCoerce co result)]
   where
-    ty_args = mkTyVarTys $ tyConTyVars vect_tc
-    el_ty   = mkTyConApp vect_tc ty_args
-
+    ty_args    = mkTyVarTys $ tyConTyVars vect_tc
+    el_ty      = mkTyConApp vect_tc ty_args
     [pdata_dc] = tyConDataCons pdata_tc
 
-
-    to_sum EmptySum 
-     = do pvoid <- builtin pvoidVar
-          return ([], Var pvoid)
-
-    to_sum (UnarySum r) = to_con r
-
-    to_sum (Sum { repr_psum_tc = psum_tc
-                , repr_sel_ty  = sel_ty
-                , repr_con_tys = tys
-                , repr_cons    = cons
-                })
-     = do (vars, exprs) <- mapAndUnzipM to_con cons
-          sel <- newLocalVar (fsLit "sel") sel_ty
-          return (sel : concat vars, mk_result (Var sel) exprs)
-      where
-        [psum_con] = tyConDataCons psum_tc
-        mk_result sel exprs = wrapFamInstBody psum_tc tys
-                            $ mkConApp psum_con
-                            $ map Type tys ++ (sel : exprs)
-
-    to_con (ConRepr _ r) = to_prod r
-
-    to_prod EmptyProd
-     = do pvoid <- builtin pvoidVar
-          return ([], Var pvoid)
-
-    to_prod (UnaryProd r)
-     = do pty  <- mkPDataType (compOrigType r)
-          var  <- newLocalVar (fsLit "x") pty
-          expr <- to_comp (Var var) r
-          return ([var], expr)
-
-    to_prod (Prod { repr_ptup_tc  = ptup_tc
-                  , repr_comp_tys = tys
-                  , repr_comps    = comps })
-     = do ptys <- mapM (mkPDataType . compOrigType) comps
-          vars <- newLocalVars (fsLit "x") ptys
-          es   <- zipWithM to_comp (map Var vars) comps
-          return (vars, mk_result es)
-      where
-        [ptup_con] = tyConDataCons ptup_tc
-        mk_result exprs = wrapFamInstBody ptup_tc tys
-                        $ mkConApp ptup_con
-                        $ map Type tys ++ exprs
-
-    to_comp expr (Keep _ _) = return expr
+    to_sum ss
+     = case ss of
+        EmptySum    -> builtin pvoidVar >>= \pvoid -> return ([], Var pvoid) 
+        UnarySum r  -> to_con r
+        Sum{}
+         -> do  let psum_tc     =  repr_psum_tc ss
+                let [psum_con]  =  tyConDataCons psum_tc
+                (vars, exprs)   <- mapAndUnzipM to_con (repr_cons ss)
+                sel             <- newLocalVar (fsLit "sel") (repr_sel_ty ss)
+                return ( sel : concat vars
+                       , wrapFamInstBody psum_tc (repr_con_tys ss)
+                         $ mkConApp psum_con 
+                         $ map Type (repr_con_tys ss) ++ (Var sel : exprs))
+
+    to_prod ss
+     = case ss of
+        EmptyProd    -> builtin pvoidVar >>= \pvoid -> return ([], Var pvoid)
+        UnaryProd r
+         -> do  pty  <- mkPDataType (compOrigType r)
+                var  <- newLocalVar (fsLit "x") pty
+                expr <- to_comp (Var var) r
+                return ([var], expr)
+        Prod{}
+         -> do  let [ptup_con]  = tyConDataCons (repr_ptup_tc ss)
+                ptys   <- mapM (mkPDataType . compOrigType) (repr_comps ss)
+                vars   <- newLocalVars (fsLit "x") ptys
+                exprs  <- zipWithM to_comp (map Var vars) (repr_comps ss)
+                return ( vars
+                       , wrapFamInstBody (repr_ptup_tc ss) (repr_comp_tys ss)
+                         $ mkConApp ptup_con
+                         $ map Type (repr_comp_tys ss) ++ exprs)
+
+    to_con (ConRepr _ r)    = to_prod r
 
     -- FIXME: this is bound to be wrong!
+    to_comp expr (Keep _ _) = return expr
     to_comp expr (Wrap ty)
       = do
           wrap_tc  <- builtin wrapTyCon
@@ -317,62 +301,58 @@ buildFromArrPRepr vect_tc prepr_tc pdata_tc _ r
 
       pdata_co <- mkBuiltinCo pdataTyCon
       let Just repr_co = tyConFamilyCoercion_maybe prepr_tc
-          co           = mkAppCo pdata_co
+      let co           = mkAppCo pdata_co
                        $ mkAxInstCo repr_co var_tys
 
-          scrut  = mkCoerce co (Var arg)
+      let scrut        = mkCoerce co (Var arg)
 
-          mk_result args = wrapFamInstBody pdata_tc var_tys
-                         $ mkConApp pdata_con
-                         $ map Type var_tys ++ args
+      let mk_result args
+            = wrapFamInstBody pdata_tc var_tys
+            $ mkConApp pdata_con
+            $ map Type var_tys ++ args
 
       (expr, _) <- fixV $ \ ~(_, args) ->
                      from_sum res_ty (mk_result args) scrut r
 
       return $ Lam arg expr
  where
-    var_tys = mkTyVarTys $ tyConTyVars vect_tc
-    el_ty   = mkTyConApp vect_tc var_tys
-
+    var_tys     = mkTyVarTys $ tyConTyVars vect_tc
+    el_ty       = mkTyConApp vect_tc var_tys
     [pdata_con] = tyConDataCons pdata_tc
 
-    from_sum _ res _ EmptySum = return (res, [])
-    from_sum res_ty res expr (UnarySum r) = from_con res_ty res expr r
-    from_sum res_ty res expr (Sum { repr_psum_tc = psum_tc
-                                  , repr_sel_ty  = sel_ty
-                                  , repr_con_tys = tys
-                                  , repr_cons    = cons })
-     = do sel  <- newLocalVar (fsLit "sel") sel_ty
-          ptys <- mapM mkPDataType tys
-          vars <- newLocalVars (fsLit "xs") ptys
-          (res', args) <- fold from_con res_ty res (map Var vars) cons
-          let scrut = unwrapFamInstScrut psum_tc tys expr
-              body  = mkWildCase scrut (exprType scrut) res_ty
-                      [(DataAlt psum_con, sel : vars, res')]
-          return (body, Var sel : args)
-      where
-        [psum_con] = tyConDataCons psum_tc
-
+    from_sum res_ty res expr ss
+     = case ss of
+        EmptySum    -> return (res, [])
+        UnarySum r  -> from_con res_ty res expr r
+        Sum {}
+         -> do  let psum_tc    =  repr_psum_tc ss
+                let [psum_con] =  tyConDataCons psum_tc
+                sel            <- newLocalVar (fsLit "sel") (repr_sel_ty ss)
+                ptys           <- mapM mkPDataType (repr_con_tys ss)
+                vars           <- newLocalVars (fsLit "xs") ptys
+                (res', args)   <- fold from_con res_ty res (map Var vars) (repr_cons ss)
+                let scrut      =  unwrapFamInstScrut psum_tc (repr_con_tys ss) expr
+                let body       =  mkWildCase scrut (exprType scrut) res_ty
+                                    [(DataAlt psum_con, sel : vars, res')]
+                return (body, Var sel : args)
+
+    from_prod res_ty res expr ss
+     = case ss of
+        EmptyProd   -> return (res, [])
+        UnaryProd r -> from_comp res_ty res expr r
+        Prod {}
+         -> do  let ptup_tc    =  repr_ptup_tc ss
+                let [ptup_con] =  tyConDataCons ptup_tc
+                ptys           <- mapM mkPDataType (repr_comp_tys ss)
+                vars           <- newLocalVars (fsLit "ys") ptys
+                (res', args)   <- fold from_comp res_ty res (map Var vars) (repr_comps ss)
+                let scrut      =  unwrapFamInstScrut ptup_tc (repr_comp_tys ss) expr
+                let body       =  mkWildCase scrut (exprType scrut) res_ty
+                                    [(DataAlt ptup_con, vars, res')]
+                return (body, args)      
 
     from_con res_ty res expr (ConRepr _ r) = from_prod res_ty res expr r
 
-    from_prod _ res _ EmptyProd
-      = return (res, [])
-    from_prod res_ty res expr (UnaryProd r)
-      = from_comp res_ty res expr r
-    from_prod res_ty res expr (Prod { repr_ptup_tc  = ptup_tc
-                                    , repr_comp_tys = tys
-                                    , repr_comps    = comps })
-     = do ptys <- mapM mkPDataType tys
-          vars <- newLocalVars (fsLit "ys") ptys
-          (res', args) <- fold from_comp res_ty res (map Var vars) comps
-          let scrut = unwrapFamInstScrut ptup_tc tys expr
-              body  = mkWildCase scrut (exprType scrut) res_ty
-                      [(DataAlt ptup_con, vars, res')]
-          return (body, args)
-      where
-        [ptup_con] = tyConDataCons ptup_tc
-
     from_comp _ res expr (Keep _ _) = return (res, [expr])
     from_comp _ res expr (Wrap ty)
      = do wrap_tc       <- builtin wrapTyCon
@@ -380,7 +360,8 @@ buildFromArrPRepr vect_tc prepr_tc pdata_tc _ r
           return (res, [unwrapNewTypeBody pwrap_tc [ty]
                         $ unwrapFamInstScrut pwrap_tc [ty] expr])
 
-    fold f res_ty res exprs rs = foldrM f' (res, []) (zip exprs rs)
+    fold f res_ty res exprs rs
+      = foldrM f' (res, []) (zip exprs rs)
       where
         f' (expr, r) (res, args) 
          = do (res', args') <- f res_ty res expr r
@@ -393,47 +374,171 @@ buildFromArrPRepr vect_tc prepr_tc pdata_tc _ r
 buildToArrPReprs :: PAInstanceBuilder
 buildToArrPReprs vect_tc prepr_tc _ pdatas_tc r
  = do
-        -- The element type of the argument.
-        --  eg: 'Tree a b'.
-        let ty_args = mkTyVarTys $ tyConTyVars vect_tc
-        let el_ty   = mkTyConApp vect_tc ty_args
-
-        -- The argument type of the instance.
-        --  eg: 'PDatas (Tree a b)'
-        Just arg_ty      <- mkPDatasType el_ty
-
-        -- The result type. 
-        --  eg: 'PDatas (PRepr (Tree a b))'
-        Just res_ty      <- mkPDatasType =<< mkPReprType el_ty
-        
-        -- Variable to bind the argument to the instance
-        -- eg: (xss :: PDatas (Tree a b))
-        varg         <- newLocalVar (fsLit "xss") arg_ty
-        
-        return  $ Lam varg (Var varg)
+    -- The argument type of the instance.
+    --  eg: 'PDatas (Tree a b)'
+    arg_ty    <- mkPDatasType el_ty
+
+    -- The result type. 
+    --  eg: 'PDatas (PRepr (Tree a b))'
+    res_ty    <- mkPDatasType =<< mkPReprType el_ty
+
+    -- Variable to bind the argument to the instance
+    -- eg: (xss :: PDatas (Tree a b))
+    varg      <- newLocalVar (fsLit "xss") arg_ty
+
+    -- Coersion to case between the (PRepr a) type and its instance.
+    pdatas_co <- mkBuiltinCo pdatasTyCon
+    let Just repr_co = tyConFamilyCoercion_maybe prepr_tc
+    let co           = mkAppCo pdatas_co
+                     . mkSymCo
+                     $ mkAxInstCo repr_co ty_args
 
+    let scrut        = unwrapFamInstScrut pdatas_tc ty_args (Var varg)
+    (vars, result)  <- to_sum r
+
+    return  $ Lam varg
+            $ mkWildCase scrut (mkTyConApp pdatas_tc ty_args) res_ty
+                    [(DataAlt pdatas_dc, vars, mkCoerce co result)]        
+
+ where
+    -- The element type of the argument.
+    --  eg: 'Tree a b'.
+    ty_args = mkTyVarTys $ tyConTyVars vect_tc
+    el_ty   = mkTyConApp vect_tc ty_args
         
+    -- PDatas data constructor
+    [pdatas_dc] = tyConDataCons pdatas_tc
+         
+    to_sum ss
+     = case ss of       -- BROKEN: should be
+        EmptySum    -> builtin pvoidVar >>= \pvoid -> return ([], Var pvoid) 
+        UnarySum r  -> to_con r
+        Sum{}
+          -> do let psums_tc     = repr_psums_tc ss
+                let [psums_con]  = tyConDataCons psums_tc
+                (vars, exprs)   <- mapAndUnzipM to_con (repr_cons ss)
+                sel             <- newLocalVar (fsLit "sel") (repr_sel_ty ss) -- BROKEN: should be vector
+                return ( sel : concat vars
+                       , wrapFamInstBody psums_tc (repr_con_tys ss)
+                         $ mkConApp psums_con 
+                         $ map Type (repr_con_tys ss) ++ (Var sel : exprs))        
+
+    to_prod ss
+     = case ss of       -- BROKEN: should be pvoids
+        EmptyProd    -> builtin pvoidVar >>= \pvoid -> return ([], Var pvoid)
+        UnaryProd r
+         -> do  pty  <- mkPDatasType (compOrigType r)
+                var  <- newLocalVar (fsLit "x") pty
+                expr <- to_comp (Var var) r
+                return ([var], expr)
+        Prod{}
+         -> do  let [ptups_con]  = tyConDataCons (repr_ptups_tc ss)
+                ptys   <- mapM (mkPDatasType . compOrigType) (repr_comps ss)
+                vars   <- newLocalVars (fsLit "x") ptys
+                exprs  <- zipWithM to_comp (map Var vars) (repr_comps ss)
+                return ( vars
+                       , wrapFamInstBody (repr_ptups_tc ss) (repr_comp_tys ss)
+                         $ mkConApp ptups_con
+                         $ map Type (repr_comp_tys ss) ++ exprs)
+
+    to_con (ConRepr _ r)    = to_prod r
+
+    -- FIXME: this is bound to be wrong!
+    to_comp expr (Keep _ _) = return expr
+    to_comp expr (Wrap ty)
+     = do wrap_tc       <- builtin wrapTyCon
+          (pwrap_tc, _) <- pdatasReprTyCon (mkTyConApp wrap_tc [ty])
+          return $ wrapNewTypeBody pwrap_tc [ty] expr
+
+
 -- buildFromArrPReprs ---------------------------------------------------------
 buildFromArrPReprs :: PAInstanceBuilder
 buildFromArrPReprs vect_tc prepr_tc _ pdatas_tc r
  = do   
-        -- The element type of the argument.
-        --  eg: 'Tree a b'.
-        let ty_args = mkTyVarTys $ tyConTyVars vect_tc
-        let el_ty   = mkTyConApp vect_tc ty_args
-
-        -- The argument type of the instance.
-        --  eg: 'PDatas (PRepr (Tree a b))'
-        Just arg_ty      <- mkPDatasType =<< mkPReprType el_ty
-
-        -- The result type. 
-        --  eg: 'PDatas (Tree a b)'
-        Just res_ty      <- mkPDatasType el_ty
+    -- The element type of the argument.
+    --  eg: 'Tree a b'.
+    let ty_args = mkTyVarTys $ tyConTyVars vect_tc
+    let el_ty   = mkTyConApp vect_tc ty_args
+
+    -- The argument type of the instance.
+    --  eg: 'PDatas (PRepr (Tree a b))'
+    arg_ty      <- mkPDatasType =<< mkPReprType el_ty
+
+    -- The result type. 
+    --  eg: 'PDatas (Tree a b)'
+    res_ty      <- mkPDatasType el_ty
         
-        -- Variable to bind the argument to the instance
-        -- eg: (xss :: PDatas (PRepr (Tree a b)))
-        varg         <- newLocalVar (fsLit "xss") arg_ty
+    -- Variable to bind the argument to the instance
+    -- eg: (xss :: PDatas (PRepr (Tree a b)))
+    varg        <- newLocalVar (fsLit "xss") arg_ty
         
-        return  $ Lam varg (Var varg)
+    -- Build the coersion between PRepr and the instance type
+    pdatas_co <- mkBuiltinCo pdatasTyCon
+    let Just repr_co = tyConFamilyCoercion_maybe prepr_tc
+    let co           = mkAppCo pdatas_co
+                     $ mkAxInstCo repr_co var_tys
+
+    let scrut        = mkCoerce co (Var varg)
+
+    let mk_result args
+            = wrapFamInstBody pdatas_tc var_tys
+            $ mkConApp pdatas_con
+            $ map Type var_tys ++ args
+
+    (expr, _) <- fixV $ \ ~(_, args) ->
+                     from_sum res_ty (mk_result args) scrut r
+
+    return $ Lam varg expr
+ where
+    var_tys      = mkTyVarTys $ tyConTyVars vect_tc
+    el_ty        = mkTyConApp vect_tc var_tys
+    [pdatas_con] = tyConDataCons pdatas_tc
+
+    from_sum res_ty res expr ss
+     = case ss of
+        EmptySum    -> return (res, [])
+        UnarySum r  -> from_con res_ty res expr r
+        Sum {}
+         -> do  let psums_tc    =  repr_psums_tc ss
+                let [psums_con] =  tyConDataCons psums_tc
+                sel             <- newLocalVar (fsLit "sel") (repr_sel_ty ss)
+                ptys            <- mapM mkPDatasType (repr_con_tys ss)
+                vars            <- newLocalVars (fsLit "xs") ptys
+                (res', args)    <- fold from_con res_ty res (map Var vars) (repr_cons ss)
+                let scrut       =  unwrapFamInstScrut psums_tc (repr_con_tys ss) expr
+                let body        =  mkWildCase scrut (exprType scrut) res_ty
+                                    [(DataAlt psums_con, sel : vars, res')]
+                return (body, Var sel : args)
+
+    from_prod res_ty res expr ss
+     = case ss of
+        EmptyProd   -> return (res, [])
+        UnaryProd r -> from_comp res_ty res expr r
+        Prod {}
+         -> do  let ptups_tc    =  repr_ptups_tc ss
+                let [ptups_con] =  tyConDataCons ptups_tc
+                ptys            <- mapM mkPDatasType (repr_comp_tys ss)
+                vars            <- newLocalVars (fsLit "ys") ptys
+                (res', args)    <- fold from_comp res_ty res (map Var vars) (repr_comps ss)
+                let scrut       =  unwrapFamInstScrut ptups_tc (repr_comp_tys ss) expr
+                let body        =  mkWildCase scrut (exprType scrut) res_ty
+                                    [(DataAlt ptups_con, vars, res')]
+                return (body, args)      
+
+    from_con res_ty res expr (ConRepr _ r)
+        = from_prod res_ty res expr r
 
+    from_comp _ res expr (Keep _ _) = return (res, [expr])
+    from_comp _ res expr (Wrap ty)
+     = do wrap_tc        <- builtin wrapTyCon
+          (pwraps_tc, _) <- pdatasReprTyCon (mkTyConApp wrap_tc [ty])
+          return (res, [unwrapNewTypeBody pwraps_tc [ty]
+                        $ unwrapFamInstScrut pwraps_tc [ty] expr])
+
+    fold f res_ty res exprs rs
+      = foldrM f' (res, []) (zip exprs rs)
+      where
+        f' (expr, r) (res, args) 
+         = do (res', args') <- f res_ty res expr r
+              return (res', args' ++ args)
 
index d9e4ff0..f10afff 100644 (file)
@@ -79,7 +79,7 @@ buildPDatasTyCon orig_tc vect_tc repr
  = fixV $ \repr_tc ->
  do name'       <- mkLocalisedName mkPDatasTyConOcc orig_name
     rhs         <- buildPDatasTyConRhs orig_name vect_tc repr_tc repr
-    Just pdatas <- builtin pdatasTyCon
+    pdatas      <- builtin pdatasTyCon
 
     liftDs $ buildAlgTyCon name'
                            tyvars
@@ -106,10 +106,7 @@ buildPDatasDataCon orig_name vect_tc repr_tc repr
  = do let tvs   = tyConTyVars vect_tc
       dc_name        <- mkLocalisedName mkPDatasDataConOcc orig_name
 
-      let mkPDatasType' t
-           = mkPDatasType t >>= (\(Just t') -> return t')
-
-      comp_tys  <- mkSumTys mkPDatasType' repr
+      comp_tys  <- mkSumTys mkPDatasType repr
 
       liftDs $ buildDataCon dc_name
                             False                  -- not infix
index 79f5712..ee0ae25 100644 (file)
@@ -39,7 +39,7 @@ import FastString
 import MonadUtils
 import Control.Monad
 import Data.List
-
+import Data.Maybe
 
 -- Note [Pragmas to vectorise tycons]
 -- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -208,6 +208,7 @@ vectTypeEnv tycons vectTypeDecls vectClassDecls
        ; repr_tcs   <- zipWith3M buildPReprTyCon  orig_tcs vect_tcs reprs
        ; pdata_tcs  <- zipWith3M buildPDataTyCon  orig_tcs vect_tcs reprs
        ; pdatas_tcs <- zipWith3M buildPDatasTyCon orig_tcs vect_tcs reprs
+
        ; let inst_tcs  = repr_tcs ++ pdata_tcs
              fam_insts = map mkLocalFamInst inst_tcs
        ; updGEnv $ extendFamEnv fam_insts
index 6472253..c4dfe5c 100644 (file)
@@ -15,9 +15,8 @@ module Vectorise.Utils.Base (
   mkBuiltinCo,
   mkVScrut,
 
-  pdataReprTyCon,
-  pdatasReprTyCon_maybe,
-  pdataReprDataCon,
+  pdataReprTyCon,   pdatasReprTyCon,
+  pdataReprDataCon, pdatasReprDataCon,
   prDFunOfTyCon
 ) where
 
@@ -37,7 +36,7 @@ import Outputable
 import FastString
 
 import Control.Monad (liftM)
-import Data.Maybe
+
 
 -- Simple Types ---------------------------------------------------------------
 voidType :: VM Type
@@ -110,17 +109,12 @@ mkPArrayType ty = mkBuiltinTyConApp parrayTyCon [ty]
 
 -- | Make an appliction of the 'PData' tycon to some argument.
 mkPDataType :: Type -> VM Type
-mkPDataType ty
-        = mkBuiltinTyConApp pdataTyCon [ty]
+mkPDataType ty  = mkBuiltinTyConApp pdataTyCon [ty]
 
 
 -- | Make an application of the 'PDatas' tycon to some argument.
-mkPDatasType :: Type -> VM (Maybe Type)
-mkPDatasType ty
- = do   mtc      <- builtin pdatasTyCon
-        case mtc of
-         Nothing        -> return Nothing
-         Just tc'       -> return $ Just $ mkTyConApp tc' [ty]
+mkPDatasType :: Type -> VM Type
+mkPDatasType ty = mkBuiltinTyConApp pdatasTyCon [ty]
 
 
 -- |Checks if a type constructor is defined in 'GHC.Prim' (e.g., 'Int#'); if so, returns it.
@@ -175,23 +169,22 @@ pdataReprTyCon :: Type -> VM (TyCon, [Type])
 pdataReprTyCon ty
         = builtin pdataTyCon >>= (`lookupFamInst` [ty])
 
+pdataReprDataCon :: Type -> VM (DataCon, [Type])
+pdataReprDataCon ty
+ = do   (tc, arg_tys) <- pdataReprTyCon ty
+        let [dc] = tyConDataCons tc
+        return (dc, arg_tys)
 
--- | Get the PDatas tycon that represents this type, if there is one.
---   Not all backends use 'PDatas', so there might not be one.
-pdatasReprTyCon_maybe :: Type -> VM (Maybe (TyCon, [Type]))
-pdatasReprTyCon_maybe ty
- = do   mtc     <- builtin pdatasTyCon
-        case mtc of
-         Nothing        -> return Nothing
-         Just tc        -> liftM Just $ lookupFamInst tc [ty]
+pdatasReprTyCon :: Type -> VM (TyCon, [Type])
+pdatasReprTyCon ty
+        = builtin pdatasTyCon >>= (`lookupFamInst` [ty])
 
+pdatasReprDataCon :: Type -> VM (DataCon, [Type])
+pdatasReprDataCon ty
+ = do   (tc, arg_tys) <- pdatasReprTyCon ty
+        let [dc] = tyConDataCons tc
+        return (dc, arg_tys)
 
-pdataReprDataCon :: Type -> VM (DataCon, [Type])
-pdataReprDataCon ty
-  = do
-      (tc, arg_tys) <- pdataReprTyCon ty
-      let [dc] = tyConDataCons tc
-      return (dc, arg_tys)
 
 prDFunOfTyCon :: TyCon -> VM CoreExpr
 prDFunOfTyCon tycon