Implement maxResults

Author: harendra-kumar

app/hfd.hs

@@ -1,5 +1,3 @@
-{-# LANGUAGE CPP #-}
-
 module Main (main) where
 
 import System.IO (stdout)
@@ -17,23 +15,20 @@ import Options.Applicative
     , long
     , metavar
     , optional
+    , option
     , progDesc
     , strArgument
     , (<**>)
     )
 import qualified Options.Applicative as OA
 import qualified Streamly.Data.Stream.Prelude as Stream
-#if defined(mingw32_HOST_OS) || defined(__MINGW32__)
-import qualified Streamly.Data.Array as Array
-import qualified Streamly.Data.Unfold as Unfold
-import qualified Streamly.Internal.Data.Stream as Stream (unfoldEachEndBy)
-import qualified Streamly.Unicode.Stream as Unicode
-#endif
 import qualified Streamly.FileSystem.Handle as Handle
 import qualified Streamly.FileSystem.Path as Path
 
 import Streamly.Coreutils.Find
     ( FindOptions
+    , findByteChunked
+    , maxResults
     , parallelInterleaved
     , parallelOrdered
     , parallelUnordered
@@ -43,15 +38,11 @@ import Streamly.Coreutils.Find
     , serialDfs
     , serialInterleaved
     )
-#if !defined(mingw32_HOST_OS) && !defined(__MINGW32__)
-import Streamly.Coreutils.Find (findByteChunked)
-#else
-import Streamly.Coreutils.Find (findChunked)
-#endif
 
 data Config = Config
     { cfgTraversal :: FindOptions -> FindOptions
     , cfgRoot :: FilePath
+    , cfgMaxResults :: Maybe Int
     }
 
 data Traversal
@@ -76,11 +67,12 @@ toTraversalConfig traversal =
         TraversalParallelInterleaved -> parallelInterleaved
         TraversalParallelOrdered -> parallelOrdered
 
-mkConfig :: Traversal -> Maybe FilePath -> Config
-mkConfig traversal mPath =
+mkConfig :: Traversal -> Maybe Int -> Maybe FilePath -> Config
+mkConfig traversal mMaxResults mPath =
     Config
         { cfgTraversal = toTraversalConfig traversal
         , cfgRoot = maybe "." id mPath
+        , cfgMaxResults = mMaxResults
         }
 
 traversalParser :: Parser Traversal
@@ -107,6 +99,11 @@ configParser :: Parser Config
 configParser =
     mkConfig
         <$> traversalParser
+        <*> optional
+                (option (OA.eitherReader parsePositiveInt)
+                    (long "max-results"
+                        <> metavar "N"
+                        <> help "Stop after emitting N results"))
         <*> optional
                 (strArgument
                     (metavar "PATH" <> help "Root path to search"))
@@ -120,22 +117,20 @@ parserInfo =
             <> progDesc "A basic fd-like driver for Streamly.Coreutils.Find."
             <> header "hfd")
 
-#if defined(mingw32_HOST_OS) || defined(__MINGW32__)
-#endif
-
 main :: IO ()
 main = do
     cfg <- execParser parserInfo
     path <- Path.fromString (cfgRoot cfg)
-#if !defined(mingw32_HOST_OS) && !defined(__MINGW32__)
+    let applyConfig opts =
+            maybe id maxResults (cfgMaxResults cfg) $
+                cfgTraversal cfg opts
     Stream.fold (Handle.writeChunks stdout)
-        $ findByteChunked (cfgTraversal cfg) path
-#else
-    Stream.fold (Handle.writeWith 32000 stdout)
-        $ Unicode.encodeUtf8
-        $ Unicode.decodeUtf16le
-        $ Stream.unfoldEachEndBy 10 Array.reader
-        $ fmap Path.toArray
-        $ Stream.unfoldEach Unfold.fromList
-        $ findChunked (cfgTraversal cfg) path
-#endif
+        $ findByteChunked applyConfig path
+
+parsePositiveInt :: String -> Either String Int
+parsePositiveInt str =
+    case reads str of
+        [(n, "")]
+            | n > 0 -> Right n
+            | otherwise -> Left "N must be positive"
+        _ -> Left "N must be an integer"

src/Streamly/Coreutils/Find.hs

@@ -85,9 +85,12 @@ module Streamly.Coreutils.Find
 
     -- * Options
     , FindOptions
+    , maxResults
     )
 where
 
+import Data.Function ((&))
+import Data.Functor.Identity (runIdentity)
 import Data.Maybe (fromJust)
 import Data.Word (Word8)
 import Streamly.Data.Array (Array)
@@ -100,12 +103,17 @@ import System.IO (stdout, hSetBuffering, BufferMode(LineBuffering))
 import qualified Streamly.Data.Stream.Prelude as Stream
 import qualified Streamly.Data.Array as Array
 import qualified Streamly.FileSystem.DirIO as DirIO
-import qualified Streamly.Internal.Data.Array as GArray (compactMax')
+import qualified Streamly.Internal.Data.Array as GArray
+    ( compactMax'
+    , read
+    , unsafeSliceOffLen
+    )
 import qualified Streamly.Internal.Data.Stream as Stream
     ( unfoldEachEndBy
     , concatIterate
     , bfsConcatIterate
     , altBfsConcatIterate
+    , postscanlMaybe
     )
 import qualified Streamly.Data.StreamK as StreamK
 import qualified Streamly.Internal.Data.StreamK as StreamK
@@ -125,6 +133,11 @@ import qualified Streamly.Internal.FileSystem.Posix.ReadDir as Dir
 import qualified Streamly.Unicode.Stream as Stream
 #endif
 
+import Streamly.Internal.Data.Scanl (Step(..), Scanl(..))
+import qualified Streamly.Internal.Data.Scanl as Scanl
+import qualified Streamly.Internal.Data.Fold as Fold
+import qualified Streamly.Internal.Data.Array as Array
+
 --
 -- Running on a sample directory tree the concurrent rust "fd" tool took 150 ms
 -- (real time). On the same tree the fastest variant using Haskell streamly
@@ -147,10 +160,17 @@ data FindTraversal
     | FindParallelInterleaved
     | FindParallelOrdered
 
-newtype FindOptions = FindOptions {findTraversal :: FindTraversal}
+data FindOptions = FindOptions
+    { findTraversal :: FindTraversal
+    , findMaxResults :: Maybe Int
+    }
 
 defaultConfig :: FindOptions
-defaultConfig = FindOptions FindSerialDfs
+defaultConfig =
+    FindOptions
+        { findTraversal = FindSerialDfs
+        , findMaxResults = Nothing
+        }
 
 serialDfs :: FindOptions -> FindOptions
 serialDfs cfg = cfg {findTraversal = FindSerialDfs}
@@ -176,6 +196,9 @@ parallelInterleaved cfg = cfg {findTraversal = FindParallelInterleaved}
 parallelOrdered :: FindOptions -> FindOptions
 parallelOrdered cfg = cfg {findTraversal = FindParallelOrdered}
 
+maxResults :: Int -> FindOptions -> FindOptions
+maxResults n cfg = cfg {findMaxResults = Just (max 0 n)}
+
 {-# INLINE recReadOpts #-}
 recReadOpts :: ReadOptions -> ReadOptions
 {-# INLINE reEncode #-}
@@ -196,12 +219,70 @@ recReadOpts =
 reEncode = id
 #endif
 
+data Counts = Counts !Int !Int deriving Show
+
+{-# INLINE countStep #-}
+countStep :: Monad m => Counts -> Word8 -> m (Step Counts (Either Int Int))
+countStep (Counts l c) ch =
+    let l1 = if ch == 10 then l - 1 else l
+     in if l1 == 0
+        then return $ Done $ Left (c + 1)
+        else return $ Partial $ Counts l1 (c + 1)
+
+{-# INLINE countExtract #-}
+countExtract :: Monad m => Counts -> m (Either a Int)
+countExtract (Counts l _) = return $ Right l
+
+{-# INLINE count #-}
+count :: Monad m => Int -> Fold.Fold m Word8 (Either Int Int)
+count l = Fold.foldtM' countStep (return $ Partial (Counts l 0 )) countExtract
+
+-- XXX Scanl is an awkward abstraction for the case when we are emitting every
+-- element and just need to transform the elements using the state. We need a
+-- smapM instead for this case. In the scan we are forced to use a Maybe and
+-- then catMaybe unnecessarily to store the elements. Because only in the
+-- initial state we do not have an element.
+--
+{-# INLINE scanStep #-}
+scanStep :: Monad m =>
+       (Int, Maybe (Array Word8))
+    -> Array Word8
+    -> m (Step (Int, Maybe (Array Word8)) (Maybe (Array Word8)))
+scanStep (n, _) arr = do
+    r <- Array.read arr & Stream.fold (count n)
+    case r of
+        Left len -> return $ Done $ Just (Array.unsafeSliceOffLen 0 len arr)
+        Right cnt ->
+            if cnt /= 0
+            then return $ Partial (cnt, Just arr)
+            else return $ Done (Just arr)
+
+{-# INLINE scanExtract #-}
+scanExtract :: Monad m => (Int, Maybe (Array Word8)) -> m (Maybe (Array Word8))
+scanExtract (_, arr) = return arr
+
+{-# INLINE scanFinal #-}
+scanFinal :: Monad m => (Int, Maybe (Array Word8)) -> m (Maybe (Array Word8))
+scanFinal (_, arr) = return arr
+
+{-# INLINE takeN #-}
+takeN :: Int -> Stream IO (Array Word8) -> Stream IO (Array Word8)
+takeN n
+    | n <= 0 = const Stream.nil
+    | otherwise =
+        Stream.postscanlMaybe
+            (Scanl
+                scanStep
+                (return (Partial (n, Nothing)))
+                scanExtract
+                scanFinal)
+
 #if !defined(mingw32_HOST_OS) && !defined(__MINGW32__)
 -- Fastest implementation, only works for posix as of now.
 findByteChunked :: (FindOptions -> FindOptions) -> Path -> Stream IO (Array Word8)
 findByteChunked f path =
-        Stream.catRights $
-            case findTraversal (f defaultConfig) of
+        transform $ Stream.catRights $
+            case findTraversal opts of
                 FindSerialDfs ->
                     Stream.concatIterate streamDirMaybe -- 154 ms
                     $ Stream.fromPure (Left [path])
@@ -232,6 +313,11 @@ findByteChunked f path =
 
     where
 
+    {-# INLINE transform #-}
+    transform s = maybe s (\n -> takeN n s) (findMaxResults opts)
+
+    opts = f defaultConfig
+
     concatIterateWith combine =
           StreamK.toStream
         . StreamK.concatIterateWith combine (StreamK.fromStream . streamDir)
@@ -253,6 +339,7 @@ findByteChunked f path =
 -- Faster than the find implementation below
 findChunked :: (FindOptions -> FindOptions) -> Path -> Stream IO [Path]
 findChunked f path =
+        -- XXX implement maxResults
         Stream.catRights $
             case findTraversal (f defaultConfig) of
                 FindSerialDfs ->
@@ -300,6 +387,7 @@ findChunked f path =
 
 find :: (FindOptions -> FindOptions) -> Path -> Stream IO Path
 find f path =
+        -- XXX implement maxResults
         Stream.catRights $
             case findTraversal (f defaultConfig) of
                 FindSerialDfs ->