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

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
@@ -1393,7 +1394,8 @@ promotePat _ (DLitP lit) = (, ADLitP lit) <$> promoteLitPat lit
 promotePat m_ki (DVarP name) = do
       -- term vars can be symbols... type vars can't!
   tyName <- mkTyName name
-  tell $ PromDPatInfos [(name, (tyName, m_ki))] OSet.empty
+  let lv = LocalVar { lvName = tyName, lvKind = m_ki }
+  tell $ PromDPatInfos [(name, lv)] OSet.empty
   return (DVarT tyName, ADVarP name)
 promotePat _ (DConP name tys pats) = do
   opts <- getOptions
@@ -1420,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)
@@ -1584,7 +1589,8 @@ dTypeFamilyHead_with_locals tf_nm local_vars arg_tvbs res_sig =
     -- Ensure that all references to local_nms are substituted away.
     subst1 = Map.fromList $
              zipWith
-               (\(local_nm, _) (local_nm', _) -> (local_nm, DVarT local_nm'))
+               (\(LocalVar { lvName = local_nm }) (LocalVar { lvName = local_nm' }) ->
+                 (local_nm, DVarT local_nm'))
                local_vars
                local_vars'
     (subst2, arg_tvbs') = substTvbs subst1 arg_tvbs

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/Single.hs

@@ -725,7 +725,7 @@ singPat var_proms = go
       tyname <- case Map.lookup name var_proms of
                   Nothing     ->
                     fail "Internal error: unknown variable when singling pattern"
-                  Just (tyname, _) -> return tyname
+                  Just (LocalVar { lvName = tyname }) -> return tyname
       pure $ DVarP (singledValueName opts name)
                `DSigP` (singFamily `DAppT` DVarT tyname)
     go (ADConP name tys pats) = do

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,33 +1,51 @@
 module Data.Singletons.TH.Syntax.LocalVar
-  ( 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
 
 -- | A local variable that is captured in a lambda-lifted type family. (See
 -- @Note [Tracking local variables]@ in "Data.Singletons.TH.Promote.Monad" for
 -- an explanation of how lambda lifting works.) A 'LocalVar' consists of:
 --
--- * A 'Name' that corresponds to the promoted, type-level version of a
---   term-level variable name.
+-- * An @'lvName' :: 'Name'@ that corresponds to the promoted, type-level
+--   version of a term-level variable name.
 --
--- * An optional kind (in the form of @'Maybe' 'DKind'@). When the kind of a
---   local variable is known, we can use it to generate code with more precise
---   kind information. See @Note [Local variables and kind information]@.
+-- * An optional @'lvKind' :: 'Maybe' 'DKind'@. When the kind of a local
+--   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
 -- positions in generated code (see @Note [Local variables and kind information]
 -- (Wrinkle: Binding positions versus argument positions)@), and using
 -- 'DTyVarBndrUnit's to represent type arguments feels somewhat awkward.
-type LocalVar = (Name, Maybe DKind)
+data LocalVar = LocalVar
+  { lvName :: Name
+  , lvKind :: Maybe DKind
+  } 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
@@ -36,17 +54,13 @@ type LocalVar = (Name, Maybe DKind)
 foldTypeLocalVars :: DType -> [LocalVar] -> DType
 foldTypeLocalVars ty = applyDType ty . map localVarToTypeArg
 
--- | Construct a 'LocalVar' with no kind information.
-localVarNoKind :: Name -> LocalVar
-localVarNoKind nm = (nm, 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
 -- and kind information] (Wrinkle: Binding positions versus argument
 -- positions)@.
 localVarToTvb :: flag -> LocalVar -> DTyVarBndr flag
-localVarToTvb flag (nm, mbKind) =
+localVarToTvb flag (LocalVar { lvName = nm, lvKind = mbKind }) =
   case mbKind of
     Nothing   -> DPlainTV nm flag
     Just kind -> DKindedTV nm flag kind
@@ -56,7 +70,7 @@ localVarToTvb flag (nm, mbKind) =
 -- See @Note [Local variables and kind information] (Wrinkle: Binding positions
 -- versus argument positions)@.
 localVarToType :: LocalVar -> DType
-localVarToType (local_nm, _) = DVarT local_nm
+localVarToType (LocalVar { lvName = local_nm }) = DVarT local_nm
 
 -- | Convert a 'LocalVar' used in an argument position to a 'DTypeArg'. Because
 -- this is used in an argument positions, it will not kind any kind information.
@@ -65,6 +79,13 @@ localVarToType (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]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

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

@@ -453,7 +453,8 @@ noExactTyVars = everywhere go
       = InjectivityAnn (noExactName lhs) (map noExactName rhs)
 
     fix_local_var :: LocalVar -> LocalVar
-    fix_local_var (n, mbKind) = (noExactName n, mbKind)
+    fix_local_var (LocalVar { lvName = n, lvKind =  mbKind })
+      = LocalVar { lvName = noExactName n, lvKind = mbKind }
 
 -- Changes a unique Name with a NameU or NameL namespace to a non-unique Name.
 -- See Note [Pitfalls of NameU/NameL] for why this is useful.