8ae28e2 · streamly · git.iko.soy

git.iko.soy/streamly

on master help sign in

Add array splitOn, getSlicesFromLen

8ae28e2Harendra Kumar5 years ago

src/Streamly/Internal/Data/Array/Foreign.hs+36−1

compare with master

unifiedsplit

@@ -104,7 +104,9 @@ module Streamly.Internal.Data.Array.Foreign

104 -- * Subarrays

105 , getSliceUnsafe

106 -- , getSlice

107− -- , getSlicesFromLenN

107+ , genSlicesFromLen

108+ , getSlicesFromLen

109+ , splitOn

108

110

109 -- * Streaming Operations

111 -- * Streaming Operations

110 , streamTransform

112 , streamTransform

@@ -145,6 +147,7 @@ import Streamly.Internal.Data.Unfold.Type (Unfold(..))

145import Streamly.Internal.System.IO (unsafeInlineIO)

147import Streamly.Internal.System.IO (unsafeInlineIO)

146

148

147import qualified Streamly.Internal.Data.Array.Foreign.Mut.Type as MA

149import qualified Streamly.Internal.Data.Array.Foreign.Mut.Type as MA

150+import qualified Streamly.Internal.Data.Array.Foreign.Mut as MA

148import qualified Streamly.Internal.Data.Array.Foreign.Type as A

151import qualified Streamly.Internal.Data.Array.Foreign.Type as A

149import qualified Streamly.Internal.Data.Fold as FL

152import qualified Streamly.Internal.Data.Fold as FL

150import qualified Streamly.Internal.Data.Stream.Prelude as P

153import qualified Streamly.Internal.Data.Stream.Prelude as P

@@ -377,6 +380,38 @@ getSliceUnsafe index len (Array fp e) =

377 end = unsafeForeignPtrToPtr fp1 `plusPtr` (len * size)

380 end = unsafeForeignPtrToPtr fp1 `plusPtr` (len * size)

378 in assert (end <= e) (Array fp1 end)

381 in assert (end <= e) (Array fp1 end)

379

382

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+

380-------------------------------------------------------------------------------

415-------------------------------------------------------------------------------

381-- Random reads and writes

416-- Random reads and writes

382-------------------------------------------------------------------------------

417-------------------------------------------------------------------------------

src/Streamly/Internal/Data/Array/Foreign/Mut.hs+52

@@ -25,9 +25,61 @@

25module Streamly.Internal.Data.Array.Foreign.Mut

26 (

27 module Streamly.Internal.Data.Array.Foreign.Mut.Type

28+ , splitOn

29+ , genSlicesFromLen

30+ , getSlicesFromLen

28 )

31 )

29where

32where

31import Prelude hiding (foldr, length, read, splitAt)

34import Prelude hiding (foldr, length, read, splitAt)

33import Streamly.Internal.Data.Array.Foreign.Mut.Type

36import 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

136 , wordsBy

137 , splitOnSeq

138 , splitOnSuffixSeq

139+ , sliceOnSuffix

139

140

140 -- * Transform (Nested Containers)

141 -- * Transform (Nested Containers)

141 -- | Opposite to compact in ArrayStream

142 -- | Opposite to compact in ArrayStream

@@ -159,14 +160,16 @@ import Fusion.Plugin.Types (Fuse(..))

159import GHC.Types (SPEC(..))

160import GHC.Types (SPEC(..))

160

161

161import Streamly.Internal.Data.Array.Foreign.Type (Array(..))

162import Streamly.Internal.Data.Array.Foreign.Type (Array(..))

162−import Streamly.Internal.Data.Refold.Type (Refold(..))

163+import Streamly.Internal.Data.Fold.Step (Step(..))

163import Streamly.Internal.Data.Fold.Type (Fold(..))

164import Streamly.Internal.Data.Fold.Type (Fold(..))

164import Streamly.Internal.Data.Parser (ParseError(..))

165import Streamly.Internal.Data.Parser (ParseError(..))

166+import Streamly.Internal.Data.Refold.Type (Refold(..))

165import Streamly.Internal.Data.SVar.Type (adaptState)

167import Streamly.Internal.Data.SVar.Type (adaptState)

168+import Streamly.Internal.Data.Tuple.Strict (Tuple'(..))

166import Streamly.Internal.Data.Unfold.Type (Unfold(..))

169import Streamly.Internal.Data.Unfold.Type (Unfold(..))

167

170

168import qualified Streamly.Internal.Data.Array.Foreign.Type as A

171import 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

170import qualified Streamly.Internal.Data.Parser as PR

173import qualified Streamly.Internal.Data.Parser as PR

171import qualified Streamly.Internal.Data.Parser.ParserD as PRD

174import qualified Streamly.Internal.Data.Parser.ParserD as PRD

172import qualified Streamly.Internal.Ring.Foreign as RB

175import qualified Streamly.Internal.Ring.Foreign as RB

@@ -1067,6 +1070,35 @@ refoldIterateM (Refold fstep finject fextract) initial (Stream step state) =

1067 stepOuter _ (CIterYield a next) = return $ Yield a next

1070 stepOuter _ (CIterYield a next) = return $ Yield a next

1068 stepOuter _ CIterStop = return Stop

1071 stepOuter _ CIterStop = return Stop

1069

1072

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+

1070------------------------------------------------------------------------------

1102------------------------------------------------------------------------------

1071-- Parsing

1103-- Parsing

1072------------------------------------------------------------------------------

1104------------------------------------------------------------------------------

src/Streamly/Internal/Data/Stream/StreamD/Transform.hs+20

@@ -107,6 +107,7 @@ module Streamly.Internal.Data.Stream.StreamD.Transform

107

108 -- * Searching

109 , findIndices

110+ , slicesBy

110

111

111 -- * Rolling map

112 -- * Rolling map

112 -- | Map using the previous element.

113 -- | Map using the previous element.

@@ -1174,6 +1175,25 @@ findIndices p (Stream step state) = Stream step' (state, 0)

1174 Skip s -> Skip (s, i)

1175 Skip s -> Skip (s, i)

1175 Stop -> Stop

1176 Stop -> Stop

1176

1177

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+

1177------------------------------------------------------------------------------

1197------------------------------------------------------------------------------

1178-- Rolling map

1198-- Rolling map

1179------------------------------------------------------------------------------

1199------------------------------------------------------------------------------