Close over kind variables when computing fixed type variables

Previously, `buildStab` would not consider kind variables when determining
which type variables need to be fixed in a generated `Lens`'s type signature.
This was not a problem in older versions of `lens`, which aggressively dropped
kind variables, but now that `lens` attempts to include kind variables in
generated type signatures, this problem has risen to the surface, resulting in
the problems observed in #972.

The solution is to take the set of fixed type variables in `buildStab` and
close over kind variables. For more information, refer to the comments I have
left near `closeOverKinds`.

Fixes #972.

CHANGELOG.markdown

@@ -1,3 +1,8 @@
+next [????.??.??]
+-----------------
+* Fix a bug in which `makeLenses` could produce ill kinded optics for
+  poly-kinded datatypes in certain situations.
+
 5 [2021.02.17]
 --------------
 * Support building with GHC 9.0.

lens.cabal

@@ -347,6 +347,7 @@ test-suite templates
   other-modules:
     T799
     T917
+    T972
   ghc-options: -Wall -threaded
   hs-source-dirs: tests
   default-language: Haskell2010

src/Control/Lens/Internal/FieldTH.hs

@@ -49,7 +49,7 @@ import Language.Haskell.TH.Lens
 import Language.Haskell.TH
 import qualified Language.Haskell.TH.Datatype as D
 import qualified Language.Haskell.TH.Datatype.TyVarBndr as D
-import Data.Maybe (isJust,maybeToList)
+import Data.Maybe (fromMaybe,isJust,maybeToList)
 import Data.List (nub, findIndices)
 import Data.Either (partitionEithers)
 import Data.Semigroup (Any (..))
@@ -249,9 +249,52 @@ buildStab s categorizedFields =
 
   where
   (fixedFields, targetFields) = partitionEithers categorizedFields
-  fixedTypeVars               = setOf typeVars fixedFields
-  unfixedTypeVars             = setOf typeVars s Set.\\ fixedTypeVars
 
+  fixedTypeVars, unfixedTypeVars :: Set Name
+  fixedTypeVars   = closeOverKinds $ setOf typeVars fixedFields
+  unfixedTypeVars = setOf typeVars s Set.\\ fixedTypeVars
+
+  -- Compute the kind variables that appear in the kind of a type variable
+  -- binder. For example, @kindVarsOfTvb (x :: (a, b)) = (x, {a, b})@. If a
+  -- type variable binder lacks an explicit kind annotation, this
+  -- conservatively assumes that there are no kind variables. For example,
+  -- @kindVarsOfTvb (y) = (y, {})@.
+  kindVarsOfTvb :: D.TyVarBndr_ flag -> (Name, Set Name)
+  kindVarsOfTvb = D.elimTV (\n   -> (n, Set.empty))
+                           (\n k -> (n, setOf typeVars k))
+
+  -- For each type variable name that appears in @s@, map to the kind variables
+  -- that appear in that type variable's kind.
+  sKindVarMap :: Map Name (Set Name)
+  sKindVarMap = Map.fromList $ map kindVarsOfTvb $ D.freeVariablesWellScoped [s]
+
+  lookupSKindVars :: Name -> Set Name
+  lookupSKindVars n = fromMaybe Set.empty $ Map.lookup n sKindVarMap
+
+  -- Consider this example (adapted from #972):
+  --
+  --   data Dart (s :: k) = Dart { _arc :: Proxy s, _direction :: Int }
+  --   $(makeLenses ''Dart)
+  --
+  -- When generating a Lens for `direction`, the type variable `s` should be
+  -- fixed. But note that (s :: k), and as a result, the kind variable `k`
+  -- needs to be fixed as well. This is because a type like this would be
+  -- ill kinded:
+  --
+  --   direction :: Lens (Dart (s :: k1)) (Dart (s :: k2)) Direction Direction
+  --
+  -- However, only `s` is mentioned syntactically in the type of `_arc`, so we
+  -- have to infer that `k` is mentioned in the kind of `s`. We accomplish this
+  -- with `closeOverKinds`, which does the following:
+  --
+  -- 1. Use freeVariablesWellScoped to compute the free type variables of
+  --    `Dart (s :: k)`, which gives us `(s :: k)`.
+  -- 2. For each type variable name in `Proxy s`, the type of `_arc`, look up
+  --    the kind variables in the type variable's kind. In the case of `s`,
+  --    the only kind variable is `k`.
+  -- 3. Add these kind variables to the set of fixed type variables.
+  closeOverKinds :: Set Name -> Set Name
+  closeOverKinds st = foldl' Set.union Set.empty (Set.map lookupSKindVars st) `Set.union` st
 
 -- | Build the signature and definition for a single field optic.
 -- In the case of a singleton constructor irrefutable matches are

tests/T917.hs

@@ -18,7 +18,7 @@ import Data.Proxy
 import Data.Kind
 #endif
 
--- Like Data.Functor.Const, but redfined to ensure that it is poly-kinded
+-- Like Data.Functor.Const, but redefined to ensure that it is poly-kinded
 -- across all versions of GHC, not just 8.0+
 newtype Constant a (b :: k) = Constant a
 

tests/T972.hs

@@ -0,0 +1,24 @@
+{-# LANGUAGE CPP #-}
+{-# LANGUAGE PolyKinds #-}
+{-# LANGUAGE TemplateHaskell #-}
+
+#if __GLASGOW_HASKELL__ >= 800 && __GLASGOW_HASKELL__ < 806
+{-# LANGUAGE TypeInType #-}
+#endif
+module T972 where
+
+import Control.Lens
+#if __GLASGOW_HASKELL__ >= 800
+import Data.Proxy
+#endif
+
+newtype Arc s = Arc { _unArc :: Int }
+
+data Direction = Negative | Positive
+data Dart s = Dart { _arc :: Arc s, _direction :: Direction }
+$(makeLenses ''Dart)
+
+#if __GLASGOW_HASKELL__ >= 800
+data Fancy k (a :: k) = MkFancy { _unFancy1 :: k, _unFancy2 :: Proxy a }
+$(makeLenses ''Fancy)
+#endif

tests/templates.hs

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