Handle associated types properly in D.S.Promote.Defun (#418)

Associated type families were going through the same code path as
top-level open type families during defunctionalization. But
top-level open type families always default any type variable binders
without explicit kinds to be of kind `Type`. This approach does not
always work for associated type families, leading to the bug observed
in #414.

This fixes the issue by partitioning associated type families
separately from other open type families in `D.S.Partition` and
introducing a new `defunAssociatedTypeFamilies` function in
`D.S.Promote.Defun` to handle them. I've also added some more
comments in various places to make things a bit clearer.

Fixes #414.

src/Data/Singletons/Partition.hs

@@ -99,8 +99,8 @@ partitionDec (DClassD cxt name tvbs fds decs) = do
                                                , cd_name      = name
                                                , cd_tvbs      = tvbs
                                                , cd_fds       = fds
-                                               , cd_lde       = lde }]
-                  , pd_open_type_family_decs = otfs }
+                                               , cd_lde       = lde
+                                               , cd_atfs      = otfs}] }
 partitionDec (DInstanceD _ _ cxt ty decs) = do
   (defns, sigs) <- liftM (bimap catMaybes mconcat) $
                    mapAndUnzipM partitionInstanceDec decs
@@ -312,4 +312,9 @@ Also note that:
 1. Other uses of type synonyms in singled code will be expanded away.
 2. Other uses of type families in singled code are unlikely to work at present
    due to Trac #12564.
+3. We track open type families, closed type families, and associated type
+   families separately, as each form of type family has different kind
+   inference behavior. See defunTopLevelTypeDecls and
+   defunAssociatedTypeFamilies in D.S.Promote.Defun for how these differences
+   manifest.
 -}

src/Data/Singletons/Promote.hs

@@ -229,7 +229,7 @@ promoteDecs raw_decls = do
         , pd_closed_type_family_decs = c_tyfams
         , pd_derived_eq_decs         = derived_eq_decs } <- partitionDecs decls
 
-  defunTypeDecls ty_syns c_tyfams o_tyfams
+  defunTopLevelTypeDecls ty_syns c_tyfams o_tyfams
     -- promoteLetDecs returns LetBinds, which we don't need at top level
   _ <- promoteLetDecs noPrefix let_decs
   mapM_ promoteClassDec classes
@@ -289,6 +289,7 @@ promoteClassDec :: UClassDecl
 promoteClassDec decl@(ClassDecl { cd_name = cls_name
                                 , cd_tvbs = tvbs
                                 , cd_fds  = fundeps
+                                , cd_atfs = atfs
                                 , cd_lde  = lde@LetDecEnv
                                     { lde_defns = defaults
                                     , lde_types = meth_sigs
@@ -302,6 +303,7 @@ promoteClassDec decl@(ClassDecl { cd_name = cls_name
     sig_decs <- mapM (uncurry promote_sig) meth_sigs_list
     (default_decs, ann_rhss, prom_rhss)
       <- mapAndUnzip3M (promoteMethod DefaultMethods meth_sigs) defaults_list
+    defunAssociatedTypeFamilies tvbs atfs
 
     infix_decls' <- mapMaybeM (uncurry promoteInfixDecl) $ OMap.assocs infix_decls
     cls_infix_decls <- promoteReifiedInfixDecls $ cls_name:meth_names

src/Data/Singletons/Promote/Defun.hs

@@ -36,11 +36,14 @@ defunInfo (DTyVarI _name _ty) =
 defunInfo (DPatSynI {}) =
   fail "Building defunctionalization symbols of pattern synonyms not supported"
 
-defunTypeDecls :: [TySynDecl]
-               -> [ClosedTypeFamilyDecl]
-               -> [OpenTypeFamilyDecl]
-               -> PrM ()
-defunTypeDecls ty_syns c_tyfams o_tyfams = do
+-- Defunctionalize type families defined at the top level (i.e., not associated
+-- with a type class).
+defunTopLevelTypeDecls ::
+     [TySynDecl]
+  -> [ClosedTypeFamilyDecl]
+  -> [OpenTypeFamilyDecl]
+  -> PrM ()
+defunTopLevelTypeDecls ty_syns c_tyfams o_tyfams = do
   defun_ty_syns <-
     concatMapM (\(TySynDecl name tvbs rhs) -> buildDefunSymsTySynD name tvbs rhs) ty_syns
   defun_c_tyfams <-
@@ -49,6 +52,58 @@ defunTypeDecls ty_syns c_tyfams o_tyfams = do
     concatMapM (buildDefunSymsOpenTypeFamilyD . getTypeFamilyDecl) o_tyfams
   emitDecs $ defun_ty_syns ++ defun_c_tyfams ++ defun_o_tyfams
 
+-- Defunctionalize all the type families associated with a type class.
+defunAssociatedTypeFamilies ::
+     [DTyVarBndr]         -- The type variables bound by the parent class
+  -> [OpenTypeFamilyDecl] -- The type families associated with the parent class
+  -> PrM ()
+defunAssociatedTypeFamilies cls_tvbs atfs = do
+  defun_atfs <- concatMapM defun atfs
+  emitDecs defun_atfs
+  where
+    defun :: OpenTypeFamilyDecl -> PrM [DDec]
+    defun (TypeFamilyDecl tf_head)
+      | cls_has_cusk
+        -- If the parent class has a CUSK, so does its associated type
+        -- families. Default type variable kinds and the result kind to be Type
+        -- if they do not yet have an explicit kind.
+      = buildDefunSymsTypeFamilyHead (cuskify . ascribe_tf_tvb_kind)
+                                     (Just . defaultToTypeKind) tf_head
+      | otherwise
+        -- If the parent class lacks a CUSK, we cannot safely default kinds to
+        -- Type. Use whatever information from the parent class is available
+        -- and let kind inference do the rest.
+      = buildDefunSymsTypeFamilyHead ascribe_tf_tvb_kind id tf_head
+
+    -- A class has a CUSK when all of its type variable binders have explicit
+    -- kinds. In other words, `cls_tvb_kind_map` should have an entry for each
+    -- class-bound type variable.
+    cls_has_cusk :: Bool
+    cls_has_cusk = length cls_tvbs == Map.size cls_tvb_kind_map
+
+    -- Maps class-bound type variables to their kind annotations (if supplied).
+    -- For example, `class C (a :: Bool) b (c :: Type)` will produce
+    -- {a |-> Bool, c |-> Type}.
+    cls_tvb_kind_map :: Map Name DKind
+    cls_tvb_kind_map = Map.fromList [ (extractTvbName cls_tvb, cls_tvb_kind)
+                                    | cls_tvb <- cls_tvbs
+                                    , Just cls_tvb_kind <- [extractTvbKind cls_tvb]
+                                    ]
+
+    -- We can sometimes learn more specific information about unannotated type
+    -- family binders from the paren class, as in the following example:
+    --
+    --   class C (a :: Bool) where
+    --     type T a :: Type
+    --
+    -- Here, we know that `T :: Bool -> Type` because we can infer that the `a`
+    -- in `type T a` should be of kind `Bool` from the class header.
+    ascribe_tf_tvb_kind :: DTyVarBndr -> DTyVarBndr
+    ascribe_tf_tvb_kind tvb =
+      case tvb of
+        DKindedTV{} -> tvb
+        DPlainTV n  -> maybe tvb (DKindedTV n) $ Map.lookup n cls_tvb_kind_map
+
 buildDefunSyms :: DDec -> PrM [DDec]
 buildDefunSyms (DDataD _new_or_data _cxt _tyName _tvbs _k ctors _derivings) =
   buildDefunSymsDataD ctors
@@ -63,15 +118,21 @@ buildDefunSyms (DClassD _cxt name tvbs _fundeps _members) = do
 buildDefunSyms _ = fail $ "Defunctionalization symbols can only be built for " ++
                           "type families and data declarations"
 
+-- If a closed type family lacks an explicit kind for any of its type variables
+-- or its result kind, then it does not have a CUSK, so we must let kind
+-- inference do the rest.
 buildDefunSymsClosedTypeFamilyD :: DTypeFamilyHead -> PrM [DDec]
 buildDefunSymsClosedTypeFamilyD = buildDefunSymsTypeFamilyHead id id
 
+-- Top-level open type families always have CUSKs, so it is safe to default
+-- type variable kinds or the result kind to Type if they are not indicated
+-- explicitly.
 buildDefunSymsOpenTypeFamilyD :: DTypeFamilyHead -> PrM [DDec]
 buildDefunSymsOpenTypeFamilyD = buildDefunSymsTypeFamilyHead cuskify (Just . defaultToTypeKind)
 
 buildDefunSymsTypeFamilyHead
-  :: (DTyVarBndr -> DTyVarBndr)
-  -> (Maybe DKind -> Maybe DKind)
+  :: (DTyVarBndr -> DTyVarBndr)   -- How to default each type variable binder
+  -> (Maybe DKind -> Maybe DKind) -- How to default the result kind
   -> DTypeFamilyHead -> PrM [DDec]
 buildDefunSymsTypeFamilyHead default_tvb default_kind
     (DTypeFamilyHead name tvbs result_sig _) = do

src/Data/Singletons/Single.hs

@@ -309,7 +309,7 @@ singTopLevelDecs locals raw_decls = withLocalDeclarations locals $ do
         , pd_derived_show_decs       = derivedShowDecs } <- partitionDecs decls
 
   ((letDecEnv, classes', insts'), promDecls) <- promoteM locals $ do
-    defunTypeDecls ty_syns c_tyfams o_tyfams
+    defunTopLevelTypeDecls ty_syns c_tyfams o_tyfams
     promoteDataDecs datas
     (_, letDecEnv) <- promoteLetDecs noPrefix letDecls
     classes' <- mapM promoteClassDec classes

src/Data/Singletons/Syntax.hs

@@ -62,12 +62,18 @@ newtype TypeFamilyDecl (info :: FamilyInfo)
 -- Whether a type family is open or closed.
 data FamilyInfo = Open | Closed
 
-data ClassDecl ann = ClassDecl { cd_cxt  :: DCxt
-                               , cd_name :: Name
-                               , cd_tvbs :: [DTyVarBndr]
-                               , cd_fds  :: [FunDep]
-                               , cd_lde  :: LetDecEnv ann
-                               }
+data ClassDecl ann
+  = ClassDecl { cd_cxt  :: DCxt
+              , cd_name :: Name
+              , cd_tvbs :: [DTyVarBndr]
+              , cd_fds  :: [FunDep]
+              , cd_lde  :: LetDecEnv ann
+              , cd_atfs :: [OpenTypeFamilyDecl]
+                  -- Associated type families. Only recorded for
+                  -- defunctionalization purposes.
+                  -- See Note [Partitioning, type synonyms, and type families]
+                  -- in D.S.Partition.
+              }
 
 data InstDecl  ann = InstDecl { id_cxt     :: DCxt
                               , id_name    :: Name

tests/SingletonsTestSuite.hs

@@ -122,6 +122,7 @@ tests =
     , compileAndDumpStdTest "T402"
     , compileAndDumpStdTest "T410"
     , compileAndDumpStdTest "T412"
+    , compileAndDumpStdTest "T414"
     ],
     testCompileAndDumpGroup "Promote"
     [ compileAndDumpStdTest "Constructors"

tests/compile-and-dump/Singletons/T414.ghc88.template

@@ -0,0 +1,53 @@
+Singletons/T414.hs:(0,0)-(0,0): Splicing declarations
+    singletons
+      [d| class C1 (a :: Bool) where
+            type T1 a b
+          class C2 a where
+            type T2 a b |]
+  ======>
+    class C1 (a :: Bool) where
+      type T1 a b
+    class C2 a where
+      type T2 a b
+    type T1Sym2 (a0123456789876543210 :: Bool) (b0123456789876543210 :: GHC.Types.Type) =
+        T1 a0123456789876543210 b0123456789876543210
+    instance SuppressUnusedWarnings (T1Sym1 a0123456789876543210) where
+      suppressUnusedWarnings = snd (((,) T1Sym1KindInference) ())
+    data T1Sym1 (a0123456789876543210 :: Bool) :: (~>) GHC.Types.Type GHC.Types.Type
+      where
+        T1Sym1KindInference :: forall a0123456789876543210
+                                      b0123456789876543210
+                                      arg. SameKind (Apply (T1Sym1 a0123456789876543210) arg) (T1Sym2 a0123456789876543210 arg) =>
+                               T1Sym1 a0123456789876543210 b0123456789876543210
+    type instance Apply (T1Sym1 a0123456789876543210) b0123456789876543210 = T1 a0123456789876543210 b0123456789876543210
+    instance SuppressUnusedWarnings T1Sym0 where
+      suppressUnusedWarnings = snd (((,) T1Sym0KindInference) ())
+    data T1Sym0 :: (~>) Bool ((~>) GHC.Types.Type GHC.Types.Type)
+      where
+        T1Sym0KindInference :: forall a0123456789876543210
+                                      arg. SameKind (Apply T1Sym0 arg) (T1Sym1 arg) =>
+                               T1Sym0 a0123456789876543210
+    type instance Apply T1Sym0 a0123456789876543210 = T1Sym1 a0123456789876543210
+    class PC1 (a :: Bool)
+    type T2Sym2 a0123456789876543210 b0123456789876543210 =
+        T2 a0123456789876543210 b0123456789876543210
+    instance SuppressUnusedWarnings (T2Sym1 a0123456789876543210) where
+      suppressUnusedWarnings = snd (((,) T2Sym1KindInference) ())
+    data T2Sym1 a0123456789876543210 b0123456789876543210
+      where
+        T2Sym1KindInference :: forall a0123456789876543210
+                                      b0123456789876543210
+                                      arg. SameKind (Apply (T2Sym1 a0123456789876543210) arg) (T2Sym2 a0123456789876543210 arg) =>
+                               T2Sym1 a0123456789876543210 b0123456789876543210
+    type instance Apply (T2Sym1 a0123456789876543210) b0123456789876543210 = T2 a0123456789876543210 b0123456789876543210
+    instance SuppressUnusedWarnings T2Sym0 where
+      suppressUnusedWarnings = snd (((,) T2Sym0KindInference) ())
+    data T2Sym0 a0123456789876543210
+      where
+        T2Sym0KindInference :: forall a0123456789876543210
+                                      arg. SameKind (Apply T2Sym0 arg) (T2Sym1 arg) =>
+                               T2Sym0 a0123456789876543210
+    type instance Apply T2Sym0 a0123456789876543210 = T2Sym1 a0123456789876543210
+    class PC2 a
+    class SC1 (a :: Bool)
+    class SC2 a

tests/compile-and-dump/Singletons/T414.hs

@@ -0,0 +1,11 @@
+module T414 where
+
+import Data.Singletons.TH
+
+$(singletons [d|
+  class C1 (a :: Bool) where
+    type T1 a b
+
+  class C2 a where
+    type T2 a b
+  |])