From 591244a878b6638a5b206133aeb4d4220f175238 Mon Sep 17 00:00:00 2001
From: meooow25 <soumiksarkar.3120@gmail.com>
Date: Sun, 3 Nov 2024 15:56:10 +0530
Subject: [PATCH] Use GHC.Exts.oneShot in parseFoldr

It can improve performance depending on the foldr implementation.
The `foldrTree` benchmark improves from 8.2ms to 3.5ms (-56%).
---
 bench/Bench.hs               | 57 ++++++++++++++++++++++++++++++++++++
 src/Regex/Internal/Parser.hs |  3 +-
 2 files changed, 59 insertions(+), 1 deletion(-)

diff --git a/bench/Bench.hs b/bench/Bench.hs
index 6171141..db592d9 100644
--- a/bench/Bench.hs
+++ b/bench/Bench.hs
@@ -3,6 +3,7 @@
 module Bench (benches) where
 
 import Control.Applicative
+import Control.DeepSeq (NFData(..))
 import Control.Monad
 import Control.Monad.Random.Strict
 import qualified Data.Foldable as F
@@ -18,6 +19,7 @@ import qualified Regex.Base as RB
 benches :: Benchmark
 benches = bgroup "parser-regex"
   [ textBenches
+  , intTreeBenches
   ]
 
 ---------
@@ -176,6 +178,61 @@ bigExpr :: Text
 bigExpr = T.intercalate "+" $ map (T.pack . show) . evalR $
   replicateM 100000 (getRandom :: R Word)
 
+------------
+-- IntTree
+------------
+
+-- These benchmarks are primarily to test that parsing works well with
+-- structures that are not [] or Text, such as trees. The performance depends on
+-- the foldr implementation, so we test two possible implementations here.
+
+intTreeBenches :: Benchmark
+intTreeBenches = bgroup "IntTree"
+  [ env (pure bigTree) $ \data_ ->
+    bgroup "many anySingle"
+    [ bench "foldrTree" $ whnf (treeParse manyAnyP) data_
+    , bench "foldrTreeStack" $ whnf (treeParseStack manyAnyP) data_
+    ]
+  ]
+  where
+    !manyAnyP = RB.compile $ RB.foldlMany' (\_ _ -> ()) () RB.anySingle
+
+data IntTree = Bin !IntTree !Int !IntTree | Tip
+
+instance NFData IntTree where
+  rnf !_ = ()
+
+foldrTree :: (Int -> b -> b) -> b -> IntTree -> b
+foldrTree f z0 t = go t z0
+  where
+    go Tip z = z
+    go (Bin l x r) z = go l (f x (go r z))
+
+treeParse :: RB.Parser Int a -> IntTree -> Maybe a
+treeParse = RB.parseFoldr foldrTree
+
+data Stack = Push !Int !IntTree !Stack | Nada
+
+foldrTreeStack :: (Int -> b -> b) -> b -> IntTree -> b
+foldrTreeStack f z0 t = go (down t Nada)
+  where
+    go Nada = z0
+    go (Push x r stk) = f x (go (down r stk))
+    down Tip stk = stk
+    down (Bin l x r) stk = down l (Push x r stk)
+
+treeParseStack :: RB.Parser Int a -> IntTree -> Maybe a
+treeParseStack = RB.parseFoldr foldrTreeStack
+
+bigTree :: IntTree
+bigTree = go 100000 0
+  where
+    go 0 _ = Tip
+    go n i = Bin (go ln i) (i+ln) (go rn (i+ln+1))
+      where
+        ln = (n-1) `div` 2
+        rn = n-1-ln
+
 -----------
 -- Random
 -----------
diff --git a/src/Regex/Internal/Parser.hs b/src/Regex/Internal/Parser.hs
index 67a1c58..f6941a7 100644
--- a/src/Regex/Internal/Parser.hs
+++ b/src/Regex/Internal/Parser.hs
@@ -22,6 +22,7 @@ import Control.Monad.Fix
 import Data.Maybe (isJust)
 import Data.Primitive.SmallArray
 import qualified Data.Foldable as F
+import qualified GHC.Exts as X
 
 import Regex.Internal.Regex (RE(..), Strictness(..), Greediness(..))
 import Regex.Internal.Unique (Unique(..), UniqueSet)
@@ -368,7 +369,7 @@ type Foldr f a = forall b. (a -> b -> b) -> b -> f -> b
 parseFoldr :: Foldr f c -> Parser c a -> f -> Maybe a
 parseFoldr fr = \p xs -> fr f finishParser xs (prepareParser p)
   where
-    f c k = \ps -> stepParser ps c >>= k
+    f c k = X.oneShot (\ps -> stepParser ps c >>= k)
 {-# INLINE parseFoldr #-}
 
 ---------