a849812 · streamly · git.iko.soy

git.iko.soy/streamly

Maintain a list of top w elements in the stream

commita849812d9960d5b7868e770156ab774ee8f70157 authorRanjeet Kumar Ranjan <ranjeetr@composewell.com> dateMon 18 Apr 2022, 11:39 (5 years ago) parentc3cf9cc

5 files changed+112−11

compare with master

unifiedsplit

benchmark/Streamly/Benchmark/Data/Fold.hs+2 core/src/Streamly/Internal/Data/Array/Foreign/Mut/Type.hs+6−7 core/src/Streamly/Internal/Data/Fold.hs+59 test/Streamly/Test/Data/Array/Foreign.hs+15−4 test/Streamly/Test/Data/Fold.hs+30

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

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

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

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

289289 , 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)

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

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

292294 , 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

142142 , shuffleBy

143143 , divideBy

144144 , mergeBy

145− , bubbleAsc

145+ , bubble

146146

147147 -- * Casting

148148 , cast

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

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

23422342--

23432343-- /Pre-release/

2344−{-# INLINE bubbleAsc #-}

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

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

2347−bubbleAsc cmp0 arr =

2344+{-# INLINE bubble #-}

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

2346+bubble cmp0 arr =

23482347 when (l > 1) $ do

23492348 x <- getIndexUnsafe (l - 1) arr

23502349 go x (l - 2)

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

23572356 if i >= 0

23582357 then do

23592358 x1 <- getIndexUnsafe i arr

2360− case x1 `cmp0` x of

2361− GT -> do

2359+ case x `cmp0` x1 of

2360+ LT -> do

23622361 putIndexUnsafe (i + 1) x1 arr

23632362 go x (i - 1)

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

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

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

269269 , snoc

270270 , duplicate

271271 , finish

272+ , top

273+ , bottom

272274

273275 -- * Deprecated

274276 , sequence

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

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

287289import Data.Map.Strict (Map)

288290import Data.Maybe (isJust, fromJust)

291+import Foreign.Storable (Storable)

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

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

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

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

296299import qualified Data.Set as Set

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

298301import qualified Prelude

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

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

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

301305

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

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

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

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

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

323328

324329------------------------------------------------------------------------------

325330-- hoist

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

21212126

21222127 {-# INLINE_LATE consume #-}

21232128 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 =

150150 assert (ls == lst)

151151-}

152152

153−testBubbleAsc :: Property

154−testBubbleAsc =

153+testBubbleWith :: Bool -> Property

154+testBubbleWith asc =

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

156156 monadicIO $ action ls0

157157

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

160160 action ls = do

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

162162 lst <- MA.toList x

163− assert (sort ls == lst)

163+ if asc

164+ then assert (sort ls == lst)

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

164166

165167 fldm ls =

166168 Fold.foldlM'

167169 (\b a -> do

168170 arr <- MA.snoc b a

169− MA.bubbleAsc compare arr

171+ if asc

172+ then MA.bubble compare arr

173+ else MA.bubble (flip compare) arr

170174 return arr

171175 )

172176 (MA.newArray $ length ls)

173177

178+testBubbleAsc :: Property

179+testBubbleAsc = testBubbleWith True

180+

181+testBubbleDesc :: Property

182+testBubbleDesc = testBubbleWith False

183+

174184main :: IO ()

175185main =

176186 hspec $

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

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

183193 --prop "testAppend" testAppend

184194 prop "testBubbleAsc" testBubbleAsc

195+ prop "testBubbleDesc" testBubbleDesc

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

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

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

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

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

11module Main (main) where

3+import Data.List (sort)

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

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

56import Test.QuickCheck

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

1314 , property

1415 , vectorOf

1516 , withMaxSuccess

17+ , generate

1618 )

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

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

1821

1922import qualified Data.Map

2023import qualified Prelude

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

2225import qualified Streamly.Prelude as S

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

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

2428import qualified Streamly.Data.Fold as FL

2529

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

640644 let v2 = scanl (+) 0 lst

641645 assert (v1 == v2)

642646

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+

643671moduleName :: String

644672moduleName = "Data.Fold"

645673

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

688716 prop "any" $ Main.any predicate

689717 prop "and" Main.and

690718 prop "or" Main.or

719+ prop "top" Main.top

720+ prop "bottom" Main.bottom

691721

692722 -- Combinators

693723