From 9698fe325c47e9e3e153962191dec3c1bb56b352 Mon Sep 17 00:00:00 2001
From: Mark Karpov <markkarpov92@gmail.com>
Date: Mon, 3 Sep 2018 22:16:27 +0700
Subject: [PATCH] =?UTF-8?q?Implement=20the=20=E2=80=98attachSourcePos?=
 =?UTF-8?q?=E2=80=99=20function=20(#320)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 CHANGELOG.md                       |  2 ++
 Text/Megaparsec/Error.hs           | 33 ++++++++++++++++++++++++++----
 tests/Text/Megaparsec/ErrorSpec.hs | 19 ++++++++++++++---
 3 files changed, 47 insertions(+), 7 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 69aa258..edd2bdb 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -112,6 +112,8 @@
   removed. `parseTest'` was removed because `parseTest` always prints
   offending lines now.
 
+* Added `attachSourcePos` function in `Text.Megaparsec.Error`.
+
 * The `ShowToken` type class has been removed and its method `showTokens`
   now lives in the `Stream` type class.
 
diff --git a/Text/Megaparsec/Error.hs b/Text/Megaparsec/Error.hs
index 92a345a..52f7688 100644
--- a/Text/Megaparsec/Error.hs
+++ b/Text/Megaparsec/Error.hs
@@ -36,6 +36,7 @@ module Text.Megaparsec.Error
   , mapParseError
   , errorOffset
   , ParseErrorBundle (..)
+  , attachSourcePos
     -- * Pretty-printing
   , ShowErrorComponent (..)
   , errorBundlePretty
@@ -45,6 +46,7 @@ where
 
 import Control.DeepSeq
 import Control.Exception
+import Control.Monad.State.Strict
 import Data.Data (Data)
 import Data.List (intercalate)
 import Data.List.NonEmpty (NonEmpty (..))
@@ -233,7 +235,7 @@ mergeError e1 e2 =
 
 data ParseErrorBundle s e = ParseErrorBundle
   { bundleErrors :: NonEmpty (ParseError s e)
-    -- ^ A /sorted/ collection of 'ParseError's to display
+    -- ^ A collection of 'ParseError's that is sorted by parse error offsets
   , bundlePosState :: PosState s
     -- ^ State that is used for line\/column calculation
   } deriving (Generic)
@@ -275,6 +277,29 @@ instance ( Show s
          ) => Exception (ParseErrorBundle s e) where
   displayException = errorBundlePretty
 
+-- | Attach 'SourcePos'es to items in a 'Traversable' container given that
+-- there is a projection allowing to get an offset per item.
+--
+-- Items must be in ascending order with respect to their offsets.
+--
+-- @since 7.0.0
+
+attachSourcePos
+  :: (Traversable t, Stream s)
+  => (a -> Int)        -- ^ How to project offset from an item
+  -> t a               -- ^ The collection of items
+  -> PosState s        -- ^ Initial 'PosState'
+  -> (t (a, SourcePos), PosState s) -- ^ The collection with 'SourcePos'es
+                                    -- added and the final 'PosState'
+attachSourcePos projectOffset xs = runState (traverse f xs)
+  where
+    f a = do
+      pst <- get
+      let (spos, pst') = reachOffsetNoLine (projectOffset a) pst
+      put pst'
+      return (a, spos)
+{-# INLINEABLE attachSourcePos #-}
+
 ----------------------------------------------------------------------------
 -- Pretty-printing
 
@@ -319,10 +344,10 @@ errorBundlePretty ParseErrorBundle {..} =
     f :: (ShowS, PosState s)
       -> ParseError s e
       -> (ShowS, PosState s)
-    f (o, !pstate) e = (o . (outChunk ++), pstate')
+    f (o, !pst) e = (o . (outChunk ++), pst')
       where
-        (epos, sline, pstate') = reachOffset (errorOffset e) pstate
-        ppos = pstateSourcePos pstate
+        (epos, sline, pst') = reachOffset (errorOffset e) pst
+        ppos = pstateSourcePos pst
         outChunk =
           "\n" <> sourcePosPretty epos <> ":\n" <>
           padding <> "|\n" <>
diff --git a/tests/Text/Megaparsec/ErrorSpec.hs b/tests/Text/Megaparsec/ErrorSpec.hs
index 1217a69..8d7ef5d 100644
--- a/tests/Text/Megaparsec/ErrorSpec.hs
+++ b/tests/Text/Megaparsec/ErrorSpec.hs
@@ -6,14 +6,14 @@ module Text.Megaparsec.ErrorSpec (spec) where
 import Control.Exception (Exception (..))
 import Data.ByteString (ByteString)
 import Data.Functor.Identity
-import Data.List (isSuffixOf, isInfixOf)
+import Data.List (isSuffixOf, isInfixOf, sort)
 import Data.List.NonEmpty (NonEmpty (..))
 import Data.Void
 import Test.Hspec
+import Test.Hspec.Megaparsec
 import Test.Hspec.Megaparsec.AdHoc ()
 import Test.QuickCheck
 import Text.Megaparsec
-import Text.Megaparsec.Error.Builder
 import qualified Data.ByteString    as B
 import qualified Data.Semigroup     as S
 import qualified Data.Set           as E
@@ -69,7 +69,7 @@ spec = do
           TrivialError pos us ps <> FancyError pos xs `shouldBe`
             (FancyError pos xs :: PE)
 
-  describe "errorPos" $
+  describe "errorOffset" $
     it "returns error position" $
       property $ \e ->
         errorOffset e `shouldBe`
@@ -77,6 +77,19 @@ spec = do
             TrivialError o _ _ -> o
             FancyError   o _   -> o)
 
+  describe "attachSourcePos" $
+    it "attaches the positions correctly" $
+      property $ \xs' s -> do
+        let xs = sort $ getSmall . getPositive <$> xs'
+            pst = initialPosState (s :: String)
+            pst' =
+              if null xs
+                then pst
+                else snd $ reachOffsetNoLine (last xs) pst
+            rs = f <$> xs
+            f x = (x, fst (reachOffsetNoLine x pst))
+        attachSourcePos id (xs :: [Int]) pst `shouldBe` (rs, pst')
+
   describe "errorBundlePretty" $ do
     context "with Char tokens" $ do
       it "shows empty line correctly" $ do