Fix type of vectorised class data constructors and add dfuns into 'VectInfo'

4 files changed, 16 insertions, 28 deletions

diff --git a/compiler/vectorise/Vectorise/Env.hs b/compiler/vectorise/Vectorise/Env.hs
index ccf034b767..0020d67412 100644
--- a/compiler/vectorise/Vectorise/Env.hs
+++ b/compiler/vectorise/Vectorise/Env.hs
@@ -194,7 +194,7 @@ setPRFunsEnv ps genv = genv { global_pr_funs = mkNameEnv ps }
 -- data constructors referenced in VECTORISE pragmas, even if they are defined in an imported
 -- module.
 --
--- The variables explicitly include class selectors.
+-- The variables explicitly include class selectors and dfuns.
 --
 modVectInfo :: GlobalEnv -> [Id] -> [TyCon] -> [CoreVect]-> VectInfo -> VectInfo
 modVectInfo env mg_ids mg_tyCons vectDecls info
@@ -206,7 +206,8 @@ modVectInfo env mg_ids mg_tyCons vectDecls info
     , vectInfoScalarTyCons = global_scalar_tycons env `minusNameSet` vectInfoScalarTyCons info
     }
   where
-    vectIds         = [id    | Vect     id    _   <- vectDecls]
+    vectIds         = [id    | Vect     id    _   <- vectDecls] ++
+                      [id    | VectInst _ id      <- vectDecls]
     vectTypeTyCons  = [tycon | VectType _ tycon _ <- vectDecls] ++
                       [tycon | VectClass tycon    <- vectDecls]
     vectDataCons    = concatMap tyConDataCons vectTypeTyCons
diff --git a/compiler/vectorise/Vectorise/Exp.hs b/compiler/vectorise/Vectorise/Exp.hs
index 1a5701cc0f..bf6fe3165e 100644
--- a/compiler/vectorise/Vectorise/Exp.hs
+++ b/compiler/vectorise/Vectorise/Exp.hs
@@ -398,16 +398,6 @@ unVectDict ty e
                                        Nothing  -> panic "Vectorise.Exp.unVectDict: no class"
     selIds                         = classAllSelIds cls
 
-{-
-!!!How about 'isClassOpId_maybe'?  Do we need to treat them specially to get the class ops for
-!!!the vectorised instances or do they just work out?? (We may want to make sure that the
-!!!vectorised Ids at least get the right IdDetails...)
-!!!NB: For *locally defined* instances, the selector functions are part of the vectorised bindings,
-!!!    but not so for *imported* instances, where we need to generate the vectorised versions from
-!!!    scratch.
-!!!Also need to take care of the builtin rules for selectors (see mkDictSelId).
- -}
-
 -- | Vectorise a lambda abstraction.
 --
 vectLam :: Bool             -- ^ When the RHS of a binding, whether that binding should be inlined.
diff --git a/compiler/vectorise/Vectorise/Monad/Naming.hs b/compiler/vectorise/Vectorise/Monad/Naming.hs
index adc2d0ca01..ecf0e81306 100644
--- a/compiler/vectorise/Vectorise/Monad/Naming.hs
+++ b/compiler/vectorise/Vectorise/Monad/Naming.hs
@@ -46,8 +46,8 @@ mkLocalisedName mk_occ name =
      ; return new_name
      }
 
--- |Produce the vectorised variant of an `Id` with the given type, while taking care that vectorised
--- dfun ids must be dfuns again.
+-- |Produce the vectorised variant of an `Id` with the given vectorised type, while taking care that
+-- vectorised dfun ids must be dfuns again.
 --
 -- Force the new name to be a system name and, if the original was an external name, disambiguate
 -- the new name with the module name of the original.
diff --git a/compiler/vectorise/Vectorise/Type/TyConDecl.hs b/compiler/vectorise/Vectorise/Type/TyConDecl.hs
index f0d05b0413..859056cd1a 100644
--- a/compiler/vectorise/Vectorise/Type/TyConDecl.hs
+++ b/compiler/vectorise/Vectorise/Type/TyConDecl.hs
@@ -49,7 +49,11 @@ vectTyConDecl tycon
        ; theta' <- mapM vectType (classSCTheta cls)
 
            -- vectorise method selectors
-       ; methods' <- sequence [ vectMethod id meth | (id, meth) <- classOpItems cls]
+       ; let opItems      = classOpItems cls
+             Just datacon = tyConSingleDataCon_maybe tycon
+             argTys       = dataConRepArgTys datacon                      -- all selector types
+             opTys        = drop (length argTys - length opItems) argTys  -- only method types
+       ; methods' <- sequence [ vectMethod id meth ty | ((id, meth), ty) <- zip opItems opTys]
 
            -- keep the original recursiveness flag
        ; let rec_flag = boolToRecFlag (isRecursiveTyCon tycon)
@@ -115,24 +119,17 @@ vectTyConDecl tycon
   | otherwise
   = cantVectorise "Can't vectorise exotic type constructor" (ppr tycon)
 
--- |Vectorise a class method.  (Don't enter into the vectorisation map yet.)
+-- |Vectorise a class method.  (Don't enter it into the vectorisation map yet.)
 --
-vectMethod :: Id -> DefMeth -> VM (Name, DefMethSpec, Type)
-vectMethod id defMeth
+vectMethod :: Id -> DefMeth -> Type -> VM (Name, DefMethSpec, Type)
+vectMethod id defMeth ty
  = do {   -- Vectorise the method type.
-      ; typ' <- vectType (varType id)
+      ; ty' <- vectType ty
 
           -- Create a name for the vectorised method.
-      ; id' <- mkVectId id typ'
+      ; id' <- mkVectId id ty'
 
-          -- When we call buildClass in vectTyConDecl, it adds foralls and dictionaries
-          -- to the types of each method. However, the types we get back from vectType
-          -- above already already have these, so we need to chop them off here otherwise
-          -- we'll get two copies in the final version.
-      ; let (_tyvars, tyBody) = splitForAllTys typ'
-      ; let (_dict,   tyRest) = splitFunTy tyBody
-
-      ; return  (Var.varName id', defMethSpecOfDefMeth defMeth, tyRest)
+      ; return  (Var.varName id', defMethSpecOfDefMeth defMeth, ty')
       }
 
 -- |Vectorise the RHS of an algebraic type.