EPA: Capture '+' location for NPlusKPatwip/az/T20243-n-plus-k-patterns

The location of the plus symbol was being discarded, we now capture
it.

Closes #20243

9 files changed, 74 insertions, 8 deletions

diff --git a/compiler/GHC/Hs/Dump.hs b/compiler/GHC/Hs/Dump.hs
index 9be0f96640..247e8099da 100644
--- a/compiler/GHC/Hs/Dump.hs
+++ b/compiler/GHC/Hs/Dump.hs
@@ -62,6 +62,7 @@ showAstData bs ba a0 = blankLine $$ showAstData' a0
               `extQ` annotationEpAnnImportDecl
               `extQ` annotationAnnParen
               `extQ` annotationTrailingAnn
+              `extQ` annotationEpaLocation
               `extQ` addEpAnn
               `extQ` lit `extQ` litr `extQ` litt
               `extQ` sourceText
@@ -254,6 +255,9 @@ showAstData bs ba a0 = blankLine $$ showAstData' a0
             annotationTrailingAnn :: EpAnn TrailingAnn -> SDoc
             annotationTrailingAnn = annotation' (text "EpAnn TrailingAnn")
 
+            annotationEpaLocation :: EpAnn EpaLocation -> SDoc
+            annotationEpaLocation = annotation' (text "EpAnn EpaLocation")
+
             annotation' :: forall a .(Data a, Typeable a)
                        => SDoc -> EpAnn a -> SDoc
             annotation' tag anns = case ba of
diff --git a/compiler/GHC/Hs/Pat.hs b/compiler/GHC/Hs/Pat.hs
index f300c4a2ca..a4b3bed851 100644
--- a/compiler/GHC/Hs/Pat.hs
+++ b/compiler/GHC/Hs/Pat.hs
@@ -147,7 +147,7 @@ type instance XNPat GhcPs = EpAnn [AddEpAnn]
 type instance XNPat GhcRn = EpAnn [AddEpAnn]
 type instance XNPat GhcTc = Type
 
-type instance XNPlusKPat GhcPs = EpAnn [AddEpAnn]
+type instance XNPlusKPat GhcPs = EpAnn EpaLocation -- Of the "+"
 type instance XNPlusKPat GhcRn = NoExtField
 type instance XNPlusKPat GhcTc = Type
 
diff --git a/compiler/GHC/Hs/Utils.hs b/compiler/GHC/Hs/Utils.hs
index 87fc46ff12..ac73720456 100644
--- a/compiler/GHC/Hs/Utils.hs
+++ b/compiler/GHC/Hs/Utils.hs
@@ -309,7 +309,7 @@ mkHsCompAnns   :: HsDoFlavour -> [ExprLStmt GhcPs] -> LHsExpr GhcPs
 
 mkNPat      :: Located (HsOverLit GhcPs) -> Maybe (SyntaxExpr GhcPs) -> EpAnn [AddEpAnn]
             -> Pat GhcPs
-mkNPlusKPat :: LocatedN RdrName -> Located (HsOverLit GhcPs) -> EpAnn [AddEpAnn]
+mkNPlusKPat :: LocatedN RdrName -> Located (HsOverLit GhcPs) -> EpAnn EpaLocation
             -> Pat GhcPs
 
 -- NB: The following functions all use noSyntaxExpr: the generated expressions
diff --git a/compiler/GHC/Parser/Annotation.hs b/compiler/GHC/Parser/Annotation.hs
index 19925b0678..a914a14b71 100644
--- a/compiler/GHC/Parser/Annotation.hs
+++ b/compiler/GHC/Parser/Annotation.hs
@@ -13,7 +13,7 @@ module GHC.Parser.Annotation (
 
   -- * In-tree Exact Print Annotations
   AddEpAnn(..),
-  EpaLocation(..), epaLocationRealSrcSpan,
+  EpaLocation(..), epaLocationRealSrcSpan, epaLocationFromSrcAnn,
   DeltaPos(..), deltaPos, getDeltaLine,
 
   EpAnn(..), Anchor(..), AnchorOperation(..),
@@ -440,6 +440,10 @@ epaLocationRealSrcSpan :: EpaLocation -> RealSrcSpan
 epaLocationRealSrcSpan (EpaSpan r) = r
 epaLocationRealSrcSpan (EpaDelta _) = panic "epaLocationRealSrcSpan"
 
+epaLocationFromSrcAnn :: SrcAnn ann -> EpaLocation
+epaLocationFromSrcAnn (SrcSpanAnn EpAnnNotUsed l) = EpaSpan (realSrcSpan l)
+epaLocationFromSrcAnn (SrcSpanAnn (EpAnn anc _ _) _) = EpaSpan (anchor anc)
+
 instance Outputable EpaLocation where
   ppr (EpaSpan r) = text "EpaSpan" <+> ppr r
   ppr (EpaDelta d) = text "EpaDelta" <+> ppr d
diff --git a/compiler/GHC/Parser/PostProcess.hs b/compiler/GHC/Parser/PostProcess.hs
index 4eab0c1486..d1ec88f7fa 100644
--- a/compiler/GHC/Parser/PostProcess.hs
+++ b/compiler/GHC/Parser/PostProcess.hs
@@ -1146,11 +1146,12 @@ checkAPat loc e0 = do
    -- n+k patterns
    PatBuilderOpApp
            (L _ (PatBuilderVar (L nloc n)))
-           (L _ plus)
+           (L l plus)
            (L lloc (PatBuilderOverLit lit@(OverLit {ol_val = HsIntegral {}})))
-           anns
+           (EpAnn anc _ cs)
                      | nPlusKPatterns && (plus == plus_RDR)
-                     -> return (mkNPlusKPat (L nloc n) (L (locA lloc) lit) anns)
+                     -> return (mkNPlusKPat (L nloc n) (L (locA lloc) lit)
+                                (EpAnn anc (epaLocationFromSrcAnn l) cs))
 
    -- Improve error messages for the @-operator when the user meant an @-pattern
    PatBuilderOpApp _ op _ _ | opIsAt (unLoc op) -> do
diff --git a/testsuite/tests/printer/Makefile b/testsuite/tests/printer/Makefile
index 4d2dad7b86..ef8aa2858a 100644
--- a/testsuite/tests/printer/Makefile
+++ b/testsuite/tests/printer/Makefile
@@ -729,6 +729,12 @@ PprCommentPlacement2:
 	$(CHECK_PPR)   $(LIBDIR) PprCommentPlacement2.hs
 	$(CHECK_EXACT) $(LIBDIR) PprCommentPlacement2.hs
 
+
+.PHONY: Test20243
+Test20243:
+	$(CHECK_PPR)   $(LIBDIR) Test20243.hs
+	$(CHECK_EXACT) $(LIBDIR) Test20243.hs
+
 .PHONY: Test20258
 Test20258:
 	$(CHECK_PPR)   $(LIBDIR) Test20258.hs
diff --git a/testsuite/tests/printer/Test20243.hs b/testsuite/tests/printer/Test20243.hs
new file mode 100644
index 0000000000..b78591f42d
--- /dev/null
+++ b/testsuite/tests/printer/Test20243.hs
@@ -0,0 +1,42 @@
+{-# LANGUAGE NPlusKPatterns #-}
+module Test20243 where
+
+singleline :: Integer -> Integer
+singleline (n + 1) = n
+
+multiline :: Integer -> Integer
+multiline(n
+  + 1) = n
+
+m :: Integer
+(m + 1) = 3
+
+erroR :: Int
+erroR = n where
+                (n+1,_) = (5,2)
+
+g :: Int -> Int
+g (x+1) = x
+g y     = y
+g _     = 0     -- Overlapped
+
+h :: Int -> Int
+h (x+1) = x
+h _     = 0     -- Not overlapped
+
+kh (n+2) x | x > n = x * 2
+kh (x+1) (m+1) = m
+
+takeList :: Int -> [a] -> [a]
+takeList 0     _      = []
+takeList (n+1) []     = []
+takeList (n+1) (x:xs) = x : takeList n xs
+
+(^^^^)          :: (Num a, Integral b) => a -> b -> a
+x ^^^^ 0                =  1
+x ^^^^ (n+1)    =  f x n x
+                   where f _ 0 y = y
+                         f x n y = g x n  where
+                                   g x n | even n  = g (x*x) (n `quot` 2)
+                                         | otherwise = f x (n-1) (x*y)
+_ ^^^^ _                = error "(^^^^){prelude}: negative exponent"
diff --git a/testsuite/tests/printer/all.T b/testsuite/tests/printer/all.T
index 4b9c04c01c..380e71c723 100644
--- a/testsuite/tests/printer/all.T
+++ b/testsuite/tests/printer/all.T
@@ -170,4 +170,5 @@ test('PprBracesSemiDataDecl', ignore_stderr, makefile_test, ['PprBracesSemiDataD
 test('PprUnicodeSyntax', ignore_stderr, makefile_test, ['PprUnicodeSyntax'])
 test('PprCommentPlacement2', ignore_stderr, makefile_test, ['PprCommentPlacement2'])
 
+test('Test20243', ignore_stderr, makefile_test, ['Test20243'])
 test('Test20258', ignore_stderr, makefile_test, ['Test20258'])
diff --git a/utils/check-exact/ExactPrint.hs b/utils/check-exact/ExactPrint.hs
index fc04e24332..29512e84b5 100644
--- a/utils/check-exact/ExactPrint.hs
+++ b/utils/check-exact/ExactPrint.hs
@@ -455,6 +455,9 @@ printStringAtMkw :: Maybe EpaLocation -> String -> EPP ()
 printStringAtMkw (Just aa) s = printStringAtAA aa s
 printStringAtMkw Nothing s = printStringAtLsDelta (SameLine 1) s
 
+printStringAtAnn :: EpAnn a -> (a -> EpaLocation) -> String -> EPP ()
+printStringAtAnn EpAnnNotUsed  _ _  = return ()
+printStringAtAnn (EpAnn _ a _) f str = printStringAtAA (f a) str
 
 printStringAtAA :: EpaLocation -> String -> EPP ()
 printStringAtAA (EpaSpan r) s = printStringAtKw' r s
@@ -3603,12 +3606,17 @@ instance ExactPrint (Pat GhcPs) where
     markAnnotated ol
 
   -- | NPlusKPat an n lit1 lit2 _ _)
+  exact (NPlusKPat an n k _lit2 _ _) = do
+    markAnnotated n
+    printStringAtAnn an id "+"
+    markAnnotated k
+
+
   exact (SigPat an pat sig) = do
     markAnnotated pat
     markEpAnn an AnnDcolon
     markAnnotated sig
-  -- exact x = withPpr x
-  exact x = error $ "missing match for Pat:" ++ showAst x
+  -- exact x = error $ "missing match for Pat:" ++ showAst x
 
 -- ---------------------------------------------------------------------