on master help sign in
Refactor splitOnSuffixSeq

* Deduplicate some code * Reformat some parts esp. multi-way if statements * Change branches in some cases to keep the likely branch first

commitfa492b8029b40b6d5d74af1b96ebbc3cc4dc63c1 authorHarendra Kumar <harendra@composewell.com> dateThu 12 Nov 2020, 04:09 (6 years ago) parente399354
1 file changed+75−147
src/Streamly/Internal/Data/Stream/StreamD.hs+75−147
src/Streamly/Internal/Data/Stream/StreamD.hs+75−147
@@ -1853,7 +1853,7 @@ splitOnSeq patArr (Fold fstep initial done) (Stream step state) =
18531853 else if patLen == 1
18541854 then do
18551855 acc <- initial
1856− let !pat = A.unsafeIndex patArr 0
1856+ pat <- liftIO $ A.unsafeIndexIO patArr 0
18571857 return $ Skip $ SplitOnSeqSingle acc state pat
18581858 else if sizeOf (undefined :: a) * patLen
18591859 <= sizeOf (undefined :: Word)
@@ -2109,6 +2109,19 @@ splitOnSuffixSeq withSep patArr (Fold fstep initial done) (Stream step state) =
21092109
21102110 addToWord wd a = (wd `shiftL` elemBits) .|. fromIntegral (fromEnum a)
21112111
2112+ -- For single element pattern case
2113+ {-# INLINE processYieldSingle #-}
2114+ processYieldSingle pat x s fs =
2115+ if pat == x
2116+ then do
2117+ fs1 <- if withSep then fstep fs x else return fs
2118+ r <- done fs1
2119+ let next = SplitOnSuffixSeqSingleInit s pat
2120+ skip $ SplitOnSuffixSeqYield r next
2121+ else do
2122+ fs1 <- fstep fs x
2123+ skip $ SplitOnSuffixSeqSingle fs1 s pat
2124+
21122125 -- For Rabin-Karp search
21132126 k = 2891336453 :: Word32
21142127 coeff = k ^ patLen
@@ -2126,14 +2139,14 @@ splitOnSuffixSeq withSep patArr (Fold fstep initial done) (Stream step state) =
21262139 {-# INLINE_LATE stepOuter #-}
21272140 stepOuter _ SplitOnSuffixSeqInit =
21282141 if patLen == 0
2129− then return $ Skip $ SplitOnSuffixSeqEmpty state
2142+ then skip $ SplitOnSuffixSeqEmpty state
21302143 else if patLen == 1
21312144 then do
2132− let !pat = A.unsafeIndex patArr 0
2133− return $ Skip $ SplitOnSuffixSeqSingleInit state pat
2145+ pat <- liftIO $ A.unsafeIndexIO patArr 0
2146+ skip $ SplitOnSuffixSeqSingleInit state pat
21342147 else if sizeOf (undefined :: a) * patLen
21352148 <= sizeOf (undefined :: Word)
2136− then return $ Skip $ SplitOnSuffixSeqWordInit state
2149+ then skip $ SplitOnSuffixSeqWordInit state
21372150 else do
21382151 (rb, rhead) <- liftIO $ RB.new patLen
21392152 skip $ SplitOnSuffixSeqKRInit 0 state rb rhead
@@ -2152,7 +2165,7 @@ splitOnSuffixSeq withSep patArr (Fold fstep initial done) (Stream step state) =
21522165 acc1 <- fstep acc x
21532166 r <- done acc1
21542167 skip $ SplitOnSuffixSeqYield r (SplitOnSuffixSeqEmpty s)
2155− Skip s -> return $ Skip (SplitOnSuffixSeqEmpty s)
2168+ Skip s -> skip (SplitOnSuffixSeqEmpty s)
21562169 Stop -> return Stop
21572170
21582171 -----------------
@@ -2168,49 +2181,18 @@ splitOnSuffixSeq withSep patArr (Fold fstep initial done) (Stream step state) =
21682181 stepOuter gst (SplitOnSuffixSeqSingleInit st pat) = do
21692182 res <- step (adaptState gst) st
21702183 case res of
2171− Yield x s -> do
2172− fs <- initial
2173− if pat == x
2174− then do
2175− if withSep
2176− then do
2177− fs1 <- fstep fs x
2178− r <- done fs1
2179− let next = SplitOnSuffixSeqSingleInit s pat
2180− skip $ SplitOnSuffixSeqYield r next
2181− else do
2182− r <- done fs
2183− let next = SplitOnSuffixSeqSingleInit s pat
2184− skip $ SplitOnSuffixSeqYield r next
2185− else do
2186− fs1 <- fstep fs x
2187− skip $ SplitOnSuffixSeqSingle fs1 s pat
2188− Skip s -> return $ Skip $ (SplitOnSuffixSeqSingleInit s pat)
2184+ Yield x s -> initial >>= processYieldSingle pat x s
2185+ Skip s -> skip $ SplitOnSuffixSeqSingleInit s pat
21892186 Stop -> return Stop
21902187
21912188 stepOuter gst (SplitOnSuffixSeqSingle fs st pat) = do
21922189 res <- step (adaptState gst) st
21932190 case res of
2194− Yield x s -> do
2195− if pat == x
2196− then do
2197− if withSep
2198− then do
2199− fs1 <- fstep fs x
2200− r <- done fs1
2201− let next = SplitOnSuffixSeqSingleInit s pat
2202− skip $ SplitOnSuffixSeqYield r next
2203− else do
2204− r <- done fs
2205− let next = SplitOnSuffixSeqSingleInit s pat
2206− skip $ SplitOnSuffixSeqYield r next
2207− else do
2208− fs1 <- fstep fs x
2209− return $ Skip $ (SplitOnSuffixSeqSingle fs1 s pat)
2210− Skip s -> return $ Skip $ SplitOnSuffixSeqSingle fs s pat
2191+ Yield x s -> processYieldSingle pat x s fs
2192+ Skip s -> skip $ SplitOnSuffixSeqSingle fs s pat
22112193 Stop -> do
22122194 r <- done fs
2213− return $ Skip $ SplitOnSuffixSeqYield r SplitOnSuffixSeqDone
2195+ skip $ SplitOnSuffixSeqYield r SplitOnSuffixSeqDone
22142196
22152197 ---------------------------
22162198 -- Short Pattern - Shift Or
@@ -2230,12 +2212,9 @@ splitOnSuffixSeq withSep patArr (Fold fstep initial done) (Stream step state) =
22302212 Yield x s -> do
22312213 fs <- initial
22322214 let wrd = addToWord 0 x
2233− if withSep
2234− then do
2235− fs1 <- fstep fs x
2236− go SPEC 1 wrd s fs1
2237− else go SPEC 1 wrd s fs
2238− Skip s -> return $ Skip (SplitOnSuffixSeqWordInit s)
2215+ fs1 <- if withSep then fstep fs x else return fs
2216+ go SPEC 1 wrd s fs1
2217+ Skip s -> skip (SplitOnSuffixSeqWordInit s)
22392218 Stop -> return Stop
22402219
22412220 where
@@ -2246,28 +2225,15 @@ splitOnSuffixSeq withSep patArr (Fold fstep initial done) (Stream step state) =
22462225 case res of
22472226 Yield x s -> do
22482227 let wrd1 = addToWord wrd x
2249− if withSep
2250− then do
2251− fs1 <- fstep fs x
2252− if idx == maxIndex
2253− then do
2254− if wrd1 .&. wordMask == wordPat
2255− then do
2256− r <- done fs
2257− let next = SplitOnSuffixSeqWordInit s
2258− skip $ SplitOnSuffixSeqYield r next
2259− else skip $ SplitOnSuffixSeqWordLoop wrd1 s fs1
2260− else go SPEC (idx + 1) wrd1 s fs1
2228+ fs1 <- if withSep then fstep fs x else return fs
2229+ if idx /= maxIndex
2230+ then go SPEC (idx + 1) wrd1 s fs1
2231+ else if wrd1 .&. wordMask /= wordPat
2232+ then skip $ SplitOnSuffixSeqWordLoop wrd1 s fs1
22612233 else do
2262− if idx == maxIndex
2263− then do
2264− if wrd1 .&. wordMask == wordPat
2265− then do
2266− r <- done fs
2267− let next = SplitOnSuffixSeqWordInit s
2268− skip $ SplitOnSuffixSeqYield r next
2269− else skip $ SplitOnSuffixSeqWordLoop wrd1 s fs
2270− else go SPEC (idx + 1) wrd1 s fs
2234+ r <- done fs
2235+ let next = SplitOnSuffixSeqWordInit s
2236+ skip $ SplitOnSuffixSeqYield r next
22712237 Skip s -> go SPEC idx wrd s fs
22722238 Stop -> skip $ SplitOnSuffixSeqWordDone idx fs wrd
22732239
@@ -2284,33 +2250,25 @@ splitOnSuffixSeq withSep patArr (Fold fstep initial done) (Stream step state) =
22842250 let wrd1 = addToWord wrd x
22852251 old = (wordMask .&. wrd)
22862252 `shiftR` (elemBits * (patLen - 1))
2287− if withSep
2253+ fs1 <-
2254+ if withSep
2255+ then fstep fs x
2256+ else fstep fs (toEnum $ fromIntegral old)
2257+ if wrd1 .&. wordMask == wordPat
22882258 then do
2289− fs1 <- fstep fs x
2290− if wrd1 .&. wordMask == wordPat
2291− then do
2292− r <- done fs1
2293− let next = SplitOnSuffixSeqWordInit s
2294− skip $ SplitOnSuffixSeqYield r next
2295− else go SPEC wrd1 s fs1
2296− else do
2297− fs1 <- fstep fs (toEnum $ fromIntegral old)
2298− if wrd1 .&. wordMask == wordPat
2299− then do
2300− r <- done fs1
2301− let next = SplitOnSuffixSeqWordInit s
2302− skip $ SplitOnSuffixSeqYield r next
2303− else go SPEC wrd1 s fs1
2259+ r <- done fs1
2260+ let next = SplitOnSuffixSeqWordInit s
2261+ skip $ SplitOnSuffixSeqYield r next
2262+ else go SPEC wrd1 s fs1
23042263 Skip s -> go SPEC wrd s fs
23052264 Stop ->
23062265 if wrd .&. wordMask == wordPat
23072266 then return Stop
2308− else
2309− if withSep
2310− then do
2311− r <- done fs
2312− skip $ SplitOnSuffixSeqYield r SplitOnSuffixSeqDone
2313− else skip $ SplitOnSuffixSeqWordDone patLen fs wrd
2267+ else if withSep
2268+ then do
2269+ r <- done fs
2270+ skip $ SplitOnSuffixSeqYield r SplitOnSuffixSeqDone
2271+ else skip $ SplitOnSuffixSeqWordDone patLen fs wrd
23142272
23152273 -------------------------------
23162274 -- General Pattern - Karp Rabin
@@ -2322,16 +2280,13 @@ splitOnSuffixSeq withSep patArr (Fold fstep initial done) (Stream step state) =
23222280 Yield x s -> do
23232281 rh1 <- liftIO $ RB.unsafeInsert rb rh0 x
23242282 fs <- initial
2325− if withSep
2326− then do
2327− fs1 <- fstep fs x
2328− skip $ SplitOnSuffixSeqKRInit1 fs1 s rb rh1
2329− else skip $ SplitOnSuffixSeqKRInit1 fs s rb rh1
2283+ fs1 <- if withSep then fstep fs x else return fs
2284+ skip $ SplitOnSuffixSeqKRInit1 fs1 s rb rh1
23302285 Skip s -> skip $ SplitOnSuffixSeqKRInit idx0 s rb rh0
23312286 Stop -> return Stop
23322287
23332288 stepOuter gst (SplitOnSuffixSeqKRInit1 fs0 st0 rb rh0) = do
2334− go SPEC (1 :: Int) rh0 st0 fs0
2289+ go SPEC 1 rh0 st0 fs0
23352290
23362291 where
23372292
@@ -2340,45 +2295,26 @@ splitOnSuffixSeq withSep patArr (Fold fstep initial done) (Stream step state) =
23402295 case res of
23412296 Yield x s -> do
23422297 rh1 <- liftIO (RB.unsafeInsert rb rh x)
2343− if withSep
2344− then do
2345− fs1 <- fstep fs x
2346− if idx == maxIndex
2347− then do
2348− let fold = RB.unsafeFoldRing (RB.ringBound rb)
2349− let !ringHash = fold addCksum 0 rb
2350− if ringHash == patHash
2351− then skip $ SplitOnSuffixSeqKRCheck fs1 s rb rh1
2352− else skip $ SplitOnSuffixSeqKRLoop
2353− fs1 s rb rh1 ringHash
2354− else go SPEC (idx + 1) rh1 s fs1
2355− else
2356− if idx == maxIndex
2357− then do
2358− let fold = RB.unsafeFoldRing (RB.ringBound rb)
2359− let !ringHash = fold addCksum 0 rb
2360− if ringHash == patHash
2361− then skip $ SplitOnSuffixSeqKRCheck fs s rb rh1
2362− else skip $ SplitOnSuffixSeqKRLoop
2363− fs s rb rh1 ringHash
2364− else go SPEC (idx + 1) rh1 s fs
2298+ fs1 <- if withSep then fstep fs x else return fs
2299+ if idx /= maxIndex
2300+ then go SPEC (idx + 1) rh1 s fs1
2301+ else skip $
2302+ let fold = RB.unsafeFoldRing (RB.ringBound rb)
2303+ !ringHash = fold addCksum 0 rb
2304+ in if ringHash == patHash
2305+ then SplitOnSuffixSeqKRCheck fs1 s rb rh1
2306+ else SplitOnSuffixSeqKRLoop fs1 s rb rh1 ringHash
23652307 Skip s -> go SPEC idx rh s fs
23662308 Stop -> do
23672309 -- do not issue a blank segment when we end at pattern
23682310 if (idx == maxIndex) && RB.unsafeEqArray rb rh patArr
23692311 then return Stop
2370− else
2371− if withSep
2372− then do
2373− r <- done fs
2374− skip $ SplitOnSuffixSeqYield r SplitOnSuffixSeqDone
2375− else skip $ SplitOnSuffixSeqKRDone
2376− idx fs rb (RB.startOf rb)
2312+ else if withSep
2313+ then do
2314+ r <- done fs
2315+ skip $ SplitOnSuffixSeqYield r SplitOnSuffixSeqDone
2316+ else skip $ SplitOnSuffixSeqKRDone idx fs rb (RB.startOf rb)
23772317
2378− -- XXX The recursive "go" is more efficient than the state based recursion
2379− -- code commented out below. Perhaps its more efficient because of
2380− -- factoring out "rb" outside the loop.
2381− --
23822318 stepOuter gst (SplitOnSuffixSeqKRLoop fs0 st0 rb rh0 cksum0) =
23832319 go SPEC fs0 st0 rh0 cksum0
23842320
@@ -2391,27 +2327,19 @@ splitOnSuffixSeq withSep patArr (Fold fstep initial done) (Stream step state) =
23912327 old <- liftIO $ peek rh
23922328 rh1 <- liftIO (RB.unsafeInsert rb rh x)
23932329 let cksum1 = deltaCksum cksum old x
2394− if withSep
2395− then do
2396− fs1 <- fstep fs x
2397− if (cksum1 == patHash)
2398− then skip $ SplitOnSuffixSeqKRCheck fs1 s rb rh1
2399− else go SPEC fs1 s rh1 cksum1
2400− else do
2401− fs1 <- fstep fs old
2402− if (cksum1 == patHash)
2403− then skip $ SplitOnSuffixSeqKRCheck fs1 s rb rh1
2404− else go SPEC fs1 s rh1 cksum1
2330+ fs1 <- if withSep then fstep fs x else fstep fs old
2331+ if (cksum1 /= patHash)
2332+ then go SPEC fs1 s rh1 cksum1
2333+ else skip $ SplitOnSuffixSeqKRCheck fs1 s rb rh1
24052334 Skip s -> go SPEC fs s rh cksum
24062335 Stop ->
24072336 if RB.unsafeEqArray rb rh patArr
24082337 then return Stop
2409− else
2410− if withSep
2411− then do
2412− r <- done fs
2413− skip $ SplitOnSuffixSeqYield r SplitOnSuffixSeqDone
2414− else skip $ SplitOnSuffixSeqKRDone patLen fs rb rh
2338+ else if withSep
2339+ then do
2340+ r <- done fs
2341+ skip $ SplitOnSuffixSeqYield r SplitOnSuffixSeqDone
2342+ else skip $ SplitOnSuffixSeqKRDone patLen fs rb rh
24152343
24162344 stepOuter _ (SplitOnSuffixSeqKRCheck fs st rb rh) = do
24172345 if RB.unsafeEqArray rb rh patArr