a849812 · streamly · git.iko.soy

git.iko.soy/streamly

on master help sign in

Maintain a list of top w elements in the stream

a849812Ranjeet Kumar Ranjan5 years ago

benchmark/Streamly/Benchmark/Data/Fold.hs+2

compare with master

unifiedsplit

@@ -287,6 +287,8 @@ o_1_space_serial_elimination value =

287 , benchIOSink value "drainN" (S.fold (FL.drainN value))

288 , benchIOSink value "last" (S.fold FL.last)

289 , benchIOSink value "length" (S.fold FL.length)

290+ , benchIOSink value "top" (S.fold $ FL.top 10)

291+ , benchIOSink value "bottom" (S.fold $ FL.bottom 10)

290 , benchIOSink value "sum" (S.fold FL.sum)

292 , benchIOSink value "sum" (S.fold FL.sum)

291 , benchIOSink value "sum (foldMap)" (S.fold (FL.foldMap Sum))

293 , benchIOSink value "sum (foldMap)" (S.fold (FL.foldMap Sum))

292 , benchIOSink value "product" (S.fold FL.product)

294 , benchIOSink value "product" (S.fold FL.product)

core/src/Streamly/Internal/Data/Array/Foreign/Mut/Type.hs+6−7

@@ -142,7 +142,7 @@ module Streamly.Internal.Data.Array.Foreign.Mut.Type

142 , shuffleBy

143 , divideBy

144 , mergeBy

145− , bubbleAsc

145+ , bubble

146

147 -- * Casting

148 , cast

@@ -2341,10 +2341,9 @@ strip eq arr@Array{..} = do

2341-- that the array remains sorted in ascending order.

2342--

2343-- /Pre-release/

2344−{-# INLINE bubbleAsc #-}

2344+{-# INLINE bubble #-}

2345−bubbleAsc :: (MonadIO m, Storable a) =>

2346− (a -> a -> Ordering) -> Array a -> m ()

2345+bubble :: (MonadIO m, Storable a) => (a -> a -> Ordering) -> Array a -> m ()

2347−bubbleAsc cmp0 arr =

2346+bubble cmp0 arr =

2348 when (l > 1) $ do

2347 when (l > 1) $ do

2349 x <- getIndexUnsafe (l - 1) arr

2348 x <- getIndexUnsafe (l - 1) arr

2350 go x (l - 2)

2349 go x (l - 2)

@@ -2357,8 +2356,8 @@ bubbleAsc cmp0 arr =

2357 if i >= 0

2356 if i >= 0

2358 then do

2357 then do

2359 x1 <- getIndexUnsafe i arr

2358 x1 <- getIndexUnsafe i arr

2360− case x1 `cmp0` x of

2359+ case x `cmp0` x1 of

2361− GT -> do

2360+ LT -> do

2362 putIndexUnsafe (i + 1) x1 arr

2361 putIndexUnsafe (i + 1) x1 arr

2363 go x (i - 1)

2362 go x (i - 1)

2364 _ -> putIndexUnsafe (i + 1) x arr

2363 _ -> putIndexUnsafe (i + 1) x arr

core/src/Streamly/Internal/Data/Fold.hs+59

@@ -269,6 +269,8 @@ module Streamly.Internal.Data.Fold

269 , snoc

270 , duplicate

271 , finish

272+ , top

273+ , bottom

272

274

273 -- * Deprecated

275 -- * Deprecated

274 , sequence

276 , sequence

@@ -286,6 +288,7 @@ import Data.Int (Int64)

286import Data.IORef (newIORef, readIORef, writeIORef)

288import Data.IORef (newIORef, readIORef, writeIORef)

287import Data.Map.Strict (Map)

289import Data.Map.Strict (Map)

288import Data.Maybe (isJust, fromJust)

290import Data.Maybe (isJust, fromJust)

291+import Foreign.Storable (Storable)

289import Streamly.Internal.Data.IsMap (IsMap(..))

292import Streamly.Internal.Data.IsMap (IsMap(..))

290import Streamly.Internal.Data.Pipe.Type (Pipe (..), PipeState(..))

293import Streamly.Internal.Data.Pipe.Type (Pipe (..), PipeState(..))

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

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

@@ -296,6 +299,7 @@ import qualified Data.IntSet as IntSet

296import qualified Data.Set as Set

299import qualified Data.Set as Set

297import qualified Streamly.Internal.Data.IsMap as IsMap

300import qualified Streamly.Internal.Data.IsMap as IsMap

298import qualified Prelude

301import qualified Prelude

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

299import qualified Streamly.Internal.Data.Pipe.Type as Pipe

303import qualified Streamly.Internal.Data.Pipe.Type as Pipe

300import qualified Streamly.Internal.Data.Stream.StreamD.Type as StreamD

304import qualified Streamly.Internal.Data.Stream.StreamD.Type as StreamD

301

305

@@ -320,6 +324,7 @@ import Streamly.Internal.Data.Fold.Type

320-- >>> import qualified Streamly.Internal.Data.Parser as Parser

324-- >>> import qualified Streamly.Internal.Data.Parser as Parser

321-- >>> import Streamly.Internal.Data.Stream.Serial (SerialT(..))

325-- >>> import Streamly.Internal.Data.Stream.Serial (SerialT(..))

322-- >>> import Data.IORef (newIORef, readIORef, writeIORef)

326-- >>> import Data.IORef (newIORef, readIORef, writeIORef)

327+-- >>> import qualified Streamly.Internal.Data.Array.Foreign.Mut.Type as MA

323

328

324------------------------------------------------------------------------------

329------------------------------------------------------------------------------

325-- hoist

330-- hoist

@@ -2121,3 +2126,57 @@ unfoldMany (Unfold ustep inject) (Fold fstep initial extract) =

2121

2126

2122 {-# INLINE_LATE consume #-}

2127 {-# INLINE_LATE consume #-}

2123 consume s a = inject a >>= produce s

2128 consume s a = inject a >>= produce s

2129+

2130+{-# INLINE topBy #-}

2131+topBy :: (MonadIO m, Storable a) =>

2132+ (a -> a -> Ordering)

2133+ -> Int

2134+ -> Fold m a (MA.Array a)

2135+topBy cmp n = Fold step initial extract

2136+

2137+ where

2138+

2139+ initial = do

2140+ arr <- MA.newArray n

2141+ if n <= 0

2142+ then return $ Done arr

2143+ else return $ Partial (arr, 0)

2144+

2145+ step (arr, i) x =

2146+ if i < n

2147+ then do

2148+ arr' <- MA.snoc arr x

2149+ MA.bubble cmp arr'

2150+ return $ Partial (arr', i + 1)

2151+ else do

2152+ x1 <- MA.getIndexUnsafe (i - 1) arr

2153+ case x `cmp` x1 of

2154+ LT -> do

2155+ MA.putIndexUnsafe (i - 1) x arr

2156+ MA.bubble cmp arr

2157+ return $ Partial (arr, i)

2158+ _ -> return $ Partial (arr, i)

2159+

2160+ extract = return . fst

2161+

2162+-- | Fold the input stream to top n elements.

2163+--

2164+-- >>> stream = Stream.fromList [2::Int,7,9,3,1,5,6,11,17]

2165+-- >>> Stream.fold (Fold.top 3) stream >>= MA.toList

2166+-- [17,11,9]

2167+--

2168+-- /Pre-release/

2169+{-# INLINE top #-}

2170+top :: (MonadIO m, Storable a, Ord a) => Int -> Fold m a (MA.Array a)

2171+top = topBy $ flip compare

2172+

2173+-- | Fold the input stream to bottom n elements.

2174+--

2175+-- >>> stream = Stream.fromList [2::Int,7,9,3,1,5,6,11,17]

2176+-- >>> Stream.fold (Fold.bottom 3) stream >>= MA.toList

2177+-- [1,2,3]

2178+--

2179+-- /Pre-release/

2180+{-# INLINE bottom #-}

2181+bottom :: (MonadIO m, Storable a, Ord a) => Int -> Fold m a (MA.Array a)

2182+bottom = topBy compare

test/Streamly/Test/Data/Array/Foreign.hs+15−4

@@ -150,8 +150,8 @@ testAppend =

150 assert (ls == lst)

151-}

152

153−testBubbleAsc :: Property

153+testBubbleWith :: Bool -> Property

154−testBubbleAsc =

154+testBubbleWith asc =

155 forAll (listOf (chooseInt (-50, 100))) $ \ls0 ->

156 monadicIO $ action ls0

157

@@ -160,17 +160,27 @@ testBubbleAsc =

160 action ls = do

161 x <- S.fold (fldm ls) $ S.fromList ls

162 lst <- MA.toList x

163+ if asc

164+ then assert (sort ls == lst)

163− assert (sort ls == lst)

165+ else assert (sort ls == reverse lst)

164

166

165 fldm ls =

167 fldm ls =

166 Fold.foldlM'

168 Fold.foldlM'

167 (\b a -> do

169 (\b a -> do

168 arr <- MA.snoc b a

170 arr <- MA.snoc b a

171+ if asc

172+ then MA.bubble compare arr

169− MA.bubbleAsc compare arr

173+ else MA.bubble (flip compare) arr

170 return arr

174 return arr

171 )

175 )

172 (MA.newArray $ length ls)

176 (MA.newArray $ length ls)

173

177

178+testBubbleAsc :: Property

179+testBubbleAsc = testBubbleWith True

180+

181+testBubbleDesc :: Property

182+testBubbleDesc = testBubbleWith False

183+

174main :: IO ()

184main :: IO ()

175main =

185main =

176 hspec $

186 hspec $

@@ -182,6 +192,7 @@ main =

182 -- XXX There is an issue https://github.com/composewell/streamly/issues/1577

192 -- XXX There is an issue https://github.com/composewell/streamly/issues/1577

183 --prop "testAppend" testAppend

193 --prop "testAppend" testAppend

184 prop "testBubbleAsc" testBubbleAsc

194 prop "testBubbleAsc" testBubbleAsc

195+ prop "testBubbleDesc" testBubbleDesc

185 prop "length . fromStream === n" testLengthFromStream

196 prop "length . fromStream === n" testLengthFromStream

186 prop "toStream . fromStream === id" testFromStreamToStream

197 prop "toStream . fromStream === id" testFromStreamToStream

187 prop "read . write === id" testFoldUnfold

198 prop "read . write === id" testFoldUnfold

test/Streamly/Test/Data/Fold.hs+30

@@ -1,5 +1,6 @@

1module Main (main) where

3+import Data.List (sort)

3import Data.Semigroup (Sum(..), getSum)

4import Data.Semigroup (Sum(..), getSum)

4import Streamly.Test.Common (checkListEqual, listEquals)

5import Streamly.Test.Common (checkListEqual, listEquals)

5import Test.QuickCheck

6import Test.QuickCheck

@@ -13,13 +14,16 @@ import Test.QuickCheck

13 , property

14 , property

14 , vectorOf

15 , vectorOf

15 , withMaxSuccess

16 , withMaxSuccess

17+ , generate

16 )

18 )

19+import Control.Monad.IO.Class (liftIO)

17import Test.QuickCheck.Monadic (monadicIO, assert, run)

20import Test.QuickCheck.Monadic (monadicIO, assert, run)

19import qualified Data.Map

22import qualified Data.Map

20import qualified Prelude

23import qualified Prelude

21import qualified Streamly.Internal.Data.Fold as F

24import qualified Streamly.Internal.Data.Fold as F

22import qualified Streamly.Prelude as S

25import qualified Streamly.Prelude as S

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

23import qualified Streamly.Internal.Data.Stream.IsStream as Stream

27import qualified Streamly.Internal.Data.Stream.IsStream as Stream

24import qualified Streamly.Data.Fold as FL

28import qualified Streamly.Data.Fold as FL

@@ -640,6 +644,30 @@ scan = forAll (listOf (chooseInt (0, 100))) $ \lst ->

640 let v2 = scanl (+) 0 lst

644 let v2 = scanl (+) 0 lst

641 assert (v1 == v2)

645 assert (v1 == v2)

642

646

647+topBy :: Bool -> Property

648+topBy isTop = forAll (listOf (chooseInt (-50, 100))) $ \ls0 ->

649+ monadicIO $ action ls0

650+

651+ where

652+

653+ action ls = do

654+

655+ let n0 = Prelude.length ls

656+ n <- liftIO $ generate $ chooseInt (-2, n0 + 2)

657+ if isTop

658+ then do

659+ lst <- S.fold (F.top n) (S.fromList ls) >>= MA.toList

660+ assert ((Prelude.take n . Prelude.reverse . sort) ls == lst)

661+ else do

662+ lst <- S.fold (F.bottom n) (S.fromList ls) >>= MA.toList

663+ assert ((Prelude.take n . sort) ls == lst)

664+

665+top :: Property

666+top = topBy True

667+

668+bottom :: Property

669+bottom = topBy False

670+

643moduleName :: String

671moduleName :: String

644moduleName = "Data.Fold"

672moduleName = "Data.Fold"

645

673

@@ -688,6 +716,8 @@ main = hspec $ do

688 prop "any" $ Main.any predicate

716 prop "any" $ Main.any predicate

689 prop "and" Main.and

717 prop "and" Main.and

690 prop "or" Main.or

718 prop "or" Main.or

719+ prop "top" Main.top

720+ prop "bottom" Main.bottom

691

721

692 -- Combinators

722 -- Combinators

693

723