singletons-th: Make pr_scoped_vars a list of LocalVars, not Names

Doing so makes it possible to track the kinds of scoped type variables with
greater precision during lambda lifting. See the golden output for the new
`T613` test case to see an example of this in action.

Fixes #613.

singletons-base/tests/SingletonsBaseTestSuite.hs

@@ -168,6 +168,7 @@ tests =
     , compileAndDumpStdTest "T361"
     , compileAndDumpStdTest "T601a"
     , compileAndDumpStdTest "T605"
+    , compileAndDumpStdTest "T613"
     ],
     testGroup "Database client"
     [ compileAndDumpTest "GradingClient/Database" ghcOpts

singletons-base/tests/compile-and-dump/Promote/T613.golden

@@ -0,0 +1,31 @@
+Promote/T613.hs:(0,0)-(0,0): Splicing declarations
+    promote
+      [d| f :: forall k (a :: k). Proxy a -> Proxy a
+          f x
+            = y
+            where
+                y = x |]
+  ======>
+    f :: forall k (a :: k). Proxy a -> Proxy a
+    f x
+      = y
+      where
+          y = x
+    type family Let0123456789876543210YSym0 k0123456789876543210 (a0123456789876543210 :: k0123456789876543210) (x0123456789876543210 :: Proxy a0123456789876543210) where
+      Let0123456789876543210YSym0 k0123456789876543210 a0123456789876543210 x0123456789876543210 = Let0123456789876543210Y k0123456789876543210 a0123456789876543210 x0123456789876543210
+    type family Let0123456789876543210Y k0123456789876543210 (a0123456789876543210 :: k0123456789876543210) (x0123456789876543210 :: Proxy a0123456789876543210) where
+      Let0123456789876543210Y k a x = x
+    type FSym0 :: forall k (a :: k). (~>) (Proxy a) (Proxy a)
+    data FSym0 :: (~>) (Proxy a) (Proxy a)
+      where
+        FSym0KindInference :: SameKind (Apply FSym0 arg) (FSym1 arg) =>
+                              FSym0 a0123456789876543210
+    type instance Apply @(Proxy a) @(Proxy a) FSym0 a0123456789876543210 = F a0123456789876543210
+    instance SuppressUnusedWarnings FSym0 where
+      suppressUnusedWarnings = snd ((,) FSym0KindInference ())
+    type FSym1 :: forall k (a :: k). Proxy a -> Proxy a
+    type family FSym1 @k @(a :: k) (a0123456789876543210 :: Proxy a) :: Proxy a where
+      FSym1 a0123456789876543210 = F a0123456789876543210
+    type F :: forall k (a :: k). Proxy a -> Proxy a
+    type family F @k @(a :: k) (a :: Proxy a) :: Proxy a where
+      F @k @a (x :: Proxy a) = Let0123456789876543210YSym0 k a x

singletons-base/tests/compile-and-dump/Promote/T613.hs

@@ -0,0 +1,12 @@
+module T613 where
+
+import Data.Proxy
+import Data.Proxy.Singletons
+import Data.Singletons.TH
+
+$(promote [d|
+  f :: forall k (a :: k). Proxy a -> Proxy a
+  f x = y
+    where
+      y = x
+  |])

singletons-th/src/Data/Singletons/TH/Promote.hs

@@ -658,7 +658,7 @@ promoteMethod meth_sort orig_sigs_map cls_tvbs (meth_name, meth_rhs) = do
   scoped_tvs_ext <- qIsExtEnabled LangExt.ScopedTypeVariables
   let all_meth_scoped_tvs
         | scoped_tvs_ext
-        = OSet.fromList (map extractTvbName cls_tvbs) `OSet.union` meth_scoped_tvs
+        = OSet.fromList (map tvbToLocalVar cls_tvbs) `OSet.union` meth_scoped_tvs
         | otherwise
         = OSet.empty
   meth_arg_tvs <- replicateM (length meth_arg_kis) (qNewName "a")
@@ -768,7 +768,7 @@ promoteMethod meth_sort orig_sigs_map cls_tvbs (meth_name, meth_rhs) = do
     --   variable binders primarily as a way to annotate the kinds of each
     --   variable, so it is possible for the helper type family to bind a kind
     --   variable in a `forall` without it scoping over the body.
-    promote_meth_ty :: PrM (OSet Name, [DTyVarBndrSpec], [DKind], DKind)
+    promote_meth_ty :: PrM (OSet LocalVar, [DTyVarBndrSpec], [DKind], DKind)
     promote_meth_ty =
       case meth_sort of
         DefaultMethods ->
@@ -783,7 +783,7 @@ promoteMethod meth_sort orig_sigs_map cls_tvbs (meth_name, meth_rhs) = do
               -- instance signature's type directly, and no substitution for
               -- class variables is required.
               (kvbs, arg_kis, res_ki) <- promoteUnraveled ty
-              pure (OSet.fromList (map extractTvbName kvbs), kvbs, arg_kis, res_ki)
+              pure (OSet.fromList (map tvbToLocalVar kvbs), kvbs, arg_kis, res_ki)
             Nothing -> do
               -- We don't have an InstanceSig, so we must compute the kind to use
               -- ourselves.
@@ -800,15 +800,15 @@ promoteMethod meth_sort orig_sigs_map cls_tvbs (meth_name, meth_rhs) = do
               pure (OSet.empty, kvbs', arg_kis', res_ki')
 
     -- Attempt to look up a class method's original type.
-    lookup_meth_ty :: PrM (OSet Name, [DTyVarBndrSpec], [DKind], DKind)
+    lookup_meth_ty :: PrM (OSet LocalVar, [DTyVarBndrSpec], [DKind], DKind)
     lookup_meth_ty = do
       opts <- getOptions
       let proName = promotedTopLevelValueName opts meth_name
       case OMap.lookup meth_name orig_sigs_map of
         Just ty -> do
           -- The type of the method is in scope, so promote that.
           (kvbs, arg_kis, res_ki) <- promoteUnraveled ty
-          pure (OSet.fromList (map extractTvbName kvbs), kvbs, arg_kis, res_ki)
+          pure (OSet.fromList (map tvbToLocalVar kvbs), kvbs, arg_kis, res_ki)
         Nothing -> do
           -- If the type of the method is not in scope, the only other option
           -- is to try reifying the promoted method name.
@@ -987,7 +987,8 @@ data LetDecRHSSort
     -- The right-hand side of a class method (either a default method or a
     -- method in an instance declaration).
   | ClassMethodRHS
-      (OSet Name) -- The scoped type variables to bring into scope over
+      (OSet LocalVar)
+                  -- The scoped type variables to bring into scope over
                   -- the RHS of the class method. See
                   -- Note [Scoped type variables and class methods]
                   -- in D.S.TH.Promote.Monad.
@@ -1159,7 +1160,7 @@ promoteLetDecRHS rhs_sort type_env fix_env mb_let_uniq name let_dec_rhs = do
           -- from the outermost forall into scope over the RHS.
           scoped_tvs_ext <- qIsExtEnabled LangExt.ScopedTypeVariables
           let scoped_tvs | scoped_tvs_ext
-                         = OSet.fromList (map extractTvbName tvbs)
+                         = OSet.fromList (map tvbToLocalVar tvbs)
                          | otherwise
                          = OSet.empty
           -- invariant: count_args ty == length argKs
@@ -1194,7 +1195,7 @@ data LetDecRHSKindInfo =
                          -- This will be Nothing if the let-dec RHS has local
                          -- variables that it closes over.
                          -- See Note [No SAKs for let-decs with local variables]
-        (OSet Name)      -- The scoped type variables to bring into scope over
+        (OSet LocalVar)  -- The scoped type variables to bring into scope over
                          -- the RHS of the let-dec. This will be a subset of the
                          -- type variables of the kind (see the field below),
                          -- but not necessarily the same. See
@@ -1421,7 +1422,10 @@ promotePat _ (DSigP pat ty) = do
   wildless_pat <- removeWilds pat
   ki <- promoteType ty
   (promoted, pat') <- promotePat (Just ki) wildless_pat
-  tell $ PromDPatInfos [] (fvDType ki)
+  tell $ PromDPatInfos []
+       $ OSet.fromList
+       $ map tvbToLocalVar
+       $ toposortTyVarsOf [ki]
   return (DSigT promoted ki, ADSigP promoted pat' ki)
 promotePat _ DWildP = return (DWildCardT, ADWildP)
 promotePat _ p@(DTypeP _) = fail ("Embedded type patterns cannot be promoted: " ++ show p)

singletons-th/src/Data/Singletons/TH/Promote/Monad.hs

@@ -30,7 +30,7 @@ import Data.Singletons.TH.Syntax
 -- environment during promotion
 data PrEnv =
   PrEnv { pr_options     :: Options
-        , pr_scoped_vars :: OSet Name
+        , pr_scoped_vars :: OSet LocalVar
           -- ^ The set of scoped type variables currently in scope.
           -- See @Note [Scoped type variables]@.
         , pr_lambda_vars :: OMap Name LocalVar
@@ -69,7 +69,7 @@ allLocals :: MonadReader PrEnv m => m [LocalVar]
 allLocals = do
   scoped <- asks (F.toList . pr_scoped_vars)
   lambdas <- asks (OMap.assocs . pr_lambda_vars)
-  return $ map localVarNoKind scoped ++ map snd lambdas
+  return $ scoped ++ map snd lambdas
 
 emitDecs :: MonadWriter [DDec] m => [DDec] -> m ()
 emitDecs = tell
@@ -79,10 +79,10 @@ emitDecsM action = do
   decs <- action
   emitDecs decs
 
--- ^ Bring a list of type variables into scope for the duration the supplied
+-- ^ Bring a set of type variables into scope for the duration the supplied
 -- computation. See @Note [Tracking local variables]@ and
 -- @Note [Scoped type variables]@.
-scopedBind :: OSet Name -> PrM a -> PrM a
+scopedBind :: OSet LocalVar -> PrM a -> PrM a
 scopedBind binds =
   local (\env@(PrEnv { pr_scoped_vars = scoped }) ->
     env { pr_scoped_vars = binds `OSet.union` scoped })
@@ -586,7 +586,7 @@ anymore, but it would require more work to special-case class methods in
 `promoteClause` to avoid this, and it doesn't hurt anything to leave `@b` in
 place).
 
-The `OSet Name` fields of `ClassMethodRHS` dictates which type variables to
+The `OSet LocalVar` field of `ClassMethodRHS` dictates which type variables to
 bring into scope via `scopedBind`. There are multiple places in the code which
 determine which type variables get put into the `OSet`:
 

singletons-th/src/Data/Singletons/TH/Syntax.hs

@@ -33,7 +33,7 @@ type VarPromotions = [(Name, LocalVar)]
 data PromDPatInfos = PromDPatInfos
   { prom_dpat_vars    :: VarPromotions
       -- Maps term-level pattern variables to their promoted, type-level counterparts.
-  , prom_dpat_sig_kvs :: OSet Name
+  , prom_dpat_sig_kvs :: OSet LocalVar
       -- Kind variables bound by DSigPas.
       -- See Note [Scoped type variables] in Data.Singletons.TH.Promote.Monad.
   }

singletons-th/src/Data/Singletons/TH/Syntax/LocalVar.hs

@@ -1,13 +1,14 @@
 module Data.Singletons.TH.Syntax.LocalVar
   ( LocalVar(..)
   , foldTypeLocalVars
-  , localVarNoKind
   , localVarToTvb
   , localVarToType
   , localVarToTypeArg
+  , tvbToLocalVar
   ) where
 
 import Data.Data (Data)
+import Data.Function (on)
 import Language.Haskell.TH.Desugar
 import Language.Haskell.TH.Syntax
 
@@ -22,6 +23,9 @@ import Language.Haskell.TH.Syntax
 --   variable is known, we can use it to generate code with more precise kind
 --   information. See @Note [Local variables and kind information]@.
 --
+-- We maintain the invariant that if two 'LocalVar' values share the same
+-- 'lvName', then they should both have the same 'lvKind' value.
+--
 -- A 'LocalVar' is very close in design to a 'DTyVarBndrUnit', as both contain
 -- 'Name's and optional 'DKind's. We use a separate 'LocalVar' type to represent
 -- local variables because 'LocalVar's can occur both in binding and argument
@@ -31,7 +35,17 @@ import Language.Haskell.TH.Syntax
 data LocalVar = LocalVar
   { lvName :: Name
   , lvKind :: Maybe DKind
-  } deriving Data
+  } deriving (Data, Show)
+
+-- Because of the invariant described in the Haddocks for 'LocalVar', it
+-- suffices to only check the 'lvName's when checking 'LocalVar's for equality.
+instance Eq LocalVar where
+  (==) = (==) `on` lvName
+
+-- Because of the invariant described in the Haddocks for 'LocalVar', it
+-- suffices to only compare the 'lvName's when comparing 'LocalVar's.
+instance Ord LocalVar where
+  compare = compare `on` lvName
 
 -- | Apply a 'DType' to a list of 'LocalVar' arguments. Because these
 -- 'LocalVar's occur in argument positions, they will not contain any kind
@@ -40,10 +54,6 @@ data LocalVar = LocalVar
 foldTypeLocalVars :: DType -> [LocalVar] -> DType
 foldTypeLocalVars ty = applyDType ty . map localVarToTypeArg
 
--- | Construct a 'LocalVar' with no kind information.
-localVarNoKind :: Name -> LocalVar
-localVarNoKind nm = LocalVar { lvName = nm, lvKind = Nothing }
-
 -- | Convert a 'LocalVar' used in a binding position to a 'DTyVarBndr' using the
 -- supplied @flag@. Because this is used in a binding position, we include kind
 -- information (if available) in the 'DTyVarBndr'. See @Note [Local variables
@@ -69,6 +79,13 @@ localVarToType (LocalVar { lvName = local_nm }) = DVarT local_nm
 localVarToTypeArg :: LocalVar -> DTypeArg
 localVarToTypeArg = DTANormal . localVarToType
 
+-- | Convert a 'DTyVarBndr' to a 'LocalVar'.
+tvbToLocalVar :: DTyVarBndr flag -> LocalVar
+tvbToLocalVar (DPlainTV nm _) =
+  LocalVar { lvName = nm, lvKind = Nothing }
+tvbToLocalVar (DKindedTV nm _ kind) =
+  LocalVar { lvName = nm, lvKind = Just kind }
+
 {-
 Note [Local variables and kind information]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~