Make the TH machinery handle PolyKinds more robustly (#945)

This is a collection of various Template Haskell–related fixes that, when all
put together, fixes #917. This does the following:

* Rather than use `th-abstraction`'s `datatypeType` function, which strips off
  important kind information from type arguments, I defined a similar
  `datatypeTypeKinded` function that preserves kinds when reasonable.
* `Control.Lens.Internal.{FieldTH,PrismTH}` is now more careful to use
  `freeVariablesWellScoped` (from `th-abstraction`) instead of `typeVars`
  to ensure that the resulting types are well scoped. This is particularly
  important for poly-kinded types, as the kind variables must always appear
  before the type variables.
* I deleted the `close` function from `Control.Lens.Internal.PrismTH` in favor
  of `quantifyType` and `quantifyType'`, which I have moved to
  `Control.Lens.Internal.TH` so that they may be used by `FieldTH` and
  `PrismTH` alike. Moreover, I now use `quantifyType'` in the definition of
  `PrismTH.makeClassyPrismClass` so that any type variables bound by the class
  itself do not get requantified in any class methods. The previous code was
  not doing this at all, which was just plain wrong.

CHANGELOG.markdown

@@ -21,6 +21,13 @@
     swapped = iso Swap.swap Swap.swap
   #endif
   ```
+* Make the functions in `Control.Lens.TH` work more robustly with poly-kinded
+  data types. This can cause a breaking change under certain situations:
+  * TH-generated optics for poly-kinded data types are now much more likely to
+    mention kind variables in their definitions, which will require enabling
+    the `PolyKinds` extension at use sites in order to typecheck.
+  * Because TH-generated optics now quantify more kind variables than they did
+    previously, this can affect the order of visible type applications.
 * Add `Control.Lens.Profunctor` with conversion functions to and from
   profunctor optic representation
 * Add `Control.Lens.Review.reviewing`, which is like `review` but with a more

lens.cabal

@@ -212,7 +212,7 @@ library
     tagged                    >= 0.4.4    && < 1,
     template-haskell          >= 2.9.0.0  && < 2.18,
     these                     >= 1.1.1.1  && < 1.2,
-    th-abstraction            >= 0.4      && < 0.5,
+    th-abstraction            >= 0.4.1    && < 0.5,
     text                      >= 1.2.3.0  && < 1.3,
     transformers              >= 0.3.0.0  && < 0.6,
     transformers-compat       >= 0.4      && < 1,
@@ -353,11 +353,16 @@ library
 test-suite templates
   type: exitcode-stdio-1.0
   main-is: templates.hs
-  other-modules: T799
+  other-modules:
+    T799
+    T917
   ghc-options: -Wall -threaded
   hs-source-dirs: tests
   default-language: Haskell2010
 
+  if impl(ghc >= 8.6)
+    ghc-options: -Wno-star-is-type
+
   if flag(dump-splices)
     ghc-options: -ddump-splices
 

src/Control/Lens/Internal/FieldTH.hs

@@ -95,9 +95,9 @@ makeFieldOpticsForDatatype rules info =
                      return (concat decss)
 
   where
-  tyName = D.datatypeName info
-  s      = D.datatypeType info
-  cons   = D.datatypeCons info
+  tyName = D.datatypeName     info
+  s      = datatypeTypeKinded info
+  cons   = D.datatypeCons     info
 
   -- Traverse the field labels of a normalized constructor
   normFieldLabels :: Traversal [(Name,[(a,Type)])] [(Name,[(b,Type)])] a b
@@ -320,12 +320,13 @@ makeClassyClass className methodName s defs = do
   let ss   = map (stabToS . view (_2 . _2)) defs
   (sub,s') <- unifyTypes (s : ss)
   c <- newName "c"
-  let vars = toListOf typeVars s'
+  let vars     = D.freeVariablesWellScoped [s']
+      varNames = map D.tvName vars
       fd   | null vars = []
-           | otherwise = [FunDep [c] vars]
+           | otherwise = [FunDep [c] varNames]
 
 
-  classD (cxt[]) className (map D.plainTV (c:vars)) fd
+  classD (cxt[]) className (D.plainTV c:vars) fd
     $ sigD methodName (return (lens'TypeName `conAppsT` [VarT c, s']))
     : concat
       [ [sigD defName (return ty)
@@ -334,7 +335,7 @@ makeClassyClass className methodName s defs = do
         inlinePragma defName
       | (TopName defName, (_, stab, _)) <- defs
       , let body = appsE [varE composeValName, varE methodName, varE defName]
-      , let ty   = quantifyType' (Set.fromList (c:vars))
+      , let ty   = quantifyType' (Set.fromList (c:varNames))
                                  (stabToContext stab)
                  $ stabToOptic stab `conAppsT`
                        [VarT c, applyTypeSubst sub (stabToA stab)]
@@ -356,8 +357,8 @@ makeClassyInstance rules className methodName s defs = do
            : map return (concat methodss)
 
   where
-  instanceHead = className `conAppsT` (s : map VarT vars)
-  vars         = toListOf typeVars s
+  instanceHead = className `conAppsT` (s : map tvbToType vars)
+  vars         = D.freeVariablesWellScoped [s]
   rules'       = rules { _generateSigs    = False
                        , _generateClasses = False
                        }
@@ -620,23 +621,3 @@ type HasFieldClasses = StateT (Set Name) Q
 
 addFieldClassName :: Name -> HasFieldClasses ()
 addFieldClassName n = modify $ Set.insert n
-
-------------------------------------------------------------------------
--- Miscellaneous utility functions
-------------------------------------------------------------------------
-
-
--- | Template Haskell wants type variables declared in a forall, so
--- we find all free type variables in a given type and declare them.
-quantifyType :: Cxt -> Type -> Type
-quantifyType = quantifyType' Set.empty
-
--- | This function works like 'quantifyType' except that it takes
--- a list of variables to exclude from quantification.
-quantifyType' :: Set Name -> Cxt -> Type -> Type
-quantifyType' exclude c t = ForallT vs c t
-  where
-  vs = map D.plainTVSpecified
-     $ filter (`Set.notMember` exclude)
-     $ nub -- stable order
-     $ toListOf typeVars t

src/Control/Lens/Internal/PrismTH.hs

@@ -23,7 +23,6 @@ module Control.Lens.Internal.PrismTH
   ) where
 
 import Control.Applicative
-import Control.Lens.Fold
 import Control.Lens.Getter
 import Control.Lens.Internal.TH
 import Control.Lens.Lens
@@ -39,6 +38,7 @@ import qualified Language.Haskell.TH.Datatype.TyVarBndr as D
 import Language.Haskell.TH.Lens
 import qualified Data.Map as Map
 import qualified Data.Set as Set
+import Data.Set (Set)
 import Prelude
 
 -- | Generate a 'Prism' for each constructor of a data type.
@@ -137,7 +137,7 @@ makePrisms' normal typeName =
      let cls | normal    = Nothing
              | otherwise = Just (D.datatypeName info)
          cons = D.datatypeCons info
-     makeConsPrisms (D.datatypeType info) (map normalizeCon cons) cls
+     makeConsPrisms (datatypeTypeKinded info) (map normalizeCon cons) cls
 
 
 -- | Generate prisms for the given 'Dec'
@@ -147,7 +147,7 @@ makeDecPrisms normal dec =
      let cls | normal    = Nothing
              | otherwise = Just (D.datatypeName info)
          cons = D.datatypeCons info
-     makeConsPrisms (D.datatypeType info) (map normalizeCon cons) cls
+     makeConsPrisms (datatypeTypeKinded info) (map normalizeCon cons) cls
 
 
 -- | Generate prisms for the given type, normalized constructors, and
@@ -166,7 +166,7 @@ makeConsPrisms t cons Nothing =
        stab <- computeOpticType t cons con
        let n = prismName conName
        sequenceA
-         ( [ sigD n (close (stabToType stab))
+         ( [ sigD n (return (quantifyType [] (stabToType Set.empty stab)))
            , valD (varP n) (normalB (makeConOpticExp stab cons con)) []
            ]
            ++ inlinePragma n
@@ -197,17 +197,15 @@ simplifyStab (Stab cx ty _ t _ b) = Stab cx ty t t b b
 stabSimple :: Stab -> Bool
 stabSimple (Stab _ _ s t a b) = s == t && a == b
 
-stabToType :: Stab -> Type
-stabToType stab@(Stab cx ty s t a b) = ForallT vs cx $
-  case ty of
-    PrismType  | stabSimple stab -> prism'TypeName  `conAppsT` [t,b]
-               | otherwise       -> prismTypeName   `conAppsT` [s,t,a,b]
-    ReviewType                   -> reviewTypeName  `conAppsT` [t,b]
-
+stabToType :: Set Name -> Stab -> Type
+stabToType clsTVBNames stab@(Stab cx ty s t a b) =
+  quantifyType' clsTVBNames cx stabTy
   where
-  vs = map D.plainTVSpecified
-     $ nub -- stable order
-     $ toListOf typeVars cx
+  stabTy =
+    case ty of
+      PrismType  | stabSimple stab -> prism'TypeName  `conAppsT` [t,b]
+                 | otherwise       -> prismTypeName   `conAppsT` [s,t,a,b]
+      ReviewType                   -> reviewTypeName  `conAppsT` [t,b]
 
 stabType :: Stab -> OpticType
 stabType (Stab _ o _ _ _ _) = o
@@ -256,7 +254,7 @@ computeIsoType t' fields =
             | otherwise    = appsT (conT isoTypeName) [s,t,a,b]
 #endif
 
-     close =<< ty
+     quantifyType [] <$> ty
 
 
 
@@ -420,28 +418,30 @@ makeClassyPrismClass t className methodName cons =
 #ifndef HLINT
      let methodType = appsT (conT prism'TypeName) [varT r,return t]
 #endif
-     methodss <- traverse (mkMethod (VarT r)) cons'
-     classD (cxt[]) className (map D.plainTV (r : vs)) (fds r)
+     methodss <- traverse (mkMethod r) cons'
+     classD (cxt[]) className (D.plainTV r : vs) (fds r)
        ( sigD methodName methodType
        : map return (concat methodss)
        )
 
   where
   mkMethod r con =
     do Stab cx o _ _ _ b <- computeOpticType t cons con
-       let stab' = Stab cx o r r b b
+       let rTy   = VarT r
+           stab' = Stab cx o rTy rTy b b
            defName = view nconName con
            body    = appsE [varE composeValName, varE methodName, varE defName]
        sequenceA
-         [ sigD defName        (return (stabToType stab'))
+         [ sigD defName        (return (stabToType (Set.fromList (r:vNames)) stab'))
          , valD (varP defName) (normalB body) []
          ]
 
   cons'         = map (over nconName prismName) cons
-  vs            = Set.toList (setOf typeVars t)
+  vs            = D.freeVariablesWellScoped [t]
+  vNames        = map D.tvName vs
   fds r
     | null vs   = []
-    | otherwise = [FunDep [r] vs]
+    | otherwise = [FunDep [r] vNames]
 
 
 
@@ -457,8 +457,8 @@ makeClassyPrismInstance ::
   [NCon] {- Constructors    -} ->
   DecQ
 makeClassyPrismInstance s className methodName cons =
-  do let vs = Set.toList (setOf typeVars s)
-         cls = className `conAppsT` (s : map VarT vs)
+  do let vs = D.freeVariablesWellScoped [s]
+         cls = className `conAppsT` (s : map tvbToType vs)
 
      instanceD (cxt[]) (return cls)
        (   valD (varP methodName)
@@ -541,9 +541,3 @@ prismName' sameNameAsCon n =
     prefix :: Char -> String -> String
     prefix char str | sameNameAsCon = char:char:str
                     | otherwise     =      char:str
-
--- | Quantify all the free variables in a type.
-close :: Type -> TypeQ
-close t = forallT (map D.plainTVSpecified (Set.toList vs)) (cxt[]) (return t)
-  where
-  vs = setOf typeVars t

src/Control/Lens/Internal/TH.hs

@@ -22,7 +22,11 @@
 module Control.Lens.Internal.TH where
 
 import Data.Functor.Contravariant
+import qualified Data.Set as Set
+import Data.Set (Set)
 import Language.Haskell.TH
+import qualified Language.Haskell.TH.Datatype as D
+import qualified Language.Haskell.TH.Datatype.TyVarBndr as D
 import Language.Haskell.TH.Syntax
 #ifndef CURRENT_PACKAGE_KEY
 import Data.Version (showVersion)
@@ -88,6 +92,80 @@ unfoldType = go []
 #endif
     go acc ty               = (ty, acc)
 
+-- Construct a 'Type' using the datatype's type constructor and type
+-- parameters. Unlike 'D.datatypeType', kind signatures are preserved to
+-- some extent. (See the comments for 'dropSigsIfNonDataFam' below for more
+-- details on this.)
+datatypeTypeKinded :: D.DatatypeInfo -> Type
+datatypeTypeKinded di
+  = foldl AppT (ConT (D.datatypeName di))
+  $ dropSigsIfNonDataFam
+  $ D.datatypeInstTypes di
+  where
+    {-
+    In an effort to prevent users from having to enable KindSignatures every
+    time that they use lens' TH functionality, we strip off reified kind
+    annotations from when:
+
+    1. The kind of a type does not contain any kind variables. If it *does*
+       contain kind variables, we want to preserve them so that we can generate
+       type signatures that preserve the dependency order of kind and type
+       variables. (The data types in test/T917.hs contain examples where this
+       is important.) This will require enabling `PolyKinds`, but since
+       `PolyKinds` implies `KindSignatures`, we can at least accomplish two
+       things at once.
+    2. The data type is not an instance of a data family. We make an exception
+       for data family instances, since the presence or absence of a kind
+       annotation can be the difference between typechecking or not.
+       (See T917DataFam in tests/T917.hs for an example.) Moreover, the
+       `TypeFamilies` extension implies `KindSignatures`.
+    -}
+    dropSigsIfNonDataFam :: [Type] -> [Type]
+    dropSigsIfNonDataFam
+      | isDataFamily (D.datatypeVariant di) = id
+      | otherwise                           = map dropSig
+
+    dropSig :: Type -> Type
+    dropSig (SigT t k) | null (D.freeVariables k) = t
+    dropSig t                                     = t
+
+-- | Template Haskell wants type variables declared in a forall, so
+-- we find all free type variables in a given type and declare them.
+quantifyType :: Cxt -> Type -> Type
+quantifyType = quantifyType' Set.empty
+
+-- | This function works like 'quantifyType' except that it takes
+-- a list of variables to exclude from quantification.
+quantifyType' :: Set Name -> Cxt -> Type -> Type
+quantifyType' exclude c t = ForallT vs c t
+  where
+  vs = filter (\tvb -> D.tvName tvb `Set.notMember` exclude)
+     $ D.changeTVFlags D.SpecifiedSpec
+     $ D.freeVariablesWellScoped (t:concatMap predTypes c) -- stable order
+
+  predTypes :: Pred -> [Type]
+#if MIN_VERSION_template_haskell(2,10,0)
+  predTypes p = [p]
+#else
+  predTypes (ClassP _ ts)  = ts
+  predTypes (EqualP t1 t2) = [t1, t2]
+#endif
+
+-- | Convert a 'TyVarBndr' into its corresponding 'Type'.
+tvbToType :: D.TyVarBndr_ flag -> Type
+tvbToType = D.elimTV VarT (SigT . VarT)
+
+-- | Peel off a kind signature from a Type (if it has one).
+unSigT :: Type -> Type
+unSigT (SigT t _) = t
+unSigT t          = t
+
+isDataFamily :: D.DatatypeVariant -> Bool
+isDataFamily D.Datatype        = False
+isDataFamily D.Newtype         = False
+isDataFamily D.DataInstance    = True
+isDataFamily D.NewtypeInstance = True
+
 ------------------------------------------------------------------------
 -- Manually quoted names
 ------------------------------------------------------------------------

tests/T917.hs

@@ -0,0 +1,45 @@
+{-# LANGUAGE CPP #-}
+{-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE FunctionalDependencies #-}
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE PolyKinds #-}
+{-# LANGUAGE TemplateHaskell #-}
+{-# LANGUAGE TypeFamilies #-}
+
+#if __GLASGOW_HASKELL__ >= 800 && __GLASGOW_HASKELL__ < 806
+{-# LANGUAGE TypeInType #-}
+#endif
+module T917 where
+
+import Control.Lens
+import Data.Proxy
+
+#if __GLASGOW_HASKELL__ >= 800 && __GLASGOW_HASKELL__ < 806
+import Data.Kind
+#endif
+
+-- Like Data.Functor.Const, but redfined to ensure that it is poly-kinded
+-- across all versions of GHC, not just 8.0+
+newtype Constant a (b :: k) = Constant a
+
+data T917OneA (a :: k -> *) (b :: k -> *) = MkT917OneA
+data T917OneB a b = MkT917OneB (T917OneA a (Const b))
+$(makePrisms ''T917OneB)
+
+data T917TwoA (a :: k -> *) (b :: k -> *) = MkT917TwoA
+data T917TwoB a b = MkT917TwoB (T917TwoA a (Const b))
+$(makeClassyPrisms ''T917TwoB)
+
+data family   T917DataFam (a :: k)
+data instance T917DataFam (a :: *) = MkT917DataFam { _unT917DataFam :: Proxy a }
+$(makeLenses 'MkT917DataFam)
+
+#if __GLASGOW_HASKELL__ >= 800
+data T917GadtOne (a :: k) where
+  MkT917GadtOne :: T917GadtOne (a :: *)
+$(makePrisms ''T917GadtOne)
+
+data T917GadtTwo (a :: k) where
+  MkT917GadtTwo :: T917GadtTwo (a :: *)
+$(makePrisms ''T917GadtTwo)
+#endif

tests/templates.hs

@@ -27,6 +27,7 @@ module Main where
 import Control.Lens
 -- import Test.QuickCheck (quickCheck)
 import T799 ()
+import T917 ()
 
 data Bar a b c = Bar { _baz :: (a, b) }
 makeLenses ''Bar