git.iko.soy/streamly
on master help sign in
Add array splitOn, getSlicesFromLen

N-ary array slicing operations.

commit8ae28e29fc78707fa7682cc1da1455054bdb64ea authorHarendra Kumar <harendra@composewell.com> dateThu 4 Nov 2021, 00:03 (5 years ago) parentdb6fd11
4 files changed+142−3
compare with master
unifiedsplit
src/Streamly/Internal/Data/Array/Foreign.hs+36−1 src/Streamly/Internal/Data/Array/Foreign/Mut.hs+52 src/Streamly/Internal/Data/Stream/StreamD/Nesting.hs+34−2 src/Streamly/Internal/Data/Stream/StreamD/Transform.hs+20
src/Streamly/Internal/Data/Array/Foreign.hs+36−1
@@ -104,7 +104,9 @@ module Streamly.Internal.Data.Array.Foreign
104104 -- * Subarrays
105105 , getSliceUnsafe
106106 -- , getSlice
107− -- , getSlicesFromLenN
107+ , genSlicesFromLen
108+ , getSlicesFromLen
109+ , splitOn
108110
109111 -- * Streaming Operations
110112 , streamTransform
@@ -145,6 +147,7 @@ import Streamly.Internal.Data.Unfold.Type (Unfold(..))
145147import Streamly.Internal.System.IO (unsafeInlineIO)
146148
147149import qualified Streamly.Internal.Data.Array.Foreign.Mut.Type as MA
150+import qualified Streamly.Internal.Data.Array.Foreign.Mut as MA
148151import qualified Streamly.Internal.Data.Array.Foreign.Type as A
149152import qualified Streamly.Internal.Data.Fold as FL
150153import qualified Streamly.Internal.Data.Stream.Prelude as P
@@ -377,6 +380,38 @@ getSliceUnsafe index len (Array fp e) =
377380 end = unsafeForeignPtrToPtr fp1 `plusPtr` (len * size)
378381 in assert (end <= e) (Array fp1 end)
379382
383+-- | Split the array into a stream of slices using a predicate. The element
384+-- matching the predicate is dropped.
385+--
386+-- /Pre-release/
387+{-# INLINE splitOn #-}
388+splitOn :: (Monad m, Storable a) =>
389+ (a -> Bool) -> Array a -> SerialT m (Array a)
390+splitOn predicate arr =
391+ fmap A.unsafeFreeze $ MA.splitOn predicate (A.unsafeThaw arr)
392+
393+{-# INLINE genSlicesFromLen #-}
394+genSlicesFromLen :: forall m a. (Monad m, Storable a)
395+ => Int -- ^ from index
396+ -> Int -- ^ length of the slice
397+ -> Unfold m (Array a) (Int, Int)
398+genSlicesFromLen from len =
399+ Unfold.lmap A.unsafeThaw (MA.genSlicesFromLen from len)
400+
401+-- | Generate a stream of slices of specified length from an array, starting
402+-- from the supplied array index. The last slice may be shorter than the
403+-- requested length.
404+--
405+-- /Pre-release//
406+{-# INLINE getSlicesFromLen #-}
407+getSlicesFromLen :: forall m a. (Monad m, Storable a)
408+ => Int -- ^ from index
409+ -> Int -- ^ length of the slice
410+ -> Unfold m (Array a) (Array a)
411+getSlicesFromLen from len =
412+ fmap A.unsafeFreeze
413+ $ Unfold.lmap A.unsafeThaw (MA.getSlicesFromLen from len)
414+
380415-------------------------------------------------------------------------------
381416-- Random reads and writes
382417-------------------------------------------------------------------------------
src/Streamly/Internal/Data/Array/Foreign/Mut.hs+52
@@ -25,9 +25,61 @@
2525module Streamly.Internal.Data.Array.Foreign.Mut
2626 (
2727 module Streamly.Internal.Data.Array.Foreign.Mut.Type
28+ , splitOn
29+ , genSlicesFromLen
30+ , getSlicesFromLen
2831 )
2932where
3033
3134import Prelude hiding (foldr, length, read, splitAt)
3235
3336import Streamly.Internal.Data.Array.Foreign.Mut.Type
37+import qualified Streamly.Internal.Data.Stream.StreamD as D
38+import Foreign.Storable (Storable)
39+import Streamly.Internal.Data.Stream.IsStream.Type (SerialT)
40+import qualified Streamly.Internal.Data.Stream.IsStream.Type as IsStream
41+import qualified Streamly.Internal.Data.Unfold as Unfold
42+import Streamly.Internal.Data.Unfold.Type (Unfold(..))
43+
44+-- | Split the array into a stream of slices using a predicate. The element
45+-- matching the predicate is dropped.
46+--
47+-- /Pre-release/
48+{-# INLINE splitOn #-}
49+splitOn :: (Monad m, Storable a) =>
50+ (a -> Bool) -> Array a -> SerialT m (Array a)
51+splitOn predicate arr =
52+ IsStream.fromStreamD
53+ $ fmap (\(i, len) -> getSliceUnsafe i len arr)
54+ $ D.sliceOnSuffix predicate (toStreamD arr)
55+
56+-- | Generate a stream of array slice descriptors ((index, len)) of specified
57+-- length from an array, starting from the supplied array index. The last slice
58+-- may be shorter than the requested length depending on the array length.
59+--
60+-- /Pre-release/
61+{-# INLINE genSlicesFromLen #-}
62+genSlicesFromLen :: forall m a. (Monad m, Storable a)
63+ => Int -- ^ from index
64+ -> Int -- ^ length of the slice
65+ -> Unfold m (Array a) (Int, Int)
66+genSlicesFromLen from len =
67+ let fromThenTo n = (from, from + len, n - 1)
68+ mkSlice n i = return (i, min len (n - i))
69+ in Unfold.lmap length
70+ $ Unfold.mapMWithInput mkSlice
71+ $ Unfold.lmap fromThenTo Unfold.enumerateFromThenTo
72+
73+-- | Generate a stream of slices of specified length from an array, starting
74+-- from the supplied array index. The last slice may be shorter than the
75+-- requested length depending on the array length.
76+--
77+-- /Pre-release/
78+{-# INLINE getSlicesFromLen #-}
79+getSlicesFromLen :: forall m a. (Monad m, Storable a)
80+ => Int -- ^ from index
81+ -> Int -- ^ length of the slice
82+ -> Unfold m (Array a) (Array a)
83+getSlicesFromLen from len =
84+ let mkSlice arr (i, n) = return $ getSliceUnsafe i n arr
85+ in Unfold.mapMWithInput mkSlice (genSlicesFromLen from len)
src/Streamly/Internal/Data/Stream/StreamD/Nesting.hs+34−2
@@ -136,6 +136,7 @@ module Streamly.Internal.Data.Stream.StreamD.Nesting
136136 , wordsBy
137137 , splitOnSeq
138138 , splitOnSuffixSeq
139+ , sliceOnSuffix
139140
140141 -- * Transform (Nested Containers)
141142 -- | Opposite to compact in ArrayStream
@@ -159,14 +160,16 @@ import Fusion.Plugin.Types (Fuse(..))
159160import GHC.Types (SPEC(..))
160161
161162import Streamly.Internal.Data.Array.Foreign.Type (Array(..))
162−import Streamly.Internal.Data.Refold.Type (Refold(..))
163+import Streamly.Internal.Data.Fold.Step (Step(..))
163164import Streamly.Internal.Data.Fold.Type (Fold(..))
164165import Streamly.Internal.Data.Parser (ParseError(..))
166+import Streamly.Internal.Data.Refold.Type (Refold(..))
165167import Streamly.Internal.Data.SVar.Type (adaptState)
168+import Streamly.Internal.Data.Tuple.Strict (Tuple'(..))
166169import Streamly.Internal.Data.Unfold.Type (Unfold(..))
167170
168171import qualified Streamly.Internal.Data.Array.Foreign.Type as A
169−import qualified Streamly.Internal.Data.Fold.Type as FL
172+import qualified Streamly.Internal.Data.Fold as FL
170173import qualified Streamly.Internal.Data.Parser as PR
171174import qualified Streamly.Internal.Data.Parser.ParserD as PRD
172175import qualified Streamly.Internal.Ring.Foreign as RB
@@ -1067,6 +1070,35 @@ refoldIterateM (Refold fstep finject fextract) initial (Stream step state) =
10671070 stepOuter _ (CIterYield a next) = return $ Yield a next
10681071 stepOuter _ CIterStop = return Stop
10691072
1073+-- "n" elements at the end are dropped by the fold.
1074+{-# INLINE sliceBy #-}
1075+sliceBy :: Monad m => Fold m a Int -> Int -> Refold m (Int, Int) a (Int, Int)
1076+sliceBy (Fold step1 initial1 extract1) n = Refold step inject extract
1077+
1078+ where
1079+
1080+ inject (i, len) = do
1081+ r <- initial1
1082+ return $ case r of
1083+ Partial s -> Partial $ Tuple' (i + len + n) s
1084+ Done l -> Done (i, l)
1085+
1086+ step (Tuple' i s) x = do
1087+ r <- step1 s x
1088+ return $ case r of
1089+ Partial s1 -> Partial $ Tuple' i s1
1090+ Done len -> Done (i, len)
1091+
1092+ extract (Tuple' i s) = (i,) <$> extract1 s
1093+
1094+{-# INLINE sliceOnSuffix #-}
1095+sliceOnSuffix :: Monad m => (a -> Bool) -> Stream m a -> Stream m (Int, Int)
1096+sliceOnSuffix predicate =
1097+ -- Scan the stream with the given refold
1098+ refoldIterateM
1099+ (sliceBy (FL.takeEndBy_ predicate FL.length) 1)
1100+ (return (-1, 0))
1101+
10701102------------------------------------------------------------------------------
10711103-- Parsing
10721104------------------------------------------------------------------------------
src/Streamly/Internal/Data/Stream/StreamD/Transform.hs+20
@@ -107,6 +107,7 @@ module Streamly.Internal.Data.Stream.StreamD.Transform
107107
108108 -- * Searching
109109 , findIndices
110+ , slicesBy
110111
111112 -- * Rolling map
112113 -- | Map using the previous element.
@@ -1174,6 +1175,25 @@ findIndices p (Stream step state) = Stream step' (state, 0)
11741175 Skip s -> Skip (s, i)
11751176 Stop -> Stop
11761177
1178+{-# INLINE_NORMAL slicesBy #-}
1179+slicesBy :: Monad m => (a -> Bool) -> Stream m a -> Stream m (Int, Int)
1180+slicesBy p (Stream step1 state1) = Stream step (Just (state1, 0, 0))
1181+
1182+ where
1183+
1184+ {-# INLINE_LATE step #-}
1185+ step gst (Just (st, i, len)) = i `seq` len `seq` do
1186+ r <- step1 (adaptState gst) st
1187+ return
1188+ $ case r of
1189+ Yield x s ->
1190+ if p x
1191+ then Yield (i, len + 1) (Just (s, i + len + 1, 0))
1192+ else Skip (Just (s, i, len + 1))
1193+ Skip s -> Skip (Just (s, i, len))
1194+ Stop -> if len == 0 then Stop else Yield (i, len) Nothing
1195+ step _ Nothing = return Stop
1196+
11771197------------------------------------------------------------------------------
11781198-- Rolling map
11791199------------------------------------------------------------------------------