WIP Get rid of `SrcSpan` in the base ASTwip/lhs-token-for-hs-arg

I beleive `HsToken` is the correct thing to use instead.

Progress towards #19623 and #19218

14 files changed, 103 insertions, 83 deletions

diff --git a/compiler/GHC/Hs/Instances.hs b/compiler/GHC/Hs/Instances.hs
index a0c588413b..c4a9f96da8 100644
--- a/compiler/GHC/Hs/Instances.hs
+++ b/compiler/GHC/Hs/Instances.hs
@@ -514,10 +514,10 @@ deriving instance Data thing => Data (HsScaled GhcPs thing)
 deriving instance Data thing => Data (HsScaled GhcRn thing)
 deriving instance Data thing => Data (HsScaled GhcTc thing)
 
-deriving instance (Data a, Data b) => Data (HsArg a b)
--- deriving instance Data (HsArg (Located (HsType GhcPs)) (Located (HsKind GhcPs)))
--- deriving instance Data (HsArg (Located (HsType GhcRn)) (Located (HsKind GhcRn)))
--- deriving instance Data (HsArg (Located (HsType GhcTc)) (Located (HsKind GhcTc)))
+-- deriving instance (DataId p, Data a, Data b) => Data (HsArg p a b)
+deriving instance (Data a, Data b) => Data (HsArg GhcPs a b)
+deriving instance (Data a, Data b) => Data (HsArg GhcRn a b)
+deriving instance (Data a, Data b) => Data (HsArg GhcTc a b)
 
 -- deriving instance (DataIdLR p p) => Data (ConDeclField p)
 deriving instance Data (ConDeclField GhcPs)
diff --git a/compiler/GHC/Hs/Type.hs b/compiler/GHC/Hs/Type.hs
index 770a91b35a..22df100ee4 100644
--- a/compiler/GHC/Hs/Type.hs
+++ b/compiler/GHC/Hs/Type.hs
@@ -1,5 +1,5 @@
-
 {-# LANGUAGE ConstraintKinds #-}
+{-# LANGUAGE DataKinds #-}
 {-# LANGUAGE FlexibleContexts #-}
 {-# LANGUAGE FlexibleInstances #-}
 {-# LANGUAGE DeriveDataTypeable #-}
@@ -312,7 +312,7 @@ type instance XIParamTy        (GhcPass _) = EpAnn [AddEpAnn]
 type instance XStarTy          (GhcPass _) = NoExtField
 type instance XKindSig         (GhcPass _) = EpAnn [AddEpAnn]
 
-type instance XAppKindTy       (GhcPass _) = SrcSpan -- Where the `@` lives
+type instance XAppKindTy       (GhcPass _) = NoExtField
 
 type instance XSpliceTy        GhcPs = NoExtField
 type instance XSpliceTy        GhcRn = HsUntypedSpliceResult (HsType GhcRn)
@@ -489,10 +489,10 @@ mkHsAppTys :: LHsType (GhcPass p) -> [LHsType (GhcPass p)]
            -> LHsType (GhcPass p)
 mkHsAppTys = foldl' mkHsAppTy
 
-mkHsAppKindTy :: XAppKindTy (GhcPass p) -> LHsType (GhcPass p) -> LHsType (GhcPass p)
+mkHsAppKindTy :: LHsType (GhcPass p) -> LHsToken "@" (GhcPass p) -> LHsType (GhcPass p)
               -> LHsType (GhcPass p)
-mkHsAppKindTy ext ty k
-  = addCLocAA ty k (HsAppKindTy ext ty k)
+mkHsAppKindTy ty at k
+  = addCLocAA ty k (HsAppKindTy NoExtField ty at k)
 
 {-
 ************************************************************************
@@ -541,7 +541,7 @@ hsTyGetAppHead_maybe = go
   where
     go (L _ (HsTyVar _ _ ln))          = Just ln
     go (L _ (HsAppTy _ l _))           = go l
-    go (L _ (HsAppKindTy _ t _))       = go t
+    go (L _ (HsAppKindTy _ t _ _))     = go t
     go (L _ (HsOpTy _ _ _ ln _))       = Just ln
     go (L _ (HsParTy _ t))             = go t
     go (L _ (HsKindSig _ t _))         = go t
@@ -553,12 +553,12 @@ hsTyGetAppHead_maybe = go
 lhsTypeArgSrcSpan :: LHsTypeArg (GhcPass pass) -> SrcSpan
 lhsTypeArgSrcSpan arg = case arg of
   HsValArg  tm    -> getLocA tm
-  HsTypeArg at ty -> at `combineSrcSpans` getLocA ty
-  HsArgPar  sp    -> sp
+  HsTypeArg at ty -> tokenSrcSpan (getLoc at) `combineSrcSpans` getLocA ty
+  HsArgPar  sp    -> tokenSrcSpan (getLoc sp)
 
 --------------------------------
 
-numVisibleArgs :: [HsArg tm ty] -> Arity
+numVisibleArgs :: [HsArg p tm ty] -> Arity
 numVisibleArgs = count is_vis
   where is_vis (HsValArg _) = True
         is_vis _            = False
@@ -576,7 +576,7 @@ numVisibleArgs = count is_vis
 -- pprHsArgsApp (++) Infix [HsValArg Char, HsValArg Double, HsVarArg Ordering] = (Char ++ Double) Ordering
 -- @
 pprHsArgsApp :: (OutputableBndr id, Outputable tm, Outputable ty)
-             => id -> LexicalFixity -> [HsArg tm ty] -> SDoc
+             => id -> LexicalFixity -> [HsArg p tm ty] -> SDoc
 pprHsArgsApp thing fixity (argl:argr:args)
   | Infix <- fixity
   = let pp_op_app = hsep [ ppr_single_hs_arg argl
@@ -591,7 +591,7 @@ pprHsArgsApp thing _fixity args
 
 -- | Pretty-print a prefix identifier to a list of 'HsArg's.
 ppr_hs_args_prefix_app :: (Outputable tm, Outputable ty)
-                        => SDoc -> [HsArg tm ty] -> SDoc
+                        => SDoc -> [HsArg p tm ty] -> SDoc
 ppr_hs_args_prefix_app acc []         = acc
 ppr_hs_args_prefix_app acc (arg:args) =
   case arg of
@@ -601,7 +601,7 @@ ppr_hs_args_prefix_app acc (arg:args) =
 
 -- | Pretty-print an 'HsArg' in isolation.
 ppr_single_hs_arg :: (Outputable tm, Outputable ty)
-                  => HsArg tm ty -> SDoc
+                  => HsArg p tm ty -> SDoc
 ppr_single_hs_arg (HsValArg tm)    = ppr tm
 ppr_single_hs_arg (HsTypeArg _ ty) = char '@' <> ppr ty
 -- GHC shouldn't be constructing ASTs such that this case is ever reached.
@@ -611,10 +611,10 @@ ppr_single_hs_arg (HsArgPar{})     = empty
 
 -- | This instance is meant for debug-printing purposes. If you wish to
 -- pretty-print an application of 'HsArg's, use 'pprHsArgsApp' instead.
-instance (Outputable tm, Outputable ty) => Outputable (HsArg tm ty) where
+instance (Outputable tm, Outputable ty) => Outputable (HsArg (GhcPass p) tm ty) where
   ppr (HsValArg tm)     = text "HsValArg"  <+> ppr tm
-  ppr (HsTypeArg sp ty) = text "HsTypeArg" <+> ppr sp <+> ppr ty
-  ppr (HsArgPar sp)     = text "HsArgPar"  <+> ppr sp
+  ppr (HsTypeArg sp ty) = text "HsTypeArg" <+> ppr (getLoc sp) <+> ppr ty
+  ppr (HsArgPar sp)     = text "HsArgPar"  <+> ppr (getLoc sp)
 
 --------------------------------
 
@@ -1180,8 +1180,8 @@ ppr_mono_ty (HsStarTy _ isUni)  = char (if isUni then '★' else '*')
 
 ppr_mono_ty (HsAppTy _ fun_ty arg_ty)
   = hsep [ppr_mono_lty fun_ty, ppr_mono_lty arg_ty]
-ppr_mono_ty (HsAppKindTy _ ty k)
-  = ppr_mono_lty ty <+> char '@' <> ppr_mono_lty k
+ppr_mono_ty (HsAppKindTy _ ty at k)
+  = ppr_mono_lty ty <+> ppr at <> ppr_mono_lty k
 ppr_mono_ty (HsOpTy _ prom ty1 (L _ op) ty2)
   = sep [ ppr_mono_lty ty1
         , sep [pprOccWithTick Infix prom op, ppr_mono_lty ty2 ] ]
@@ -1295,7 +1295,7 @@ lhsTypeHasLeadingPromotionQuote ty
     go (HsWildCardTy{})      = False
     go (HsStarTy{})          = False
     go (HsAppTy _ t _)       = goL t
-    go (HsAppKindTy _ t _)   = goL t
+    go (HsAppKindTy _ t _ _) = goL t
     go (HsParTy{})           = False
     go (HsDocTy _ t _)       = goL t
     go (XHsType{})           = False
diff --git a/compiler/GHC/Hs/Utils.hs b/compiler/GHC/Hs/Utils.hs
index 3e74eea3db..cfec92e601 100644
--- a/compiler/GHC/Hs/Utils.hs
+++ b/compiler/GHC/Hs/Utils.hs
@@ -621,13 +621,13 @@ nlHsTyConApp prom fixity tycon tys
     mk_app fun@(L _ (HsOpTy {})) arg = mk_app (noLocA $ HsParTy noAnn fun) arg
       -- parenthesize things like `(A + B) C`
     mk_app fun (HsValArg ty) = noLocA (HsAppTy noExtField fun (parenthesizeHsType appPrec ty))
-    mk_app fun (HsTypeArg _ ki) = noLocA (HsAppKindTy noSrcSpan fun (parenthesizeHsType appPrec ki))
+    mk_app fun (HsTypeArg _ ki) = noLocA (HsAppKindTy noExtField fun noHsTok (parenthesizeHsType appPrec ki))
     mk_app fun (HsArgPar _) = noLocA (HsParTy noAnn fun)
 
 nlHsAppKindTy ::
   LHsType (GhcPass p) -> LHsKind (GhcPass p) -> LHsType (GhcPass p)
 nlHsAppKindTy f k
-  = noLocA (HsAppKindTy noSrcSpan f (parenthesizeHsType appPrec k))
+  = noLocA (HsAppKindTy noExtField f noHsTok (parenthesizeHsType appPrec k))
 
 {-
 Tuples.  All these functions are *pre-typechecker* because they lack
diff --git a/compiler/GHC/HsToCore/Quote.hs b/compiler/GHC/HsToCore/Quote.hs
index ac122446b7..9e58680744 100644
--- a/compiler/GHC/HsToCore/Quote.hs
+++ b/compiler/GHC/HsToCore/Quote.hs
@@ -1375,7 +1375,7 @@ repTy (HsAppTy _ f a)       = do
                                 f1 <- repLTy f
                                 a1 <- repLTy a
                                 repTapp f1 a1
-repTy (HsAppKindTy _ ty ki) = do
+repTy (HsAppKindTy _ ty _ ki) = do
                                 ty1 <- repLTy ty
                                 ki1 <- repLTy ki
                                 repTappKind ty1 ki1
diff --git a/compiler/GHC/Iface/Ext/Ast.hs b/compiler/GHC/Iface/Ext/Ast.hs
index 64eac53af0..1b204cb448 100644
--- a/compiler/GHC/Iface/Ext/Ast.hs
+++ b/compiler/GHC/Iface/Ext/Ast.hs
@@ -552,6 +552,9 @@ instance HasLoc (LocatedA a) where
 instance HasLoc (LocatedN a) where
   loc (L la _) = locA la
 
+instance HasLoc (GenLocated TokenLocation a) where
+  loc (L tl _) = tokenSrcSpan tl
+
 instance HasLoc a => HasLoc [a] where
   loc [] = noSrcSpan
   loc xs = foldl1' combineSrcSpans $ map loc xs
@@ -563,10 +566,10 @@ instance (HasLoc a, HiePass p) => HasLoc (FamEqn (GhcPass p) a) where
     HsOuterExplicit{hso_bndrs = tvs} ->
       foldl1' combineSrcSpans [loc a, loc tvs, loc b, loc c]
 
-instance (HasLoc tm, HasLoc ty) => HasLoc (HsArg tm ty) where
+instance (HiePass p, HasLoc tm, HasLoc ty) => HasLoc (HsArg (GhcPass p) tm ty) where
   loc (HsValArg tm) = loc tm
   loc (HsTypeArg _ ty) = loc ty
-  loc (HsArgPar sp)  = sp
+  loc (HsArgPar sp)  = loc sp
 
 instance HasLoc (HsDataDefn GhcRn) where
   loc def@(HsDataDefn{}) = loc $ dd_cons def
@@ -595,6 +598,9 @@ instance (ToHie a) => ToHie (Bag a) where
 instance (ToHie a) => ToHie (Maybe a) where
   toHie = maybe (pure []) toHie
 
+instance ToHie (GenLocated TokenLocation (HsToken sym)) where
+  toHie = locOnly . loc
+
 instance ToHie (IEContext (LocatedA ModuleName)) where
   toHie (IEC c (L (SrcSpanAnn _ (RealSrcSpan span _)) mname)) = do
       org <- ask
@@ -1760,8 +1766,9 @@ instance ToHie (LocatedA (HsType GhcRn)) where
         [ toHie a
         , toHie b
         ]
-      HsAppKindTy _ ty ki ->
+      HsAppKindTy _ ty at ki ->
         [ toHie ty
+        , toHie at
         , toHie ki
         ]
       HsFunTy _ w a b ->
@@ -1818,10 +1825,10 @@ instance ToHie (LocatedA (HsType GhcRn)) where
       HsStarTy _ _ -> []
       XHsType _ -> []
 
-instance (ToHie tm, ToHie ty) => ToHie (HsArg tm ty) where
+instance (ToHie tm, ToHie ty) => ToHie (HsArg GhcRn tm ty) where
   toHie (HsValArg tm) = toHie tm
   toHie (HsTypeArg _ ty) = toHie ty
-  toHie (HsArgPar sp) = locOnly sp
+  toHie (HsArgPar sp) = toHie sp
 
 instance Data flag => ToHie (TVScoped (LocatedA (HsTyVarBndr flag GhcRn))) where
   toHie (TVS tsc sc (L span bndr)) = concatM $ makeNodeA bndr span : case bndr of
diff --git a/compiler/GHC/Parser.y b/compiler/GHC/Parser.y
index 904f566458..0595115ea4 100644
--- a/compiler/GHC/Parser.y
+++ b/compiler/GHC/Parser.y
@@ -2190,7 +2190,7 @@ ftype :: { forall b. DisambTD b => PV (LocatedA b) }
         | ftype tyarg                   { $1 >>= \ $1 ->
                                           mkHsAppTyPV $1 $2 }
         | ftype PREFIX_AT atype         { $1 >>= \ $1 ->
-                                          mkHsAppKindTyPV $1 (getLoc $2) $3 }
+                                          mkHsAppKindTyPV $1 (hsTok $2) $3 }
 
 tyarg :: { LHsType GhcPs }
         : atype                         { $1 }
diff --git a/compiler/GHC/Parser/Annotation.hs b/compiler/GHC/Parser/Annotation.hs
index 2f00422f8b..fc272a10c8 100644
--- a/compiler/GHC/Parser/Annotation.hs
+++ b/compiler/GHC/Parser/Annotation.hs
@@ -15,6 +15,7 @@ module GHC.Parser.Annotation (
   AddEpAnn(..),
   EpaLocation(..), epaLocationRealSrcSpan, epaLocationFromSrcAnn,
   TokenLocation(..),
+  mkTokenLocation, tokenSrcSpan,
   DeltaPos(..), deltaPos, getDeltaLine,
 
   EpAnn(..), Anchor(..), AnchorOperation(..),
@@ -413,6 +414,15 @@ data EpaLocation = EpaSpan !RealSrcSpan
 data TokenLocation = NoTokenLoc | TokenLoc !EpaLocation
                deriving (Data,Eq)
 
+mkTokenLocation :: SrcSpan -> TokenLocation
+mkTokenLocation (UnhelpfulSpan _) = NoTokenLoc
+mkTokenLocation (RealSrcSpan r _)  = TokenLoc (EpaSpan r)
+
+tokenSrcSpan :: TokenLocation -> SrcSpan
+tokenSrcSpan NoTokenLoc = UnhelpfulSpan UnhelpfulNoLocationInfo -- TODO reason in TokenLocation?
+tokenSrcSpan (TokenLoc (EpaSpan r)) = RealSrcSpan r Strict.Nothing
+tokenSrcSpan (TokenLoc _          ) = error "Not yet handled"
+
 instance Outputable a => Outputable (GenLocated TokenLocation a) where
   ppr (L _ x) = ppr x
 
@@ -462,6 +472,10 @@ instance Outputable EpaLocation where
 instance Outputable AddEpAnn where
   ppr (AddEpAnn kw ss) = text "AddEpAnn" <+> ppr kw <+> ppr ss
 
+instance Outputable TokenLocation where
+  ppr NoTokenLoc = text "NoTokenLoc"
+  ppr (TokenLoc e) = text "TokenLoc" <+> ppr e
+
 -- ---------------------------------------------------------------------
 
 -- | The exact print annotations (EPAs) are kept in the HsSyn AST for
diff --git a/compiler/GHC/Parser/PostProcess.hs b/compiler/GHC/Parser/PostProcess.hs
index 9cce37e051..2eec26852e 100644
--- a/compiler/GHC/Parser/PostProcess.hs
+++ b/compiler/GHC/Parser/PostProcess.hs
@@ -77,9 +77,6 @@ module GHC.Parser.PostProcess (
         UnpackednessPragma(..),
         mkMultTy,
 
-        -- Token location
-        mkTokenLocation,
-
         -- Help with processing exports
         ImpExpSubSpec(..),
         ImpExpQcSpec(..),
@@ -891,7 +888,7 @@ checkTyVars pp_what equals_or_where tc tparms
     check (HsTypeArg _ ki@(L loc _)) = addFatalError $ mkPlainErrorMsgEnvelope (locA loc) $
                                          (PsErrUnexpectedTypeAppInDecl ki pp_what (unLoc tc))
     check (HsValArg ty) = chkParens [] [] emptyComments ty
-    check (HsArgPar sp) = addFatalError $ mkPlainErrorMsgEnvelope sp $
+    check (HsArgPar sp) = addFatalError $ mkPlainErrorMsgEnvelope (tokenSrcSpan $ getLoc sp) $
                             (PsErrMalformedDecl pp_what (unLoc tc))
         -- Keep around an action for adjusting the annotations of extra parens
     chkParens :: [AddEpAnn] -> [AddEpAnn] -> EpAnnComments -> LHsType GhcPs
@@ -1019,7 +1016,7 @@ checkTyClHdr is_cls ty
       where
         (o,c) = mkParensEpAnn (realSrcSpan l)
     go _ (HsAppTy _ t1 t2) acc ops cps fix = goL t1 (HsValArg t2:acc) ops cps fix
-    go _ (HsAppKindTy l ty ki) acc ops cps fix = goL ty (HsTypeArg l ki:acc) ops cps fix
+    go _ (HsAppKindTy _ ty at ki) acc ops cps fix = goL ty (HsTypeArg at ki:acc) ops cps fix
     go l (HsTupleTy _ HsBoxedOrConstraintTuple ts) [] ops cps fix
       = return (L (noAnnSrcSpan l) (nameRdrName tup_name)
                , map HsValArg ts, fix, (reverse ops)++cps)
@@ -1956,7 +1953,7 @@ class DisambTD b where
   -- | Disambiguate @f x@ (function application or prefix data constructor).
   mkHsAppTyPV :: LocatedA b -> LHsType GhcPs -> PV (LocatedA b)
   -- | Disambiguate @f \@t@ (visible kind application)
-  mkHsAppKindTyPV :: LocatedA b -> SrcSpan -> LHsType GhcPs -> PV (LocatedA b)
+  mkHsAppKindTyPV :: LocatedA b -> LHsToken "@" GhcPs -> LHsType GhcPs -> PV (LocatedA b)
   -- | Disambiguate @f \# x@ (infix operator)
   mkHsOpTyPV :: PromotionFlag -> LHsType GhcPs -> LocatedN RdrName -> LHsType GhcPs -> PV (LocatedA b)
   -- | Disambiguate @{-\# UNPACK \#-} t@ (unpack/nounpack pragma)
@@ -1965,7 +1962,7 @@ class DisambTD b where
 instance DisambTD (HsType GhcPs) where
   mkHsAppTyHeadPV = return
   mkHsAppTyPV t1 t2 = return (mkHsAppTy t1 t2)
-  mkHsAppKindTyPV t l_at ki = return (mkHsAppKindTy l_at t ki)
+  mkHsAppKindTyPV t at ki = return (mkHsAppKindTy t at ki)
   mkHsOpTyPV prom t1 op t2 = return (mkLHsOpTy prom t1 op t2)
   mkUnpackednessPV = addUnpackednessP
 
@@ -2001,8 +1998,8 @@ instance DisambTD DataConBuilder where
     -- the grammar in Parser.y is written (see infixtype/ftype).
     panic "mkHsAppTyPV: InfixDataConBuilder"
 
-  mkHsAppKindTyPV lhs l_at ki =
-    addFatalError $ mkPlainErrorMsgEnvelope l_at $
+  mkHsAppKindTyPV lhs at ki =
+    addFatalError $ mkPlainErrorMsgEnvelope (tokenSrcSpan $ getLoc at) $
                       (PsErrUnexpectedKindAppInDataCon (unLoc lhs) (unLoc ki))
 
   mkHsOpTyPV prom lhs tc rhs = do
@@ -3102,10 +3099,6 @@ mkMultTy pct t@(L _ (HsTyLit _ (HsNumTy (SourceText "1") 1))) arr
     locOfPct1 = token_location_widenR (getLoc pct) (locA (getLoc t))
 mkMultTy pct t arr = HsExplicitMult pct t arr
 
-mkTokenLocation :: SrcSpan -> TokenLocation
-mkTokenLocation (UnhelpfulSpan _) = NoTokenLoc
-mkTokenLocation (RealSrcSpan r _)  = TokenLoc (EpaSpan r)
-
 -- Precondition: the TokenLocation has EpaSpan, never EpaDelta.
 token_location_widenR :: TokenLocation -> SrcSpan -> TokenLocation
 token_location_widenR NoTokenLoc _ = NoTokenLoc
diff --git a/compiler/GHC/Rename/HsType.hs b/compiler/GHC/Rename/HsType.hs
index 8a9fdf6542..c7b3961a7a 100644
--- a/compiler/GHC/Rename/HsType.hs
+++ b/compiler/GHC/Rename/HsType.hs
@@ -724,12 +724,12 @@ rnHsTyKi env (HsAppTy _ ty1 ty2)
        ; (ty2', fvs2) <- rnLHsTyKi env ty2
        ; return (HsAppTy noExtField ty1' ty2', fvs1 `plusFV` fvs2) }
 
-rnHsTyKi env (HsAppKindTy l ty k)
+rnHsTyKi env (HsAppKindTy l ty at k)
   = do { kind_app <- xoptM LangExt.TypeApplications
        ; unless kind_app (addErr (typeAppErr "kind" k))
        ; (ty', fvs1) <- rnLHsTyKi env ty
        ; (k', fvs2) <- rnLHsTyKi (env {rtke_level = KindLevel }) k
-       ; return (HsAppKindTy l ty' k', fvs1 `plusFV` fvs2) }
+       ; return (HsAppKindTy l ty' at k', fvs1 `plusFV` fvs2) }
 
 rnHsTyKi env t@(HsIParamTy x n ty)
   = do { notInKinds env t
@@ -1957,7 +1957,7 @@ extract_lty (L _ ty) acc
                                            flds
       HsAppTy _ ty1 ty2           -> extract_lty ty1 $
                                      extract_lty ty2 acc
-      HsAppKindTy _ ty k          -> extract_lty ty $
+      HsAppKindTy _ ty _ k        -> extract_lty ty $
                                      extract_lty k acc
       HsListTy _ ty               -> extract_lty ty acc
       HsTupleTy _ _ tys           -> extract_ltys tys acc
diff --git a/compiler/GHC/Tc/Gen/HsType.hs b/compiler/GHC/Tc/Gen/HsType.hs
index 3fed598f4d..e3fef98939 100644
--- a/compiler/GHC/Tc/Gen/HsType.hs
+++ b/compiler/GHC/Tc/Gen/HsType.hs
@@ -1508,12 +1508,13 @@ splitHsAppTys hs_ty
     is_app _                       = False
 
     go :: LHsType GhcRn
-       -> [HsArg (LHsType GhcRn) (LHsKind GhcRn)]
+       -> [HsArg GhcRn (LHsType GhcRn) (LHsKind GhcRn)]
        -> (LHsType GhcRn,
-           [HsArg (LHsType GhcRn) (LHsKind GhcRn)]) -- AZ temp
-    go (L _  (HsAppTy _ f a))      as = go f (HsValArg a : as)
-    go (L _  (HsAppKindTy l ty k)) as = go ty (HsTypeArg l k : as)
-    go (L sp (HsParTy _ f))        as = go f (HsArgPar (locA sp) : as)
+           [HsArg GhcRn (LHsType GhcRn) (LHsKind GhcRn)]) -- AZ temp
+    go (L _  (HsAppTy _ f a))         as = go f (HsValArg a : as)
+    go (L _  (HsAppKindTy _ ty at k)) as = go ty (HsTypeArg at k : as)
+    go (L sp (HsParTy _ f))           as = go f (HsArgPar at : as)
+       where at = L (mkTokenLocation $ locA sp) HsTok
     go (L _  (HsOpTy _ prom l op@(L sp _) r)) as
       = ( L (na2la sp) (HsTyVar noAnn prom op)
         , HsValArg l : HsValArg r : as )
@@ -1690,7 +1691,7 @@ tcInferTyApps_nosat mode orig_hs_ty fun orig_hs_args
         substed_fun_ki = substTy subst fun_ki
         hs_ty          = appTypeToArg orig_hs_ty (take (n-1) orig_hs_args)
 
-    n_initial_val_args :: [HsArg tm ty] -> Arity
+    n_initial_val_args :: [HsArg GhcRn tm ty] -> Arity
     -- Count how many leading HsValArgs we have
     n_initial_val_args (HsValArg {} : args) = 1 + n_initial_val_args args
     n_initial_val_args (HsArgPar {} : args) = n_initial_val_args args
@@ -1881,11 +1882,11 @@ unsaturated arguments: see #11246.  Hence doing this in tcInferApps.
 -}
 
 appTypeToArg :: LHsType GhcRn -> [LHsTypeArg GhcRn] -> LHsType GhcRn
-appTypeToArg f []                       = f
-appTypeToArg f (HsValArg arg    : args) = appTypeToArg (mkHsAppTy f arg) args
-appTypeToArg f (HsArgPar _      : args) = appTypeToArg f                 args
-appTypeToArg f (HsTypeArg l arg : args)
-  = appTypeToArg (mkHsAppKindTy l f arg) args
+appTypeToArg f []                        = f
+appTypeToArg f (HsValArg arg     : args) = appTypeToArg (mkHsAppTy f arg) args
+appTypeToArg f (HsArgPar _       : args) = appTypeToArg f                 args
+appTypeToArg f (HsTypeArg at arg : args)
+  = appTypeToArg (mkHsAppKindTy f at arg) args
 
 
 {- *********************************************************************
diff --git a/compiler/GHC/Tc/Gen/Sig.hs b/compiler/GHC/Tc/Gen/Sig.hs
index 4bb9a8038d..d04bf200db 100644
--- a/compiler/GHC/Tc/Gen/Sig.hs
+++ b/compiler/GHC/Tc/Gen/Sig.hs
@@ -297,7 +297,7 @@ no_anon_wc_ty lty = go lty
     go (L _ ty) = case ty of
       HsWildCardTy _                 -> False
       HsAppTy _ ty1 ty2              -> go ty1 && go ty2
-      HsAppKindTy _ ty ki            -> go ty && go ki
+      HsAppKindTy _ ty _ ki          -> go ty && go ki
       HsFunTy _ w ty1 ty2            -> go ty1 && go ty2 && go (arrowToHsType w)
       HsListTy _ ty                  -> go ty
       HsTupleTy _ _ tys              -> gos tys
diff --git a/compiler/GHC/Tc/TyCl.hs b/compiler/GHC/Tc/TyCl.hs
index ee9314e74b..8e8dadce78 100644
--- a/compiler/GHC/Tc/TyCl.hs
+++ b/compiler/GHC/Tc/TyCl.hs
@@ -3028,7 +3028,7 @@ tcTyFamInstEqn fam_tc mb_clsinfo
                               (map (const Nominal) qtvs)
                               (locA loc)) }
 
-checkTyFamInstEqn :: TcTyCon -> Name -> [HsArg tm ty] -> TcM ()
+checkTyFamInstEqn :: TcTyCon -> Name -> [HsArg GhcRn tm ty] -> TcM ()
 checkTyFamInstEqn tc_fam_tc eqn_tc_name hs_pats =
   do { -- Ensure that each equation's type constructor is for the right
        -- type family.  E.g. barf on
diff --git a/compiler/GHC/ThToHs.hs b/compiler/GHC/ThToHs.hs
index 5ba99fe7ac..622ad140d1 100644
--- a/compiler/GHC/ThToHs.hs
+++ b/compiler/GHC/ThToHs.hs
@@ -1784,7 +1784,7 @@ mk_apps head_ty type_args = do
           HsValArg ty  -> do p_ty <- add_parens ty
                              mk_apps (HsAppTy noExtField phead_ty p_ty) args
           HsTypeArg l ki -> do p_ki <- add_parens ki
-                               mk_apps (HsAppKindTy l phead_ty p_ki) args
+                               mk_apps (HsAppKindTy noExtField phead_ty l p_ki) args
           HsArgPar _   -> mk_apps (HsParTy noAnn phead_ty) args
 
   go type_args
@@ -1829,8 +1829,10 @@ split_ty_app ty = go ty []
   where
     go (AppT f a) as' = do { a' <- cvtType a; go f (HsValArg a':as') }
     go (AppKindT ty ki) as' = do { ki' <- cvtKind ki
-                                 ; go ty (HsTypeArg noSrcSpan ki':as') }
-    go (ParensT t) as' = do { loc <- getL; go t (HsArgPar loc: as') }
+                                 ; go ty (HsTypeArg noHsTok ki':as') }
+    go (ParensT t) as' = do { loc <- getL
+                            ; go t $ HsArgPar (L (mkTokenLocation loc) HsTok) : as'
+                            }
     go f as           = return (f,as)
 
 cvtTyLit :: TH.TyLit -> HsTyLit (GhcPass p)
diff --git a/compiler/Language/Haskell/Syntax/Type.hs b/compiler/Language/Haskell/Syntax/Type.hs
index 67bb8eabd3..b9e8ed5d5a 100644
--- a/compiler/Language/Haskell/Syntax/Type.hs
+++ b/compiler/Language/Haskell/Syntax/Type.hs
@@ -62,7 +62,6 @@ import Language.Haskell.Syntax.Extension
 import GHC.Types.Name.Reader ( RdrName )
 import GHC.Core.DataCon( HsSrcBang(..) )
 import GHC.Core.Type (Specificity)
-import GHC.Types.SrcLoc (SrcSpan)
 
 import GHC.Hs.Doc (LHsDoc)
 import GHC.Data.FastString (FastString)
@@ -755,6 +754,7 @@ data HsType pass
 
   | HsAppKindTy         (XAppKindTy pass) -- type level type app
                         (LHsType pass)
+                        (LHsToken "@" pass)
                         (LHsKind pass)
 
   | HsFunTy             (XFunTy pass)
@@ -1178,29 +1178,32 @@ if they correspond to a visible 'forall'.
 -}
 
 -- | Arguments in an expression/type after splitting
-data HsArg tm ty
-  = HsValArg tm   -- Argument is an ordinary expression     (f arg)
-  | HsTypeArg SrcSpan ty -- Argument is a visible type application (f @ty)
-                         -- SrcSpan is location of the `@`
-  | HsArgPar SrcSpan -- See Note [HsArgPar]
+--
+-- A HsArgPar indicates that everything to the left of this in the argument list is
+-- enclosed in parentheses together with the function itself. It is necessary so
+-- that we can recreate the parenthesis structure in the original source after
+-- typechecking the arguments.
+--
+-- The SrcSpan is the span of the original HsPar
+--
+-- @((f arg1) arg2 arg3)@ results in an input argument list of
+-- @[HsValArg arg1, HsArgPar span1, HsValArg arg2, HsValArg arg3, HsArgPar span2]@
+data HsArg pass tm ty
 
--- type level equivalent
-type LHsTypeArg p = HsArg (LHsType p) (LHsKind p)
+  -- | Argument is an ordinary expression     (f arg)
+  = HsValArg tm
 
-{-
-Note [HsArgPar]
-~~~~~~~~~~~~~~~
-A HsArgPar indicates that everything to the left of this in the argument list is
-enclosed in parentheses together with the function itself. It is necessary so
-that we can recreate the parenthesis structure in the original source after
-typechecking the arguments.
-
-The SrcSpan is the span of the original HsPar
+  -- | Argument is a visible type application (f @ty)
+  | HsTypeArg (LHsToken "@" pass) ty
 
-((f arg1) arg2 arg3) results in an input argument list of
-[HsValArg arg1, HsArgPar span1, HsValArg arg2, HsValArg arg3, HsArgPar span2]
+  -- | A closing paren.
+  --
+  -- The correponding opening parens are all at the front, so there is
+  -- no ambiguity from just storing the closing one.
+  | HsArgPar (LHsToken ")" pass)
 
--}
+-- type level equivalent
+type LHsTypeArg p = HsArg p (LHsType p) (LHsKind p)
 
 
 {-