Fix vectorisation of sum type constructors

compiler/prelude/PrelNames.lhs

@@ -223,6 +223,7 @@ ndpNames = [ parrayTyConName, paTyConName, preprTyConName, prTyConName
            , closureTyConName
            , mkClosureName, applyClosureName
            , mkClosurePName, applyClosurePName
+           , replicatePAIntPrimName, upToPAIntPrimName
            , lengthPAName, replicatePAName, emptyPAName, packPAName,
              combinePAName ]
 \end{code}
@@ -702,6 +703,9 @@ prTyConName         = tcQual   nDP_PARRAY FSLIT("PR")     prTyConKey
 parrayIntPrimTyConName = tcQual nDP_PRIM  FSLIT("PArray_Int#")
                                                           parrayIntPrimTyConKey
 mkPRName            = varQual  nDP_PARRAY FSLIT("mkPR")   mkPRIdKey
+replicatePAIntPrimName = varQual nDP_PRIM FSLIT("replicatePA_Int#")
+                                                        replicatePAIntPrimIdKey
+upToPAIntPrimName   = varQual  nDP_PRIM   FSLIT("upToPA_Int#") upToPAIntPrimIdKey
 lengthPAName        = varQual  nDP_PARRAY FSLIT("lengthPA")    lengthPAIdKey
 replicatePAName     = varQual  nDP_PARRAY FSLIT("replicatePA") replicatePAIdKey
 emptyPAName         = varQual  nDP_PARRAY FSLIT("emptyPA") emptyPAIdKey
@@ -1090,6 +1094,8 @@ emptyPAIdKey                  = mkPreludeMiscIdUnique 133
 packPAIdKey                   = mkPreludeMiscIdUnique 134
 combinePAIdKey                = mkPreludeMiscIdUnique 135
 mkPRIdKey                     = mkPreludeMiscIdUnique 136
+replicatePAIntPrimIdKey       = mkPreludeMiscIdUnique 137
+upToPAIntPrimIdKey            = mkPreludeMiscIdUnique 138
 
 ---------------- Template Haskell -------------------
 --	USES IdUniques 200-399

compiler/vectorise/VectBuiltIn.hs

@@ -53,6 +53,8 @@ data Builtins = Builtins {
                 , applyClosureVar  :: Var
                 , mkClosurePVar    :: Var
                 , applyClosurePVar :: Var
+                , replicatePAIntPrimVar :: Var
+                , upToPAIntPrimVar :: Var
                 , lengthPAVar      :: Var
                 , replicatePAVar   :: Var
                 , emptyPAVar       :: Var
@@ -93,6 +95,8 @@ initBuiltins
       applyClosureVar  <- dsLookupGlobalId applyClosureName
       mkClosurePVar    <- dsLookupGlobalId mkClosurePName
       applyClosurePVar <- dsLookupGlobalId applyClosurePName
+      replicatePAIntPrimVar <- dsLookupGlobalId replicatePAIntPrimName
+      upToPAIntPrimVar <- dsLookupGlobalId upToPAIntPrimName
       lengthPAVar      <- dsLookupGlobalId lengthPAName
       replicatePAVar   <- dsLookupGlobalId replicatePAName
       emptyPAVar       <- dsLookupGlobalId emptyPAName
@@ -117,6 +121,8 @@ initBuiltins
                , applyClosureVar  = applyClosureVar
                , mkClosurePVar    = mkClosurePVar
                , applyClosurePVar = applyClosurePVar
+               , replicatePAIntPrimVar = replicatePAIntPrimVar
+               , upToPAIntPrimVar = upToPAIntPrimVar
                , lengthPAVar      = lengthPAVar
                , replicatePAVar   = replicatePAVar
                , emptyPAVar       = emptyPAVar

compiler/vectorise/VectType.hs

@@ -297,6 +297,11 @@ arrShapeVars repr = mapM (newLocalVar FSLIT("sh")) =<< arrShapeTys repr
 
 replicateShape :: Repr -> CoreExpr -> CoreExpr -> VM [CoreExpr]
 replicateShape (ProdRepr {}) len _ = return [len]
+replicateShape (SumRepr {})  len tag
+  = do
+      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 })

compiler/vectorise/VectUtils.hs

@@ -68,7 +68,7 @@ isAnnTypeArg (_, AnnType t) = True
 isAnnTypeArg _              = False
 
 mkDataConTag :: DataCon -> CoreExpr
-mkDataConTag dc = mkConApp intDataCon [mkIntLitInt $ dataConTag dc]
+mkDataConTag = mkIntLitInt . dataConTag
 
 splitUnTy :: String -> Name -> Type -> Type
 splitUnTy s name ty