diff --git a/.gitignore b/.gitignore --- a/.gitignore +++ b/.gitignore @@ -1,27 +1,27 @@ -dist -dist-* -cabal-dev -*.o -*.hi -*.hie -*.chi -*.chs.h -*.dyn_o -*.dyn_hi -.hpc -.hsenv -.cabal-sandbox/ -cabal.sandbox.config -*.prof -*.aux -*.hp -*.eventlog -.stack-work/ -cabal.project.local -cabal.project.local~ -.HTF/ -.ghc.environment.* -*.exe -*.txt -*.vim -proust/ +dist +dist-* +cabal-dev +*.o +*.hi +*.hie +*.chi +*.chs.h +*.dyn_o +*.dyn_hi +.hpc +.hsenv +.cabal-sandbox/ +cabal.sandbox.config +*.prof +*.aux +*.hp +*.eventlog +.stack-work/ +cabal.project.local +cabal.project.local~ +.HTF/ +.ghc.environment.* +*.exe +*.txt +*.vim +proust/ diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,11 @@ -# Revision history for jet-stream - -## 0.1.0.0 -- YYYY-mm-dd - -* First version. Released on an unsuspecting world. +# Revision history for jet-stream + +## 1.1.0.0 + +* Switch to GHC2021 +* Remove dependency on 'conceit' package. +* Remove the use of LinearTypes. + +## 1.0.0.0 -- YYYY-mm-dd + +* First version. Released on an unsuspecting world. diff --git a/LICENSE b/LICENSE --- a/LICENSE +++ b/LICENSE @@ -1,30 +1,30 @@ -Copyright (c) 2021, Daniel Diaz - -All rights reserved. - -Redistribution and use in source and binary forms, with or without -modification, are permitted provided that the following conditions are met: - - * Redistributions of source code must retain the above copyright - notice, this list of conditions and the following disclaimer. - - * Redistributions in binary form must reproduce the above - copyright notice, this list of conditions and the following - disclaimer in the documentation and/or other materials provided - with the distribution. - - * Neither the name of Daniel Diaz nor the names of other - contributors may be used to endorse or promote products derived - from this software without specific prior written permission. - -THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS -"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT -LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR -A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT -OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, -SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT -LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, -DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY -THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT -(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE -OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +Copyright (c) 2021, Daniel Diaz + +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + + * Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + + * Redistributions in binary form must reproduce the above + copyright notice, this list of conditions and the following + disclaimer in the documentation and/or other materials provided + with the distribution. + + * Neither the name of Daniel Diaz nor the names of other + contributors may be used to endorse or promote products derived + from this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT +LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY +THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/README.md b/README.md --- a/README.md +++ b/README.md @@ -1,71 +1,71 @@ -# jet-stream - -This is yet another streaming library for Haskell, created to scratch the -following itches: - -- The main type is as simple as possible: the only type parameter is the type - of the yielded elements. - -- The `Monoid` / `Alternative` / `MonadPlus` methods perform concatenation, - just like with regular lists. The `Functor` `Applicative` and `Monad` - instances also resemble those of lists. - -- There are direct analogues of functions like `withFile`, `bracket`, `finally` - and `onError` that easy to integrate in a streaming pipeline, and behave - smartly when combined with functions like `take`. - -- Compatible with the [foldl](https://hackage.haskell.org/package/foldl) - library for collector-like terminal operations. (All self-respecting - streaming libraries must have this.) - -In order to achieve those objectives, the following sacrifices have been made: - -- No flexibility in the underlying monad for the stream effects: it's always - `IO`. - -- No separate "channels" that return extra information at the end of the - stream. This means exceptions are the only way of signalling errors or - unexpected conditions. - -- Elements in a stream can't be "extracted" one by one in a pull-based way, - like you can do for example in - [streaming](https://hackage.haskell.org/package/streaming-0.2.3.0/docs/Streaming-Prelude.html#v:next). - -- There's `take` and `drop`, but not at proper `splitAt`. Also, grouping - operations are cumbersome and underpowered, especially compared to libraries - like - [streaming]((https://hackage.haskell.org/package/streaming-0.2.3.0/docs/Streaming-Prelude.html#v:next)) - or - [streaming-bytestring](https://hackage.haskell.org/package/streaming-bytestring). - -## What about performance? - -I haven't run any benchmarks, but you can safely assume that this library will -move like a snail compared to -[streamly](https://hackage.haskell.org/package/streamly)'s Ferrari. - -## Some close cousins - -- [turtle](https://hackage.haskell.org/package/turtle). The `Shell` type - resembles `Jet`. One possible difference is that `Shell` doesn't seem to - provide a way for the `Shell` consumer to signal that no further values are - needed, at least judging from the docs for - [limit](https://hackage.haskell.org/package/turtle-1.5.22/docs/Turtle-Prelude.html#v:limit). - - \"turtle\" also inspired the idea of having a separate type for lines. - -- [streamly](https://hackage.haskell.org/package/streamly). I might have - reinvented a subset of streamly ([but - worse](https://www.mcmillen.dev/language_checklist.html)). - -- [Z.IO.BIO](https://hackage.haskell.org/package/Z-IO-1.0.0.0/docs/Z-IO-BIO.html) - from [Z-IO](https://hackage.haskell.org/package/Z-IO). Like `Jet`, uses a - callback-transformation approach. - -- The - [Stream](https://docs.oracle.com/en/java/javase/16/docs/api/java.base/java/util/stream/Stream.html) - type from Java is somewhat similar to this library's `Jet`. (And the - [foldl](https://hackage.haskell.org/package/foldl) library would be - the analogue of - [Collectors](https://docs.oracle.com/en/java/javase/16/docs/api/java.base/java/util/stream/Collectors.html).) - +# jet-stream + +This is yet another streaming library for Haskell, created to scratch the +following itches: + +- The main type is as simple as possible: the only type parameter is the type + of the yielded elements. + +- The `Monoid` / `Alternative` / `MonadPlus` methods perform concatenation, + just like with regular lists. The `Functor` `Applicative` and `Monad` + instances also resemble those of lists. + +- There are direct analogues of functions like `withFile`, `bracket`, `finally` + and `onError` that easy to integrate in a streaming pipeline, and behave + smartly when combined with functions like `take`. + +- Compatible with the [foldl](https://hackage.haskell.org/package/foldl) + library for collector-like terminal operations. (All self-respecting + streaming libraries must have this.) + +In order to achieve those objectives, the following sacrifices have been made: + +- No flexibility in the underlying monad for the stream effects: it's always + `IO`. + +- No separate "channels" that return extra information at the end of the + stream. This means exceptions are the only way of signalling errors or + unexpected conditions. + +- Elements in a stream can't be "extracted" one by one in a pull-based way, + like you can do for example in + [streaming](https://hackage.haskell.org/package/streaming-0.2.3.0/docs/Streaming-Prelude.html#v:next). + +- There's `take` and `drop`, but not at proper `splitAt`. Also, grouping + operations are cumbersome and underpowered, especially compared to libraries + like + [streaming]((https://hackage.haskell.org/package/streaming-0.2.3.0/docs/Streaming-Prelude.html#v:next)) + or + [streaming-bytestring](https://hackage.haskell.org/package/streaming-bytestring). + +## What about performance? + +I haven't run any benchmarks, but you can safely assume that this library will +move like a snail compared to +[streamly](https://hackage.haskell.org/package/streamly)'s Ferrari. + +## Some close cousins + +- [turtle](https://hackage.haskell.org/package/turtle). The `Shell` type + resembles `Jet`. One possible difference is that `Shell` doesn't seem to + provide a way for the `Shell` consumer to signal that no further values are + needed, at least judging from the docs for + [limit](https://hackage.haskell.org/package/turtle-1.5.22/docs/Turtle-Prelude.html#v:limit). + + \"turtle\" also inspired the idea of having a separate type for lines. + +- [streamly](https://hackage.haskell.org/package/streamly). I might have + reinvented a subset of streamly ([but + worse](https://www.mcmillen.dev/language_checklist.html)). + +- [Z.IO.BIO](https://hackage.haskell.org/package/Z-IO-1.0.0.0/docs/Z-IO-BIO.html) + from [Z-IO](https://hackage.haskell.org/package/Z-IO). Like `Jet`, uses a + callback-transformation approach. + +- The + [Stream](https://docs.oracle.com/en/java/javase/16/docs/api/java.base/java/util/stream/Stream.html) + type from Java is somewhat similar to this library's `Jet`. (And the + [foldl](https://hackage.haskell.org/package/foldl) library would be + the analogue of + [Collectors](https://docs.oracle.com/en/java/javase/16/docs/api/java.base/java/util/stream/Collectors.html).) + diff --git a/jet-stream.cabal b/jet-stream.cabal --- a/jet-stream.cabal +++ b/jet-stream.cabal @@ -1,77 +1,65 @@ -cabal-version: 3.0 -name: jet-stream -version: 1.0.0.0 - -synopsis: Yet another streaming library. -description: - This is a streaming library focused on simplicity at the cost of some - expressivity. - - Basic operations like `drop` and `take` are supported. - - The `Functor`, `Applicative` and `Monad` instances of the stream type - resemble those of pure lists. There are also `Monoid`, `Alternative` and - `MonadPlus` instances for stream concatenation. - - Provides resource-managing operations like `withFile` that are easily - integrated into streams. - - For expressive and composable terminal operations, streams can be consumed - with folds from the "foldl" library. - -license: BSD-3-Clause -license-file: LICENSE -author: Daniel Diaz Carrete -maintainer: diaz_carrete@yahoo.com -category: Streaming -extra-source-files: - CHANGELOG.md, - README.md, - .gitignore -source-repository head - type: git - location: https://github.com/danidiaz/jet-stream.git - -common common - build-depends: - base >= 4.11.0.0 && < 5, - bytestring >= 0.10, - text >= 1.2, - stm >= 2.4, - process >= 1.4.3.0, - stm-chans ^>= 3.0.0.0, - async ^>= 2.2.3, - conceit ^>= 0.5.0.0, - default-language: Haskell2010 - -library - import: common - exposed-modules: Jet - Jet.Internal - build-depends: - hs-source-dirs: lib - -test-suite tests - import: common - ghc-options: -threaded - type: exitcode-stdio-1.0 - hs-source-dirs: test - main-is: tests.hs - build-depends: - jet-stream, - time >= 1.9, - tasty ^>= 1.3.1, - tasty-hunit ^>= 0.10.0.2, - --- VERY IMPORTANT for doctests to work: https://stackoverflow.com/a/58027909/1364288 --- http://hackage.haskell.org/package/cabal-doctest -test-suite doctests - import: common - ghc-options: -threaded - type: exitcode-stdio-1.0 - hs-source-dirs: test - main-is: doctests.hs - build-depends: - jet-stream, - foldl ^>= 1.4.12, - doctest ^>= 0.18.1, +cabal-version: 3.0 +name: jet-stream +version: 1.1.0.0 + +synopsis: Yet another streaming library. +description: + This is a streaming library focused on simplicity at the cost of some + expressivity. + + Basic operations like `drop` and `take` are supported. + + The `Functor`, `Applicative` and `Monad` instances of the stream type + resemble those of pure lists. There are also `Monoid`, `Alternative` and + `MonadPlus` instances for stream concatenation. + + Provides resource-managing operations like `withFile` that are easily + integrated into streams. + + For expressive and composable terminal operations, streams can be consumed + with folds from the "foldl" library. + +license: BSD-3-Clause +license-file: LICENSE +author: Daniel Diaz Carrete +maintainer: diaz_carrete@yahoo.com +category: Streaming +tested-with: GHC ==9.6.4 +extra-source-files: + .gitignore +extra-doc-files: + CHANGELOG.md, + README.md, +source-repository head + type: git + location: https://github.com/danidiaz/jet-stream.git + +common common + build-depends: + base >= 4.11.0.0 && < 5, + bytestring >= 0.10.0 && < 0.13, + stm >= 2.5.0 && < 2.6, + text >= 2.0 && < 2.2, + process >= 1.6.0 && < 1.7, + stm-chans ^>= 3.0.0.0, + async ^>= 2.2.5 + default-language: GHC2021 + +library + import: common + exposed-modules: Jet + Jet.Internal + build-depends: + hs-source-dirs: lib + +test-suite tests + import: common + ghc-options: -threaded + type: exitcode-stdio-1.0 + hs-source-dirs: test + main-is: tests.hs + build-depends: + jet-stream, + time >= 1.9, + tasty >= 1.4.0, + tasty-hunit >= 0.10.0.2, diff --git a/lib/Jet.hs b/lib/Jet.hs --- a/lib/Jet.hs +++ b/lib/Jet.hs @@ -1,301 +1,288 @@ -{-# LANGUAGE ImportQualifiedPost #-} -{-# LANGUAGE PatternSynonyms #-} --- | A streaming library build around the 'Jet' type, which behaves as a kind of \"effectful list\". --- --- For example, here's a way to print the first ten lines of a file to @stdout@: --- --- >>> action = J.jet @Line (File "foo.txt") & J.limit 10 & J.sink stdout --- --- The code is using the 'J.jet' function to create a 'Jet' of 'Line' values --- (read using the default system encoding). 'J.jet' is part of the --- 'J.JetSource' helper typeclass. Meanwhile, 'J.sink' is part of the --- complementary 'J.JetSink' typeclass. --- --- Note also the use of '(&)', which is simply a flipped '($)'. I've found it --- useful to define forward-chained pipelines. --- --- If instead of printing to @stdout@ we wanted to store the lines in a list: --- --- >>> action = J.jet @Line (File "foo.txt") & J.limit 10 & J.toList --- --- Imagine we wanted to print the combined lines of two files, excepting the --- first 10 lines of each: --- --- >>> :{ --- action = --- do file <- J.each [File "foo.txt", File "bar.txt"] --- jet @Line file & J.drop 10 --- & J.sink stdout --- :} --- --- Here we are making use of the 'Monad' instance of 'Jet', which resembles --- that of conventional lists. We are mixing monadic do-blocks and conventional --- function application. Also we use 'J.each', a function which creates a 'Jet' --- out of any 'Foldable' container. --- --- 'Jet's are 'Monoid's too, so we could have written: --- --- >>> action = [File "foo.txt", File "bar.txt"] & foldMap (J.drop 10 . J.jet @Line) & J.sink stdout --- --- Here's an interesting use of 'sink'. Imagine we have a big utf8-encoded file --- and we want to split it into a number of files of no more than 100000 bytes --- each, with the extra condition that we don't want to split any line between --- two files. We could do it like this: --- --- >>> :{ --- action = --- let buckets = BoundedSize 100000 . File . ("result.txt." ++) . show <$> [1..] --- in jet (File "12999.txt.utf-8") --- & J.decodeUtf8 --- & J.lines --- <&> (\line -> J.lineToUtf8 line <> J.textToUtf8 J.newline) --- & J.sink buckets --- :} --- --- In this example we aren't using the default system encoding: instead of --- that, we are reading bytes, explicity decoding them with 'J.decodeUtf8' and --- finding 'J.lines'. Then we create a 'ByteBundle' for each 'Line' to signify --- that it shouldn't be broken, and end by writing to a sequence of --- 'BoundedSize' 'File's. --- -module Jet ( - -- * The Jet type - Jet, - J.run, - J.consume, - J.drain, - -- J.flatMap, - -- * Building Jets - J.each, - J.repeat, - J.repeatIO, - J.replicate, - J.replicateIO, - J.iterate, - J.iterateIO, - J.unfold, - J.unfoldIO, - J.untilEOF, - J.untilNothing, - -- * List-like functions - -- $listlike - J.toList, - J.length, - J.traverse, - J.traverse_, - J.for, - J.for_, - J.filter, - J.filterIO, - J.take, - J.limit, - J.takeWhile, - J.takeWhileIO, - J.drop, - J.dropWhile, - J.dropWhileIO, - J.mapAccum, - J.mapAccumIO, - J.intersperse, - -- * Zips - -- $zips - J.zip, - J.zipWith, - J.zipIO, - J.zipWithIO, - -- * Control operations - -- $control - J.withFile, - J.bracket, - J.bracket_, - J.bracketOnError, - J.finally, - J.onException, - -- ** Building your own - -- $doityourself - J.control, - J.unsafeCoerceControl, - J.control_, - J.unsafeCoerceControl_, - -- * Folding Jets - -- $folding - J.fold, - J.foldIO, - -- * Byte utils - J.bytes, - J.ChunkSize (..), - J.ByteBundle, - J.bundle, - J.bundleLength, - J.bundleBytes, - -- * Text and line utils - J.decodeUtf8, - J.encodeUtf8, - J.Line (Line), - J.lines, - J.unlines, - J.newline, - J.lineToText, - J.lineToUtf8, - J.textToLine, - J.textToUtf8, - J.stringToLine, - J.lineContains, - J.lineBeginsWith, - J.prefixLine, - -- * Concurrency - traverseConcurrently, - PoolConf, - defaults, - inputQueueSize, - numberOfWorkers, - outputQueueSize, - -- * Process invocation - throughProcess, - linesThroughProcess, - utf8LinesThroughProcess, - ProcConf, - bufferStdin, - readFromStderr, - handleExitCode, - -- * Conversion helpers - J.JetSource (..), - J.JetSink (..), - J.Sink (..), - J.File (..), - J.BoundedSize (..), - J.BucketOverflow (..), - -- * Some complicated stuff - -- $complicated - recast, - Splitter (..), - MealyIO(..), - SplitStepResult(..), - bytesOverBuckets, - byteBundlesOverBuckets, - Combiners, - combiners, - withCombiners, - withCombiners_, - combineIntoLists, - -- * Re-exports - -- $pipelines - (&), - (<&>), - -- $standardstreams - stdin, - stdout, - stderr, - -- $exceptions - T.UnicodeException, - -- $process - proc, - shell, - ) where - -import Data.Text.Encoding.Error qualified as T - -import System.IO (stdin, stdout, stderr) -import System.Process - -import Jet.Internal -import Jet.Internal qualified as J - -import Data.Function ((&)) -import Data.Functor ((<&>)) - --- $setup --- --- >>> :set -XTypeApplications --- >>> :set -XImportQualifiedPost --- >>> :set -XScopedTypeVariables --- >>> :set -XLambdaCase --- >>> :set -XNumDecimals --- >>> import Jet (Jet, (&)) --- >>> import Jet qualified as J --- >>> import Control.Foldl qualified as L --- >>> import Control.Concurrent --- >>> import Data.IORef --- >>> import Data.Text qualified as T - - --- $zips --- --- It's not possible to zip two 'Jet's together. But 'Jet's can be zipped with --- pure lists, or with lists of 'IO' actions. --- --- - --- $complicated --- --- I didn't manage to make this stuff simpler. --- - --- $pipelines --- I've found that the 'Data.Function.&' (reverse application) and 'Data.Functor.<&>' (reverse 'fmap') --- operators feel quite natural for building pipelines. - --- $standardstreams --- The standard streams, useful with functions like 'sink'. --- - --- $exceptions --- Thrown when decoding UTF8. --- - - --- $process --- Functions that create process specs for use with 'throughProcess'. For more control, import the whole of "System.Process". --- - --- $folding These functions can be used directly, but they're also useful for --- interfacing with the @Applicative@ folds from the --- [foldl](https://hackage.haskell.org/package/foldl) library, with the help of --- functions like @Control.Foldl.purely@ and @Control.Foldl.impurely@. --- --- @Applicative@ folds are useful because they let you run multiple --- \"analyses\" of a 'Jet' while going through it only once. - - --- $doityourself --- These are for advanced usage. --- --- Sometimes we want to lift some existing --- resource-handling operation not already covered, one that works with plain --- 'IO' values. These functions help with that. --- --- They have a linear type to statically forbid --- [\"funny\"](http://blog.ezyang.com/2012/01/monadbasecontrol-is-unsound/) --- operations like @\\x -> x *> x@ that disrupt proper threading of the --- consumer state. --- - - - --- $control --- Some 'Jet's must allocate resources to do its work. For example, opening a --- text file and yielding its lines. These resources must be promptly released --- when the 'Jet' itself finishes or the consumers stops it (for example, by --- using 'limit' on the 'Jet'). They must also be released in the face of --- exceptions. --- --- Here are various control operations like those from "Control.Exception", but --- lifted to work on 'Jet's. --- --- When put in a do-block, these operations \"protect\" every statement in the --- do-block below the operation itself. --- - - --- $listlike --- --- In these functions, the 'Jet' is working as a kind of \"effectful list\". --- The effects which produce the elements, and the effects with which we --- transform and consume the elements, are always 'IO' effects. --- --- Don't confuse these functions with similarly named functions from --- 'Data.Traversable' or 'Control.Monad', for which 'Jet' doesn't work as the --- \"container\", but as the Applicative/Monadic effect itself. --- - - +{-# LANGUAGE ImportQualifiedPost #-} +{-# LANGUAGE PatternSynonyms #-} + +-- | A streaming library build around the 'Jet' type, which behaves as a kind of \"effectful list\". +-- +-- For example, here's a way to print the first ten lines of a file to @stdout@: +-- +-- >>> action = J.jet @Line (File "foo.txt") & J.limit 10 & J.sink stdout +-- +-- The code is using the 'J.jet' function to create a 'Jet' of 'Line' values +-- (read using the default system encoding). 'J.jet' is part of the +-- 'J.JetSource' helper typeclass. Meanwhile, 'J.sink' is part of the +-- complementary 'J.JetSink' typeclass. +-- +-- Note also the use of '(&)', which is simply a flipped '($)'. I've found it +-- useful to define forward-chained pipelines. +-- +-- If instead of printing to @stdout@ we wanted to store the lines in a list: +-- +-- >>> action = J.jet @Line (File "foo.txt") & J.limit 10 & J.toList +-- +-- Imagine we wanted to print the combined lines of two files, excepting the +-- first 10 lines of each: +-- +-- >>> :{ +-- action = +-- do file <- J.each [File "foo.txt", File "bar.txt"] +-- jet @Line file & J.drop 10 +-- & J.sink stdout +-- :} +-- +-- Here we are making use of the 'Monad' instance of 'Jet', which resembles +-- that of conventional lists. We are mixing monadic do-blocks and conventional +-- function application. Also we use 'J.each', a function which creates a 'Jet' +-- out of any 'Foldable' container. +-- +-- 'Jet's are 'Monoid's too, so we could have written: +-- +-- >>> action = [File "foo.txt", File "bar.txt"] & foldMap (J.drop 10 . J.jet @Line) & J.sink stdout +-- +-- Here's an interesting use of 'sink'. Imagine we have a big utf8-encoded file +-- and we want to split it into a number of files of no more than 100000 bytes +-- each, with the extra condition that we don't want to split any line between +-- two files. We could do it like this: +-- +-- >>> :{ +-- action = +-- let buckets = BoundedSize 100000 . File . ("result.txt." ++) . show <$> [1..] +-- in jet (File "12999.txt.utf-8") +-- & J.decodeUtf8 +-- & J.lines +-- <&> (\line -> J.lineToUtf8 line <> J.textToUtf8 J.newline) +-- & J.sink buckets +-- :} +-- +-- In this example we aren't using the default system encoding: instead of +-- that, we are reading bytes, explicity decoding them with 'J.decodeUtf8' and +-- finding 'J.lines'. Then we create a 'ByteBundle' for each 'Line' to signify +-- that it shouldn't be broken, and end by writing to a sequence of +-- 'BoundedSize' 'File's. +module Jet + ( -- * The Jet type + Jet, + J.run, + J.consume, + J.drain, + -- J.flatMap, + + -- * Building Jets + J.each, + J.repeat, + J.repeatIO, + J.replicate, + J.replicateIO, + J.iterate, + J.iterateIO, + J.unfold, + J.unfoldIO, + J.untilEOF, + J.untilNothing, + + -- * List-like functions + -- $listlike + J.toList, + J.length, + J.traverse, + J.traverse_, + J.for, + J.for_, + J.filter, + J.filterIO, + J.take, + J.limit, + J.takeWhile, + J.takeWhileIO, + J.drop, + J.dropWhile, + J.dropWhileIO, + J.mapAccum, + J.mapAccumIO, + J.intersperse, + + -- * Zips + -- $zips + J.zip, + J.zipWith, + J.zipIO, + J.zipWithIO, + + -- * Control operations + -- $control + J.withFile, + J.bracket, + J.bracket_, + J.bracketOnError, + J.finally, + J.onException, + + -- ** Building your own + -- $doityourself + J.control, + J.control_, + + -- * Folding Jets + -- $folding + J.fold, + J.foldIO, + + -- * Byte utils + J.bytes, + J.ChunkSize (..), + J.ByteBundle, + J.bundle, + J.bundleLength, + J.bundleBytes, + + -- * Text and line utils + J.decodeUtf8, + J.encodeUtf8, + J.Line (Line), + J.lines, + J.unlines, + J.newline, + J.lineToText, + J.lineToUtf8, + J.textToLine, + J.textToUtf8, + J.stringToLine, + J.lineContains, + J.lineBeginsWith, + J.prefixLine, + + -- * Concurrency + traverseConcurrently, + PoolConf, + defaults, + inputQueueSize, + numberOfWorkers, + outputQueueSize, + + -- * Process invocation + throughProcess, + linesThroughProcess, + utf8LinesThroughProcess, + ProcConf, + bufferStdin, + readFromStderr, + handleExitCode, + + -- * Conversion helpers + J.JetSource (..), + J.JetSink (..), + J.Sink (..), + J.File (..), + J.BoundedSize (..), + J.BucketOverflow (..), + + -- * Some complicated stuff + -- $complicated + recast, + Splitter (..), + MealyIO (..), + SplitStepResult (..), + bytesOverBuckets, + byteBundlesOverBuckets, + Combiners, + combiners, + withCombiners, + withCombiners_, + combineIntoLists, + + -- * Re-exports + -- $pipelines + (&), + (<&>), + -- $standardstreams + stdin, + stdout, + stderr, + -- $exceptions + T.UnicodeException, + -- $process + proc, + shell, + ) +where + +import Data.Function ((&)) +import Data.Functor ((<&>)) +import Data.Text.Encoding.Error qualified as T +import Jet.Internal +import Jet.Internal qualified as J +import System.IO (stderr, stdin, stdout) +import System.Process + +-- $setup +-- +-- >>> :set -XTypeApplications +-- >>> :set -XImportQualifiedPost +-- >>> :set -XScopedTypeVariables +-- >>> :set -XLambdaCase +-- >>> :set -XNumDecimals +-- >>> import Jet (Jet, (&)) +-- >>> import Jet qualified as J +-- >>> import Control.Foldl qualified as L +-- >>> import Control.Concurrent +-- >>> import Data.IORef +-- >>> import Data.Text qualified as T + +-- $zips +-- +-- It's not possible to zip two 'Jet's together. But 'Jet's can be zipped with +-- pure lists, or with lists of 'IO' actions. + +-- $complicated +-- +-- I didn't manage to make this stuff simpler. + +-- $pipelines +-- I've found that the 'Data.Function.&' (reverse application) and 'Data.Functor.<&>' (reverse 'fmap') +-- operators feel quite natural for building pipelines. + +-- $standardstreams +-- The standard streams, useful with functions like 'sink'. + +-- $exceptions +-- Thrown when decoding UTF8. + +-- $process +-- Functions that create process specs for use with 'throughProcess'. For more control, import the whole of "System.Process". + +-- $folding These functions can be used directly, but they're also useful for +-- interfacing with the @Applicative@ folds from the +-- [foldl](https://hackage.haskell.org/package/foldl) library, with the help of +-- functions like @Control.Foldl.purely@ and @Control.Foldl.impurely@. +-- +-- @Applicative@ folds are useful because they let you run multiple +-- \"analyses\" of a 'Jet' while going through it only once. + +-- $doityourself +-- These are for advanced usage. +-- +-- Sometimes we want to lift some existing +-- resource-handling operation not already covered, one that works with plain +-- 'IO' values. These functions help with that. + +-- $control +-- Some 'Jet's must allocate resources to do its work. For example, opening a +-- text file and yielding its lines. These resources must be promptly released +-- when the 'Jet' itself finishes or the consumers stops it (for example, by +-- using 'limit' on the 'Jet'). They must also be released in the face of +-- exceptions. +-- +-- Here are various control operations like those from "Control.Exception", but +-- lifted to work on 'Jet's. +-- +-- When put in a do-block, these operations \"protect\" every statement in the +-- do-block below the operation itself. + +-- $listlike +-- +-- In these functions, the 'Jet' is working as a kind of \"effectful list\". +-- The effects which produce the elements, and the effects with which we +-- transform and consume the elements, are always 'IO' effects. +-- +-- Don't confuse these functions with similarly named functions from +-- 'Data.Traversable' or 'Control.Monad', for which 'Jet' doesn't work as the +-- \"container\", but as the Applicative/Monadic effect itself. diff --git a/lib/Jet/Internal.hs b/lib/Jet/Internal.hs --- a/lib/Jet/Internal.hs +++ b/lib/Jet/Internal.hs @@ -1,1782 +1,1787 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE BlockArguments #-} -{-# LANGUAGE DeriveFunctor #-} -{-# LANGUAGE DerivingStrategies #-} -{-# LANGUAGE DerivingVia #-} -{-# LANGUAGE GeneralizedNewtypeDeriving #-} -{-# LANGUAGE ImportQualifiedPost #-} -{-# LANGUAGE MultiWayIf #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE ViewPatterns #-} -{-# LANGUAGE TupleSections #-} -{-# LANGUAGE LinearTypes #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE LambdaCase #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE PatternSynonyms #-} -{-# LANGUAGE UndecidableInstances #-} -{-# LANGUAGE NamedFieldPuns #-} -{-# LANGUAGE GADTSyntax #-} -{-# LANGUAGE ExistentialQuantification #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE StandaloneKindSignatures #-} -{-# LANGUAGE PartialTypeSignatures #-} -{-# LANGUAGE StandaloneDeriving #-} -{-# LANGUAGE ApplicativeDo #-} -{-# OPTIONS_GHC -Wno-partial-type-signatures #-} --- | Tampering with the internals lets you write invalid 'Jet's that don't --- respect stop signals from consumers, so be careful. --- --- Also, the internals expose 'Line' and 'ByteBundle' as thin coats of paint --- over lazy text and lazy bytestring, respectively. -module Jet.Internal where - -import Control.Applicative -import Control.Monad -import Control.Monad.IO.Class -import Control.Exception -import Data.Foldable qualified -import Prelude hiding (traverse_, for_, filter, drop, dropWhile, fold, take, - takeWhile, unfold, zip, zipWith, filterM, lines, intersperse, unlines) -import Prelude qualified -import Unsafe.Coerce qualified -import System.IO (Handle, IOMode(..), hClose, openBinaryFile) -import System.IO qualified -import Data.Function ((&)) -import Data.Functor ((<&>)) - -import Data.Bifunctor -import Data.Text (Text) -import Data.Text qualified as T -import Data.Text.IO qualified as T -import Data.Text.Encoding qualified as T -import Data.Text.Encoding.Error qualified as T -import Data.Text.Lazy qualified as TL -import Data.Text.Lazy.Encoding qualified as TL -import Data.ByteString (ByteString) -import Data.ByteString qualified as B -import Data.ByteString.Lazy qualified as BL - -import Control.Concurrent -import Data.IORef -import Control.Concurrent.STM -import Control.Concurrent.MVar -import Control.Concurrent.Conceit -import Control.Concurrent.STM.TBMQueue -import Control.Concurrent.Async -import System.Process -import System.Exit -import Data.String (IsString(..)) -import Data.Typeable -import Data.Traversable qualified -import Data.Maybe -import Data.List qualified -import Data.Bifunctor (first) --- import Debug.Trace - --- $setup --- --- >>> :set -XTypeApplications --- >>> :set -XImportQualifiedPost --- >>> :set -XScopedTypeVariables --- >>> :set -XLambdaCase --- >>> :set -XNumDecimals --- >>> import Jet (Jet, (&)) --- >>> import Jet qualified as J --- >>> import Control.Foldl qualified as L --- >>> import Control.Concurrent --- >>> import Data.IORef --- >>> import Data.Text qualified as T - --- | A 'Jet' is a sequence of values produced through 'IO' effects. --- --- It allows consuming the elements as they are produced and doesn't force them --- to be present in memory all at the same time, unlike functions like --- 'Control.Monad.replicateM' from @base@. --- -newtype Jet a = Jet { - runJet :: forall s. (s -> Bool) -> (s -> a -> IO s) -> s -> IO s - } - --- | Maps over the yielded elements. '(<&>)' can be used to put the function last. --- --- >>> J.each "aa" <&> succ & J.toList --- "bb" -deriving stock instance Functor Jet - --- | Go through the elements produced by a 'Jet', while threading an --- state @s@ and possibly performing some effect. --- --- The caller is the one who chooses the type of the state @s@, and must pass --- an initial value for it. The state is kept in [weak-head normal form](https://en.wikibooks.org/wiki/Haskell/Graph_reduction#Weak_Head_Normal_Form). --- --- The caller must also provide a predicate on the state that informs the `Jet` --- when to stop producing values: whenever the predicate returns --- @True@. -run :: forall a s. Jet a -> (s -> Bool) -> (s -> a -> IO s) -> s -> IO s -run j = runJet j - --- | Like 'run', but always goes through all elements produced by the 'Jet'. --- --- Equivalent to @run (const False)@. -consume :: forall a s. Jet a -> (s -> a -> IO s) -> s -> IO s -consume j = run j (const False) - -for :: Jet a -> (a -> IO b) -> Jet b -for j k = zipWithIO (\() -> k) (Prelude.repeat (pure ())) j - -for_ :: Jet a -> (a -> IO b) -> IO () -for_ j k = consume j (\() -> void <$> k) () - --- | Apply an effectful transformation to each element in a 'Jet'. --- --- >>> :{ --- J.each "abc" --- & J.traverse (\c -> let c' = succ c in putStrLn ([c] ++ " -> " ++ [c']) *> pure c') --- & J.toList --- :} --- a -> b --- b -> c --- c -> d --- "bcd" --- -traverse :: (a -> IO b) -> Jet a -> Jet b -traverse = flip for - -traverse_ :: (a -> IO b) -> Sink a -traverse_ = flip for_ - --- | Go through the 'Jet' only for the 'IO' effects, discarding all yielded elements. -drain :: Sink a -drain = traverse_ pure - --- | Similar to the instance for pure lists, that generates combinations. --- --- >>> (,) <$> J.each "ab" <*> J.each "cd" & J.toList --- [('a','c'),('a','d'),('b','c'),('b','d')] --- -instance Applicative Jet where - pure i = Jet \stop step initial -> - if - | stop initial -> pure initial - | otherwise -> step initial i - Jet left <*> Jet right = Jet \stop step initial -> - -- Here we assume that the first Jet correctly handles the stop signal. - let step' f s a = step s (f a) - in left stop (\s f -> right stop (step' f) s) initial - --- | Similar to the instance for pure lists, that does search. --- --- >>> :{ --- do string <- J.each ["ab","cd"] --- J.each string --- & --- J.toList --- :} --- "abcd" -instance Monad Jet where - return = pure - Jet m >>= k = Jet \stop step initial -> - m stop (\s a -> runJet (k a) stop step s) initial - --- | --- >>> liftIO (putStrLn "foo") <> liftIO (putStrLn "bar") & J.toList --- foo --- bar --- [(),()] -instance MonadIO Jet where - liftIO action = Jet \stop step initial -> - if - | stop initial -> pure initial - | otherwise -> do - a <- action - step initial a - --- | 'Jet' concatenation. --- --- >>> J.each "ab" <> J.each "cd" & J.toList --- "abcd" -instance Semigroup (Jet a) where - Jet f1 <> Jet f2 = Jet \stop step s0 -> do - -- perhaps some of the stop checks are redundant, the first one in particular? - if - | stop s0 -> - pure s0 - | otherwise -> do - !s1 <- f1 stop step s0 - if - | stop s1 -> - pure s1 - | otherwise -> do - !s2 <- f2 stop step s1 - pure s2 - --- | 'mempty' is the empty 'Jet'. --- --- >>> mempty <> J.each "ab" <> mempty & J.toList --- "ab" -instance Monoid (Jet a) where - mempty = Jet \_ _ initial -> pure initial - --- | Same as 'Monoid'. -instance Alternative Jet where - (<|>) = (<>) - empty = mempty - --- | Same as 'Monoid' -instance MonadPlus Jet where - mzero = mempty - mplus = (<>) - --- | A failed pattern-match in a do-block produces 'mzero'. --- --- >>> :{ --- do Just c <- J.each [Nothing, Just 'a', Nothing, Just 'b'] --- pure c --- & J.toList --- :} --- "ab" --- -instance MonadFail Jet where - fail _ = mzero - --- | Build a 'Jet' from any 'Foldable' container --- --- >>> J.each [True,False] & J.toList --- [True,False] --- -each :: forall a f . Foldable f => f a -> Jet a -each (Data.Foldable.toList -> seed) = Jet \stop step -> - -- This could be done with Jet.unfold, but let's leave as it is. - let go b s = - if - | stop s -> - pure s - | otherwise -> - case b of - [] -> - pure s - -- see corresponding comment in unfold. - x : xs -> do - !s' <- step s x - go xs s' - in go seed - --- | --- --- >>> J.repeat True & J.take 2 & J.toList --- [True,True] --- -repeat :: a -> Jet a -repeat a = repeatIO (pure a) - - --- | --- --- >>> J.repeatIO (putStrLn "hi" *> pure True) & J.take 2 & J.toList --- hi --- hi --- [True,True] --- -repeatIO :: IO a -> Jet a -repeatIO action = untilNothing (fmap Just action) - --- | --- --- >>> J.replicate 2 True & J.toList --- [True,True] --- -replicate :: Int -> a -> Jet a -replicate n a = replicateIO n (pure a) - --- | --- >>> J.replicateIO 2 (putStrLn "hi" *> pure True) & J.toList --- hi --- hi --- [True,True] --- --- Don't confuse this with @Control.Monad.replicateM :: Int -> Jet a -> Jet [a]@ which has a combinatorial behavior. --- -replicateIO :: Int -> IO a -> Jet a -replicateIO n ioa = take n (repeatIO ioa) - --- | --- --- >>> J.iterate succ (1 :: Int) & J.take 2 & J.toList --- [1,2] --- -iterate :: (a -> a) -> a -> Jet a -iterate h = iterateIO (fmap pure h) - --- | --- --- >>> J.iterateIO (\x -> putStrLn "hi" *> pure (succ x)) (1 :: Int) & J.take 2 & J.toList --- hi --- [1,2] --- -iterateIO :: (a -> IO a) -> a -> Jet a -iterateIO h a = pure a <> unfoldIO (fmap (fmap (\x -> Just (x,x))) h) a - --- | --- >>> J.unfold (\case [] -> Nothing ; c : cs -> Just (c,cs)) "abc" & J.toList --- "abc" --- -unfold :: (b -> Maybe (a, b)) -> b -> Jet a -unfold h = unfoldIO (fmap pure h) - --- | --- >>> :{ --- J.unfoldIO (\x -> do putStrLn "hi" --- pure $ case x of --- [] -> Nothing --- c : cs -> Just (c,cs)) --- "abc" --- & J.toList --- :} --- hi --- hi --- hi --- hi --- "abc" --- -unfoldIO :: (b -> IO (Maybe (a, b))) -> b -> Jet a -unfoldIO h seed = Jet \stop step -> - let go b s = - if - | stop s -> - pure s - | otherwise -> do - next <- h b - case next of - Nothing -> - pure s - -- strictness only on the states. Good idea, or bad? - Just (a, !b') -> do - !s' <- step s a - go b' s' - in go seed - --- | --- >>> j = J.untilEOF System.IO.hIsEOF System.IO.hGetLine :: Handle -> Jet String --- -untilEOF :: (handle -> IO Bool) -> (handle -> IO a) -> handle -> Jet a -untilEOF hIsEOF' hGetLine' handle = untilNothing do - eof <- hIsEOF' handle - if - | eof -> - pure Nothing - | otherwise -> - Just <$> hGetLine' handle - --- | --- --- >>> :{ --- do ref <- newIORef "abc" --- let pop = atomicModifyIORef ref (\case [] -> ([], Nothing) --- x : xs -> (xs, Just x)) --- J.untilNothing pop & J.toList --- :} --- "abc" --- -untilNothing :: IO (Maybe a) -> Jet a -untilNothing action = unfoldIO (\() -> fmap (fmap (,())) action) () - --- | Convert to a regular list. This breaks streaming. --- --- >>> J.each "abc" & J.toList --- "abc" --- --- Alternatively, we can use 'fold' in combination with 'Control.Foldl.list' form the [foldl](https://hackage.haskell.org/package/foldl) library: --- --- >>> L.purely (J.fold (J.each "abc")) L.list --- "abc" --- --- which is more verbose, but more composable. -toList :: Jet a -> IO [a] -toList (Jet f) = do - as <- f (const False) (\xs x -> pure (x : xs)) [] - pure (reverse as) - --- | Returns the number of elements yielded by the 'Jet', exhausting it in the process. --- --- >>> J.each "abc" & J.length --- 3 --- --- Alternatively, we can use 'fold' in combination with 'Control.Foldl.length' form the [foldl](https://hackage.haskell.org/package/foldl) library: --- --- >>> L.purely (J.fold (J.each "abc")) L.length --- 3 --- --- which is more verbose, but more composable. -length :: Jet a -> IO Int -length (Jet f) = do - l <- f (const False) (\s _ -> pure (succ s)) 0 - pure l - -data Pair a b = Pair !a !b deriving Show - -pairExtract (Pair _ b) = b - -pairEnv (Pair a _) = a - -data Triple a b c = Triple !a !b !c - -tripleExtract (Triple _ _ c) = c - --- fromTuple :: (a, b) -> Pair a b --- fromTuple (a, b) -> Pair a b - --- | >>> J.each "abc" & J.drop 2 & J.toList --- "c" --- -drop :: Int -> Jet a -> Jet a -drop limit (Jet f) = Jet \stop step initial -> do - let stop' = stop . pairExtract - step' (Pair count s) a = - if - | count < limit -> do - pure (Pair (succ count) s) - | otherwise -> do - !s' <- step s a - pure (Pair count s') - initial' = Pair 0 initial - Pair _ final <- f stop' step' initial' - pure final - -data DropState = StillDropping | DroppingNoMore - --- | >>> J.each [1..5] & J.dropWhile (<3) & J.toList --- [3,4,5] --- -dropWhile :: (a -> Bool) -> Jet a -> Jet a -dropWhile p = dropWhileIO (fmap pure p) - -dropWhileIO :: (a -> IO Bool) -> Jet a -> Jet a -dropWhileIO p (Jet f) = Jet \stop step initial -> do - let stop' = stop . pairExtract - step' (Pair DroppingNoMore s) a = do - !s' <- step s a - pure (Pair DroppingNoMore s') - step' (Pair StillDropping s) a = do - keepDropping <- p a - if - | keepDropping -> - pure (Pair StillDropping s) - | otherwise -> do - !s' <- step s a - pure (Pair DroppingNoMore s') - initial' = (Pair StillDropping initial) - Pair _ final <- f stop' step' initial' - pure final - --- | >>> J.each "abc" & J.take 2 & J.toList --- "ab" --- -take :: Int -> Jet a -> Jet a -take limit (Jet f) = Jet \stop step initial -> do - let stop' (Pair count s) = - count >= limit || stop s - step' (Pair count s) a = do - !s' <- step s a - pure (Pair (succ count) s') - initial' = Pair 0 initial - Pair _ final <- f stop' step' initial' - pure final - --- | Synonym for 'take'. -limit :: Int -> Jet a -> Jet a -limit = take - -data TakeState = StillTaking | TakingNoMore - --- | >>> J.each [1..] & J.takeWhile (<5) & J.toList --- [1,2,3,4] --- -takeWhile :: (a -> Bool) -> Jet a -> Jet a -takeWhile p = takeWhileIO (fmap pure p) - -takeWhileIO :: (a -> IO Bool) -> Jet a -> Jet a -takeWhileIO p (Jet f) = Jet \stop step initial -> do - let stop' (Pair TakingNoMore _) = - True - stop' (Pair StillTaking s) = - stop s - step' (Pair internal s) a = do - keepTaking <- p a - if - | keepTaking -> do - !s' <- step s a - pure (Pair internal s') - | otherwise -> - pure (Pair TakingNoMore s) - initial' = Pair StillTaking initial - Pair _ final <- f stop' step' initial' - pure final - --- | --- >>> J.each "abc" & J.filter (=='a') & J.toList --- "a" --- -filter :: (a -> Bool) -> Jet a -> Jet a -filter p = filterIO (fmap pure p) - -filterIO :: (a -> IO Bool) -> Jet a -> Jet a -filterIO p (Jet f) = Jet \stop step initial -> do - let step' s a = do - shouldPass <- p a - if - | shouldPass -> do - !s' <- step s a - pure s' - | otherwise -> - pure s - f stop step' initial - --- | Behaves like a combination of 'fmap' and 'foldl'; it applies a function to --- each element of a structure passing an accumulating parameter from left to right. --- --- The resulting 'Jet' has the same number of elements as the original one. --- --- Unlike 'Data.Traversable.mapAccumL', it doesn't make the final state available. --- --- >>> J.each [1,2,3,4] & J.mapAccum (\a b -> (a + b,a)) 0 & J.toList --- [0,1,3,6] --- -mapAccum :: (a -> b -> (a, c)) -> a -> Jet b -> Jet c -mapAccum stepAcc = mapAccumIO (fmap (fmap pure) stepAcc) - -mapAccumIO :: (a -> b -> IO (a, c)) -> a -> Jet b -> Jet c -mapAccumIO stepAcc initialAcc (Jet f) = Jet \stop step initial -> do - let stop' = stop . pairExtract - step' (Pair acc s) b = do - (acc', c) <- stepAcc acc b - !s' <- step s c - pure (Pair acc' s') - initial' = Pair initialAcc initial - Pair _ final <- f stop' step' initial' - pure final - -data Touched = - NotYetTouched - | AlreadyTouched - --- TODO: there's a bug here!!!! - --- | --- >>> J.each "abc" & J.intersperse '-' & J.toList --- "a-b-c" --- -intersperse :: a -> Jet a -> Jet a -intersperse intrusion (Jet upstream) = Jet \stop step initial -> do - let stop' = stop . pairExtract - step' (Pair AlreadyTouched s) a = do - !s' <- step s intrusion - if - | stop s' -> - pure (Pair AlreadyTouched s') - | otherwise -> do - !s'' <- step s' a - pure (Pair AlreadyTouched s'') - step' (Pair NotYetTouched s) a = do - !s' <- step s a - pure (Pair AlreadyTouched s') - initial' = Pair NotYetTouched initial - Pair _ final <- upstream stop' step' initial' - pure final - --- | --- >>> J.each "abc" & J.zip [1..] & J.toList --- [(1,'a'),(2,'b'),(3,'c')] --- --- >>> J.each [1..] & J.zip "abc" & J.toList --- [('a',1),('b',2),('c',3)] --- -zip :: Foldable f => f a -> Jet b -> Jet (a, b) -zip = zipWith (,) - -zipWith :: Foldable f => (a -> b -> c) -> f a -> Jet b -> Jet c -zipWith zf (Data.Foldable.toList -> as0) = zipWithIO (fmap (fmap pure) zf) (fmap pure as0) - -zipIO :: Foldable f => f (IO a) -> Jet b -> Jet (a, b) -zipIO = zipWithIO (\x y -> pure (x, y)) - --- | --- Zips a list of 'IO' actions with a 'Jet', where the combining function can also have effects. --- --- If the list of actions is exhausted, the 'Jet' stops: --- --- >>> J.each [1..] <&> show & zipWithIO (\c1 c2 -> putStrLn (c1 ++ c2)) [pure "a", pure "b"] & J.toList --- a1 --- b2 --- [(),()] --- -zipWithIO :: Foldable f => (a -> b -> IO c) -> f (IO a) -> Jet b -> Jet c -zipWithIO zf (Data.Foldable.toList -> ioas0) (Jet f) = Jet \stop step initial -> do - let stop' (Pair [] _) = True - stop' (Pair _ s) = stop s - step' (Pair (ioa : ioas) s) b = do - a <- ioa - z <- zf a b - !s' <- step s z - pure (Pair ioas s') - step' (Pair [] _) _ = error "never happens" - initial' = Pair ioas0 initial - Pair _ final <- f stop' step' initial' - pure final - - --- | Opens a file and makes the 'Handle' available to all following statements --- in the do-block. --- --- Notice that it's often simpler to use the 'JetSource' (for reading) and --- 'JetSink' (for writing) instances of 'File'. -withFile :: FilePath -> IOMode -> Jet Handle -withFile path iomode = control @Handle (unsafeCoerceControl @Handle (System.IO.withFile path iomode)) - --- | --- --- >>> :{ --- do r <- J.bracket (putStrLn "allocating" *> pure "foo") (\r -> putStrLn $ "deallocating " ++ r) --- liftIO $ putStrLn $ "using resource " ++ r --- & drain --- :} --- allocating --- using resource foo --- deallocating foo --- -bracket :: forall a b . IO a -- ^ allocator - -> (a -> IO b) -- ^ finalizer - -> Jet a -bracket allocate free = control @a (unsafeCoerceControl @a (Control.Exception.bracket allocate free)) - -bracket_ :: forall a b . IO a -- ^ allocator - -> IO b -- ^ finalizer - -> Jet () -bracket_ allocate free = control_ (unsafeCoerceControl_ (Control.Exception.bracket_ allocate free)) - -bracketOnError :: forall a b . IO a -- ^ allocator - -> (a -> IO b) -- ^ finalizer - -> Jet a -bracketOnError allocate free = control @a (unsafeCoerceControl @a (Control.Exception.bracketOnError allocate free)) - --- | --- --- Notice how the finalizer runs even when we limit the 'Jet': --- --- >>> :{ --- do J.finally (putStrLn "hi") -- protects statements below --- liftIO (putStrLn "hey") --- J.each "abc" --- & J.limit 2 --- & J.toList --- :} --- hey --- hi --- "ab" --- --- But if the protected 'Jet' is not consumed at all, the finalizer might not run. --- --- >>> :{ --- do J.finally (putStrLn "hi") -- protects statements below --- liftIO (putStrLn "hey") --- J.each "abc" --- & J.limit 0 --- & J.toList --- :} --- "" --- -finally :: IO a -> Jet () -finally afterward = - control_ (unsafeCoerceControl_ (flip Control.Exception.finally afterward)) - -onException :: IO a -> Jet () -onException afterward = - control_ (unsafeCoerceControl_ (flip Control.Exception.onException afterward)) - --- | Lift a control operation (like 'Control.Exception.bracket') for which the --- callback uses the allocated resource. -control :: forall resource. (forall x. (resource -> IO x) %1 -> IO x) -> Jet resource -control f = - Jet \stop step initial -> - if - | stop initial -> - pure initial - | otherwise -> do - f (step initial) - --- | Lift a control operation (like 'Control.Exception.finally') for which the --- callback doesn't use the allocated resource. -control_ :: (forall x. IO x %1-> IO x) -> Jet () -control_ f = - Jet \stop step initial -> - if - | stop initial -> do - pure initial - | otherwise -> do - f (step initial ()) - --- | \"morally\", all control operations compatible with this library should --- execute the callback only once, which means that they should have a linear --- type. But because linear types are not widespread, they usually are given a --- less precise non-linear type. If you know what you are doing, use this --- function to give them a linear type. -unsafeCoerceControl :: forall resource . (forall x. (resource -> IO x) -> IO x) -> (forall x. (resource -> IO x) %1 -> IO x) -unsafeCoerceControl f = Unsafe.Coerce.unsafeCoerce f - --- | Line 'unsafeCoerceControl', for when the callback doesn't use the --- allocated resource. -unsafeCoerceControl_ :: (forall x. IO x -> IO x) -> (forall x. IO x %1 -> IO x) -unsafeCoerceControl_ f = Unsafe.Coerce.unsafeCoerce f - --- | --- --- >>> L.purely (J.fold (J.each "abc")) ((,) <$> L.list <*> L.length) --- ("abc",3) --- -fold :: Jet a -> (s -> a -> s) -> s -> (s -> r) -> IO r -fold (Jet f) step initial coda = do - r <- f (const False) (fmap (fmap pure) step) initial - pure $ coda r - --- | --- >>> L.impurely (J.foldIO (J.each "abc")) (L.FoldM (\() c -> putStrLn [c]) (pure ()) pure *> L.generalize L.length) --- a --- b --- c --- 3 --- -foldIO :: Jet a -> (s -> a -> IO s) -> IO s -> (s -> IO r) -> IO r -foldIO (Jet f) step initialIO coda = do - initial <- initialIO - r <- f (const False) step initial - coda r - - --- Byte Jets - --- https://stackoverflow.com/questions/49852060/how-to-choose-chunk-size-when-reading-a-large-file --- https://askubuntu.com/questions/641900/how-file-system-block-size-works --- https://stackoverflow.com/questions/1111661/8192-bytes-when-creating-file -data ChunkSize = - DefaultChunkSize - | ChunkSize Int - | ChunkSize1K - | ChunkSize4K - | ChunkSize8K - | ChunkSize16K - | ChunkSize1M - | ChunkSize2M - deriving Show - -chunkSize :: ChunkSize -> Int -chunkSize = \case - DefaultChunkSize -> 8192 - ChunkSize c -> c - ChunkSize1K -> 1024 - ChunkSize4K -> 4096 - ChunkSize8K -> 8192 - ChunkSize16K -> 16384 - ChunkSize1M -> 1048576 - ChunkSize2M -> 2097152 - --- | Helper multi-parameter typeclass for creating 'Jet' values out of a --- variety of common sources. --- --- Because there's no functional dependency, sometimes we need to use --- @TypeApplications@ to give the compiler a hint about the type of elements --- we want to produce. For example, here we want 'Line's and not, say, --- 'ByteString's: --- --- >>> action = J.jet @Line (File "foo.txt") & J.sink J.stdout --- --- -class JetSource a source where - jet :: source -> Jet a - -bytes :: ChunkSize -> Handle -> Jet ByteString -bytes (chunkSize -> count) handle = - untilEOF System.IO.hIsEOF (flip B.hGetSome count) handle - -instance JetSource ByteString Handle where - jet = bytes DefaultChunkSize - -instance JetSource a Handle => JetSource a File where - jet (File path) = do - handle <- withFile path ReadMode - jet handle - -accumByteLengths :: Jet ByteString -> Jet (Int,ByteString) -accumByteLengths = mapAccum (\acc bytes -> let acc' = acc + B.length bytes in (acc',(acc',bytes))) (0 :: Int) - -data AmIContinuing = Continuing - | NotContinuing deriving Show - --- | Splits a stream of bytes into groups bounded by maximum byte sizes. When --- one group \"fills up\", the next one is started. --- --- When the list of buckets sizes is exhausted, all incoming bytes are put into --- the same unbounded group. --- --- Useful in combination with 'recast'. -bytesOverBuckets :: [Int] -> Splitter ByteString ByteString -bytesOverBuckets buckets0 = MealyIO step mempty (pure (Pair NotContinuing buckets0)) - where - -- logStep s@(Pair c zzz) a = do - -- putStrLn "foooo!" - -- System.IO.hFlush System.IO.stdout - -- traceIO ("state: " ++ show c) - -- traceIO ("bucket: " ++ show (Prelude.take 2 zzz)) - -- traceIO ("input: " ++ show a) - -- r@(nexts, _) <- step s a - -- traceIO ("output: " ++ show nexts) - -- pure r - step :: Pair AmIContinuing [Int] -> ByteString -> IO (SplitStepResult ByteString, Pair AmIContinuing [Int]) - step splitterState b = do - (continueResult, Pair continuing' buckets', b') <- continue splitterState b - if | B.null b' -> - pure (continueResult, Pair continuing' buckets') - | otherwise -> do - (entiresResult, splitterState') <- makeEntires mempty b' buckets' - pure (continueResult <> entiresResult, splitterState') - continue :: Pair AmIContinuing [Int] -> ByteString -> IO (SplitStepResult ByteString, Pair AmIContinuing [Int], ByteString) - continue (Pair NotContinuing []) b = pure ( nextWith b , Pair NotContinuing [] , B.empty) - continue (Pair Continuing []) b = pure ( continueWith b , Pair Continuing [] , B.empty) - continue (Pair NotContinuing (bucket : buckets)) b = do - let blen = B.length b - -- traceIO ("b = " ++ show b ++ " bucket size= " ++ show bucket) - pure case compare blen bucket of - LT -> (nextWith b, Pair Continuing (bucket - blen : buckets), B.empty) - EQ -> (entireWith (singleton b), Pair NotContinuing buckets, B.empty) - GT -> let (left,right) = B.splitAt bucket b - in (entireWith (singleton left), Pair NotContinuing buckets, right) - continue (Pair Continuing (bucket : buckets)) b = do - let blen = B.length b - pure case compare blen bucket of - LT -> (continueWith b, Pair Continuing (bucket - blen : buckets), B.empty) - EQ -> (continueWith b, Pair NotContinuing buckets, B.empty) - GT -> let (left,right) = B.splitAt bucket b - in (continueWith left, Pair NotContinuing buckets, right) - makeEntires :: DList ByteString -> ByteString -> [Int] -> IO (SplitStepResult ByteString, Pair AmIContinuing [Int]) - makeEntires acc b [] = pure (entireWith acc <> nextWith b, Pair Continuing []) - makeEntires acc b (bucket : buckets) = do - let blen = B.length b - case compare blen bucket of - LT -> pure (entireWith acc <> nextWith b, Pair Continuing (bucket - blen : buckets)) - EQ -> pure (entireWith (acc <> singleton b), Pair NotContinuing buckets) - GT -> do let (left,right) = B.splitAt bucket b - makeEntires (acc <> singleton left) right buckets -- non-terminal - continueWith b = mempty { continuationOfPreviouslyStartedGroup = [b] } - entireWith bdf = mempty { entireGroups = fmap pure (closeDList bdf) } - nextWith b = mempty { startOfNewGroup = [b] } - --- | A sequence of bytes that we might want to keep together. -newtype ByteBundle = ByteBundle BL.ByteString deriving newtype (Show, Semigroup, Monoid) - --- | Constructs a 'ByteBundle' out of the bytes of some 'Foldable' container. -bundle :: Foldable f => f ByteString -> ByteBundle -bundle = ByteBundle . BL.fromChunks . Data.Foldable.toList - --- | Length in bytes. -bundleLength :: ByteBundle -> Int -bundleLength (ByteBundle value) = fromIntegral (BL.length value) -- Int64, but unlikely we'll reach the limit - -bundleBytes :: ByteBundle -> Jet ByteString -bundleBytes (ByteBundle value) = each (BL.toChunks value) - --- | Exception thrown when we try to write too much data in a size-bounded destination. -data BucketOverflow = BucketOverflow - deriving (Show, Typeable) - -instance Exception BucketOverflow - --- | Splits a stream of 'ByteBundles' into groups bounded by maximum byte --- sizes. Bytes belonging to the same 'ByteBundle' are always put in the same --- group. When one group \"fills up\", the next one is started. --- --- When the list of buckets sizes is exhausted, all incoming bytes are put into --- the same unbounded group. --- --- Useful in combination with 'recast'. --- --- __THROWS__: --- --- * 'BucketOverflow' exception if the size bound of a group turns out to be --- too small for holding even a single 'ByteBundle' value. --- --- -byteBundlesOverBuckets :: [Int] -> Splitter ByteBundle ByteString -byteBundlesOverBuckets buckets0 = MealyIO step mempty (pure (Pair NotContinuing buckets0)) - where - step :: Pair AmIContinuing [Int] -> ByteBundle -> IO (SplitStepResult ByteString, Pair AmIContinuing [Int]) - step (Pair splitterState []) (ByteBundle pieces) = - -- We assume [] means "infinite bucket" so once we enter it we'll only be able to continue. - pure ( case splitterState of - Continuing -> continueWith pieces - NotContinuing -> nextWith pieces - , Pair Continuing []) - step (Pair splitterState (bucket : buckets)) e@(ByteBundle pieces) = do - let elen = bundleLength e - case compare elen bucket of - LT -> pure ( case splitterState of - Continuing -> continueWith pieces - NotContinuing -> nextWith pieces - , Pair Continuing (bucket - elen : buckets) ) - EQ -> pure ( case splitterState of - Continuing -> continueWith pieces - NotContinuing -> entireWith pieces - , Pair NotContinuing buckets ) - -- NB: It's possible to close a bucket and open the next one in the same iteration. - GT -> case splitterState of - Continuing -> step (Pair NotContinuing buckets) e - -- If we are not continuing, that means that the brand-new bucket hasn't - -- enough space to hold a single entity. - NotContinuing -> throwIO BucketOverflow - continueWith bs = mempty { continuationOfPreviouslyStartedGroup = BL.toChunks bs } - entireWith pieces = mempty { entireGroups = [BL.toChunks pieces] } - nextWith bs = mempty { startOfNewGroup = BL.toChunks bs } - --- | Uses the default system locale. -instance JetSource Line Handle where - jet handle = - textToLine <$> untilEOF System.IO.hIsEOF T.hGetLine handle - --- --- --- Text Jets - --- | --- __THROWS__: --- --- * 'T.UnicodeException' -decodeUtf8 :: Jet ByteString -> Jet Text -decodeUtf8 (Jet f) = Jet \stop step initial -> do - let stop' = stop . pairExtract - step' (Pair leftovers s) bytes = do - T.Some !text !_ !leftovers' <- pure $ T.streamDecodeUtf8 bytes - !s' <- step s text - pure (Pair leftovers' s') - initial' = Pair leftovers0 initial - Pair leftovers final <- f stop' step' initial' - T.Some !_ !bytes !_ <- pure $ T.streamDecodeUtf8 B.empty - if | not (B.null bytes) -> - throwIO (T.DecodeError "Unconsumed leftovers at end." Nothing) - | otherwise -> - pure final - where - leftovers0 = - let T.Some _ _ g = T.streamDecodeUtf8 B.empty - in g - -encodeUtf8 :: Jet Text -> Jet ByteString -encodeUtf8 = fmap T.encodeUtf8 - --- | A line of text. --- --- While it is guaranteed that the 'Line's coming out of the 'lines' function --- do not contain newlines, that invariant is not otherwise enforced. -newtype Line = Line_ TL.Text - deriving newtype (Eq,Ord,Semigroup,Monoid,Show,IsString) - --- https://ghc.gitlab.haskell.org/ghc/doc/users_guide/exts/pattern_synonyms.html - --- | Unidirectional pattern that allows converting a 'Line' into a 'Text' --- during pattern-matching. -pattern Line text <- Line_ (TL.toStrict -> text) - --- | Converts a 'Line' back to text, without adding the newline. -lineToText :: Line -> Text -lineToText (Line_ text) = TL.toStrict text - --- | Converts a 'Line' to an utf8-encdoed 'ByteBundle', without adding the newline. -lineToUtf8 :: Line -> ByteBundle -lineToUtf8 (Line_ l) = TL.toChunks l <&> T.encodeUtf8 & bundle - -textToLine :: Text -> Line -textToLine = Line_ . TL.fromStrict - --- | @Data.Text.singleton '\\n'@ -newline :: Text -newline = T.singleton '\n' - -textToUtf8 :: Text -> ByteBundle -textToUtf8 t = ByteBundle (t & T.encodeUtf8 & BL.fromStrict) - -lineContains :: Text -> Line -> Bool -lineContains t (Line_ l) = TL.isInfixOf (TL.fromStrict t) l - -lineBeginsWith :: Text -> Line -> Bool -lineBeginsWith t (Line_ l) = TL.isPrefixOf (TL.fromStrict t) l - --- | Adds the 'Text' to the beginning of the 'Line'. -prefixLine :: Text -> Line -> Line -prefixLine t (Line_ l) = Line_ (TL.fromChunks (t : TL.toChunks l)) - --- textToLine :: Text -> Line --- textToLine text --- | Just _ <- T.find (=='\n') text = throw NewlineForbidden --- | otherwise = Line_ (removeTrailingCarriageReturn text) - -stringToLine :: String -> Line -stringToLine = Line_ . TL.pack - --- withLineText :: (Text -> r) -> Line -> r --- withLineText f (Line text) = f text - -isEmptyLine :: Line -> Bool -isEmptyLine (Line_ text) = TL.null text - -emptyLine :: Line -emptyLine = Line_ TL.empty - --- | Exception thrown when we find newlines in functions which don't accept them. --- --- A direct copy of the @NewlineForbidden@ exception from the [turtle](https://hackage.haskell.org/package/turtle) package. -data NewlineForbidden = NewlineForbidden - deriving (Show, Typeable) - -instance Exception NewlineForbidden - -removeTrailingCarriageReturn :: Text -> Text -removeTrailingCarriageReturn text - | T.null text = text - | T.last text == '\r' = T.init text - | otherwise = text - -lines :: Jet Text -> Jet Line -lines (Jet f) = Jet \stop step initial -> do - let stop' = stop . pairExtract - findLinesInCurrentBlock text - | T.null text = - [] - | otherwise = - map (textToLine . removeTrailingCarriageReturn) (T.lines text) - ++ - if - | T.last text == '\n' -> - [mempty] - | otherwise -> - [] - step' (Pair lineUnderConstruction s) (findLinesInCurrentBlock -> linesInCurrentBlock) = do - case linesInCurrentBlock of - [] -> do - pure (Pair lineUnderConstruction s) - [l] -> do - pure (Pair (lineUnderConstruction <> singleton l) s) - l : rest@(x : xs) -> do - -- Ineficcient mconcat, better strictify a lazy text here? - let completedLine = mconcat $ runDList lineUnderConstruction [l] - s' <- downstream stop step (completedLine : init rest) s - pure (Pair (singleton (last linesInCurrentBlock)) s') - initial' = Pair mempty initial - Pair (mconcat . closeDList -> lineUnderConstruction) final <- f stop' step' initial' - if - | stop final -> - pure final - | isEmptyLine lineUnderConstruction -> - pure final - | otherwise -> - step final lineUnderConstruction - -unlines :: Jet Line -> Jet Text -unlines j = do - Line text <- j - pure text <> pure (T.singleton '\n') - -downstream :: (s -> Bool) -> (s -> x -> IO s) -> [x] -> s -> IO s -downstream stop step = go - where - go [] s = - pure s - go (x : xs) s - | stop s = - pure s - | otherwise = do - !s' <- step s x - go xs s' - --- General sinks - --- | A function that consumes a 'Jet' totally or partially, without returning a result. -type Sink a = Jet a -> IO () - --- | Helper multi-parameter typeclass for creating 'Jet'-consuming functions --- out of a variety of common destinations. --- --- >>> J.each ["aaa","bbb","ccc"] <&> J.stringToLine & J.sink J.stdout --- aaa --- bbb --- ccc --- -class JetSink a target where - sink :: target -> Sink a - -instance JetSink ByteString Handle where - sink handle j = for_ j (B.hPut handle) - -instance JetSink a Handle => JetSink a File where - sink (File path) j = System.IO.withFile path System.IO.WriteMode \handle -> - sink handle j - --- | Uses the default system locale. Adds newlines. -instance JetSink Line Handle where - sink handle = traverse_ (T.hPutStrLn handle . lineToText) - --- | Uses the default system locale. -instance JetSink Text Handle where - sink handle = traverse_ (T.hPutStr handle) - --- | 'FilePaths' are plain strings. This newtype provides a small measure of --- safety over them. -newtype File = File { getFilePath :: FilePath } deriving Show - --- | The maximum size in bytes of some destination into which we write the --- bytes produced by a 'Jet'. -data BoundedSize x = BoundedSize Int x deriving stock (Show,Read) - -instance JetSink ByteBundle Handle where - sink handle j = traverse_ (B.hPut handle) do - s <- j - bundleBytes s - --- | Distributes incoming bytes through a sequence of files. Once a file is --- full, we start writing the next one. -instance JetSink ByteString [BoundedSize File] where - sink bucketFiles j = - withCombiners_ - (\handle b -> B.hPut handle b) - hClose - (makeAllocator <$> bucketFiles) - (\combiners -> drain $ recast (bytesOverBuckets bucketSizes) combiners j) - where - bucketSizes = map (\(BoundedSize size _) -> size) bucketFiles - --- | Distributes incoming bytes through a sequence of files. Once a file is --- full, we start writing the next one. --- --- Each 'ByteBundle' value is garanteed to be written to a single file. If a --- file turns out to be too small for even a single 'ByteBundle' value, a --- 'BucketOverflow' exception is thrown. -instance JetSink ByteBundle [BoundedSize File] where - sink bucketFiles j = - withCombiners_ - (\handle b -> B.hPut handle b) - hClose - (makeAllocator <$> bucketFiles) - (\combiners -> drain $ recast (byteBundlesOverBuckets bucketSizes) combiners j) - where - bucketSizes = map (\(BoundedSize size _) -> size) bucketFiles - -makeAllocator :: BoundedSize File -> IO Handle -makeAllocator (BoundedSize _ (File path)) = openBinaryFile path WriteMode - --- DList helper -newtype DList a = DList { runDList :: [a] -> [a] } - -instance Semigroup (DList a) where - DList a1 <> DList a2 = DList (a1 . a2) - -instance Monoid (DList a) where - mempty = DList id - -makeDList :: [a] -> DList a -makeDList as = DList \xs -> as ++ xs - -closeDList :: DList a -> [a] -closeDList (DList f) = f [] - -singleton :: a -> DList a -singleton a = DList $ (a :) - --- --- concurrency - --- | Process the values yielded by the upstream 'Jet' in a concurrent way, --- and return the results in the form of another 'Jet' as they are produced. --- --- __NB__: this function might scramble the order of the returned values. Right --- now there isn't a function for unscrambling them. --- --- >>> :{ --- J.each [(3,'a'), (2,'b'), (1,'c')] --- & J.traverseConcurrently (numberOfWorkers 10) (\(d,c) -> threadDelay (d*1e5) *> pure c) --- & J.toList --- :} --- "cba" --- --- What happens if we 'limit' the resulting 'Jet' and we reach that limit, or --- if we otherwise stop consuming the 'Jet' before it gets exhausted? In those --- cases, all pending @IO b@ tasks are cancelled. --- --- >>> :{ --- J.each [(9999,'a'), (2,'b'), (1,'c')] --- & J.traverseConcurrently (numberOfWorkers 10) (\(d,c) -> threadDelay (d*1e5) *> pure c) --- & J.take 2 --- & J.toList --- :} --- "cb" --- -traverseConcurrently :: (PoolConf -> PoolConf) -> (a -> IO b) -> Jet a -> Jet b --- TODO: --- It would be nice to have 0-lengh channels for which one side blocks until --- the other side takes the job. -traverseConcurrently adaptConf makeTask upstream = Jet \stop step initial -> do - if - -- If we know we aren't going to do any work, don't bother starting the - -- whole boondoggle. - | stop initial -> - pure initial - | otherwise -> do - -- At this point we know we should do at least one step. - let PoolConf {_inputQueueSize,_numberOfWorkers,_outputQueueSize} = adaptConf defaultPoolConf - input <- newTBMQueueIO _inputQueueSize - inputQueueWriterShouldStop <- newIORef False - aliveWorkers <- newIORef _numberOfWorkers - output <- newTBMQueueIO _outputQueueSize - let - -- The inputQueueWriter should *not* be interrupted aynchronously. - -- After each iteration, it reads the IORef to see if it should stop. - -- Once it stops, it closes the input queue. - inputQueueWriter = do - run - upstream - id - (\_ a -> do - atomically $ writeTBMQueue input (makeTask a) - readIORef inputQueueWriterShouldStop) - False - atomically $ closeTBMQueue input - -- Workers *can* be interrupted asynchronously. - worker = do - mtask <- atomically $ readTBMQueue input - case mtask of - Nothing -> do - remaining <- do - atomicModifyIORef' aliveWorkers \count -> - let count' = pred count - in (count', count') - if - | remaining == 0 -> do - atomically $ closeTBMQueue output - | otherwise -> do - pure () - Just task -> do - result <- task - atomically $ writeTBMQueue output result - worker - outputQueueReader s = do - if - | stop s -> do - -- tell the inserter from upstream that it should stop. is this enough? - writeIORef inputQueueWriterShouldStop True - atomically $ closeTBMQueue input -- perhaps unnecessary? - pure s - | otherwise -> do - mresult <- atomically $ readTBMQueue output - case mresult of - Nothing -> do - pure s - Just result -> do - !s' <- step s result - outputQueueReader s' - runConcurrently $ - Concurrently do - inputQueueWriter - *> - Concurrently do - finalLeft <- do - runConceit $ - -- The worker pool is always killed when the output reader finishes, - -- but for the "happy path" the workers will already be dead. - Conceit (Right <$> replicateConcurrently_ _numberOfWorkers worker) - *> - -- This Left is what kills the worker pool. - Conceit (Left <$> outputQueueReader initial) - case finalLeft of - Right () -> do - error "never happens, the Left always wins" - Left final -> do - pure final - --- | Configuration record for the worker pool. -data PoolConf = PoolConf { - _inputQueueSize :: Int, - _numberOfWorkers :: Int, - _outputQueueSize :: Int - } deriving Show - -defaultPoolConf = PoolConf { - _inputQueueSize = 1, - _numberOfWorkers = 1, - _outputQueueSize = 1 - } - --- | Size of the waiting queue into the worker pool. The default is @1@. -inputQueueSize :: Int -> PoolConf -> PoolConf -inputQueueSize size poolConf = poolConf { _inputQueueSize = size } - --- | The size of the worker pool. The default is @1@. -numberOfWorkers :: Int -> PoolConf -> PoolConf -numberOfWorkers number poolConf = poolConf { _numberOfWorkers = number } - --- | Size of the queue holding results out of the working pool before they --- are yielded downstream. The default is @1@. -outputQueueSize :: Int -> PoolConf -> PoolConf -outputQueueSize size poolConf = poolConf { _outputQueueSize = size } - --- | An alias for 'id'. Useful with functions like 'traverseConcurrently' and --- 'throughProcess', for which it means \"use the default configuration\". -defaults :: a -> a -defaults = id - --- --- process invocation - --- | Feeds the upstream 'Jet' to an external process' @stdin@ and returns the --- process' @stdout@ as another @Jet@. The feeding and reading of the standard --- streams is done concurrently in order to avoid deadlocks. --- --- What happens if we 'limit' the resulting 'Jet' and we reach that limit, or --- if we otherwise stop consuming the 'Jet' before it gets exhausted? In those --- cases, the external process is promptly terminated. -throughProcess :: (ProcConf -> ProcConf) -> CreateProcess -> Jet ByteString -> Jet ByteString -throughProcess adaptConf = throughProcess_ (adaptConf defaultProcConf) - --- | Like 'throughProcess', but feeding and reading 'Line's using the default --- system encoding. --- --- >>> :{ --- J.each ["aaa","bbb","ccc"] --- <&> J.stringToLine --- & linesThroughProcess defaults (shell "cat") --- & J.toList --- :} --- ["aaa","bbb","ccc"] --- --- An example of not reading all the lines from a long-lived process that gets cancelled: --- --- >>> :{ --- mempty --- & linesThroughProcess defaults (shell "{ printf \"aaa\\nbbb\\nccc\\n\" ; sleep infinity ; }") --- & J.limit 2 --- & J.toList --- :} --- ["aaa","bbb"] --- -linesThroughProcess :: (ProcConf -> ProcConf) -> CreateProcess -> Jet Line -> Jet Line -linesThroughProcess adaptConf procSpec = do - let textLinesProcConf = (adaptConf defaultProcConf) { - _writeToStdIn = T.hPutStrLn, - _readFromStdout = T.hGetLine - } - fmap textToLine . throughProcess_ textLinesProcConf procSpec . fmap lineToText - --- | Like 'throughProcess', but feeding and reading 'Line's encoded in UTF8. -utf8LinesThroughProcess :: (ProcConf -> ProcConf) -> CreateProcess -> Jet Line -> Jet Line -utf8LinesThroughProcess adaptConf procSpec = do - lines . decodeUtf8 . throughProcess adaptConf procSpec . encodeUtf8 . unlines - -throughProcess_ :: forall a b . ProcConf_ a b -> CreateProcess -> Jet a -> Jet b -throughProcess_ procConf procSpec upstream = Jet \stop step initial -> do - let ProcConf_ {_bufferStdin, _writeToStdIn, _readFromStdout,_readFromStderr, _handleExitCode} = procConf - if - -- If we know we aren't going to do any work, don't bother starting the - -- whole boondoggle. - | stop initial -> - pure initial - | otherwise -> do - let procSpec' = procSpec { - std_in = CreatePipe, - std_out = CreatePipe, - std_err = CreatePipe - } - input <- newTBMQueueIO @a 1 - inputQueueWriterShouldStop <- newIORef False - -- remember to drain stderr concurrently with stdout... - let inputQueueWriter = do - run - upstream - id - (\_ a -> do - atomically $ writeTBMQueue input a - readIORef inputQueueWriterShouldStop) - False - atomically $ closeTBMQueue input - finalEither <- - runConcurrently $ - Concurrently do - inputQueueWriter - *> - Concurrently do - withCreateProcess procSpec' \(Just stdin') (Just stdout') (Just stderr') phandle -> do - when (not _bufferStdin) (System.IO.hSetBuffering stdin' System.IO.NoBuffering) - let stdinWriter = do - ma <- atomically $ readTBMQueue input - case ma of - Nothing -> do - hClose stdin' - Just a -> do - _writeToStdIn stdin' a - stdinWriter - stderrReader = do - untilEOF System.IO.hIsEOF _readFromStdout stderr' & drain - stdoutReader s = do - if | stop s -> do - writeIORef inputQueueWriterShouldStop True - pure (Left s) - | otherwise -> do - eof <- System.IO.hIsEOF stdout' - if - | eof -> do - writeIORef inputQueueWriterShouldStop True - exitCode <- waitForProcess phandle - _handleExitCode exitCode - pure (Right s) - | otherwise -> do - b <- _readFromStdout stdout' - !s' <- step s b - stdoutReader s' - runConceit $ - _Conceit do stdinWriter - *> - _Conceit do stderrReader - *> - Conceit do stdoutReader initial - pure (either id id finalEither) - --- | Configuration record with some extra options in addition to those in "CreateProcess". -type ProcConf = ProcConf_ ByteString ByteString -data ProcConf_ a b = ProcConf_ { - _bufferStdin :: Bool, - _writeToStdIn :: Handle -> a -> IO (), - _readFromStdout :: Handle -> IO b, - _readFromStderr :: Handle -> IO (), - _handleExitCode :: ExitCode -> IO () - } - -defaultProcConf :: ProcConf -defaultProcConf = ProcConf_ { - _bufferStdin = False, - _writeToStdIn = B.hPut, - _readFromStdout = flip B.hGetSome 8192, - _readFromStderr = void . T.hGetLine , - _handleExitCode = \exitCode -> case exitCode of - ExitFailure _ -> throwIO exitCode - ExitSuccess -> pure () - } - --- | Should we buffer the process' @stdin@? Usually should be 'True' for --- interactive scenarios. --- --- By default, 'False'. -bufferStdin :: Bool -> ProcConf -> ProcConf -bufferStdin doBuffering procConf = procConf { _bufferStdin = doBuffering } - --- | Sets the function that reads a single line of output from the process --- @stderr@. It's called repeatedly until @stderr@ is exhausted. The reads are --- done concurrently with the reads from @stdout@. --- --- By default, lines of text are read using the system's default encoding. --- --- This is a good place to throw an exception if we don't like what comes out --- of @stderr@. -readFromStderr :: (Handle -> IO ()) -> ProcConf -> ProcConf -readFromStderr readFunc procConf = procConf { _readFromStderr = readFunc } - --- | Sets the function that handles the final `ExitCode` of the process. --- --- The default behavior is to throw the `ExitCode` as an exception if it's not --- a success. -handleExitCode :: (ExitCode -> IO ()) -> ProcConf -> ProcConf -handleExitCode handler procConf = procConf { _handleExitCode = handler } - --- --- --- complicated stufff - -data AreWeInsideGroup foldState = OutsideGroup - | InsideGroup !foldState - -data RecastState foldState = RecastState !(AreWeInsideGroup foldState) [IO foldState] - --- | This is a complex, unwieldly, yet versatile function. It can be used to --- define grouping operations, but also for decoding and other purposes. --- --- Groups are delimited in the input 'Jet' using the 'Splitter', and the --- contents of those groups are then combined using 'Combiners'. The result of --- each combiner is yielded by the return 'Jet'. --- --- If the list of combiners is finite and becomes exhausted, we stop splitting --- and the return 'Jet' stops. -recast :: forall a b c . Splitter a b -> Combiners b c -> Jet a -> Jet c -recast (MealyIO splitterStep splitterCoda splitterAlloc) - (Combiners foldStep foldCoda foldAllocs0) - (Jet upstream) = Jet \stop step initial -> do - initialSplitterState <- splitterAlloc - let -- When to stop? Either downstream says we need to stop, - -- or we are outside a group and there isn't another group consumer we - -- can use to process the next one. - stop' (Triple _ (RecastState OutsideGroup []) _) = True - stop' (Triple _ _ s) = stop s - - step' (Triple splitterState recastState s) a = do - (splitResult, splitterState') <- splitterStep splitterState a - Pair recastState' s' <- advanceRecast splitResult recastState s - pure (Triple splitterState' recastState' s') - - advanceRecast ssr@(SplitStepResult {continuationOfPreviouslyStartedGroup, entireGroups, startOfNewGroup}) (RecastState areWeInside foldAllocs) s = do - case (areWeInside, entireGroups, startOfNewGroup) of - -- If there aren't any new groups and we don't start an incomplete one, just advance the current fold - (InsideGroup foldState, [], []) -> do - -- traceIO $ "recast inside group just continuing" - foldState' <- advanceGroupWithougClosing foldState continuationOfPreviouslyStartedGroup - pure (Pair (RecastState (InsideGroup foldState') foldAllocs) s) -- main state didn't change - (InsideGroup foldState, _, _) -> do - -- traceIO $ "recast inside group closing" - !c <- processSingleGroup foldState continuationOfPreviouslyStartedGroup - !s' <- step s c - if - | stop s' -> do - -- traceIO $ "recast inside group pure" - pure (Pair (RecastState OutsideGroup foldAllocs) s') - | otherwise -> do - -- traceIO $ "recast inside group advancing" - advanceRecast ssr (RecastState OutsideGroup foldAllocs) s' - -- if we are outside of a group, the "continuationOfPreviouslyStartedGroup" is ignored. - (OutsideGroup, _, _) -> do - -- traceIO $ "recast outside group" - -- doens't return foldState becasue we close the groups - Pair foldAllocs' s' <- processEntireGroups foldAllocs s entireGroups - bail <- pure (Pair (RecastState OutsideGroup foldAllocs') s') - if - | stop s' -> do - pure bail - | otherwise -> do - case startOfNewGroup of - [] -> do - pure bail - (_ : _) -> do - case foldAllocs of - [] -> do - pure bail - alloc : allocs -> do - -- traceIO $ "recast we should be allocating here" - -- there is a next group, so let's begin it - !foldState0 <- alloc - foldState <- processBeginNextGroup foldState0 startOfNewGroup - pure (Pair (RecastState (InsideGroup foldState) allocs) s') - -- foldM ? - advanceGroupWithougClosing :: _ -> [b] -> IO _ - advanceGroupWithougClosing foldState [] = - pure foldState - advanceGroupWithougClosing foldState (b:bs) = do - !foldState' <- foldStep foldState b - advanceGroupWithougClosing foldState' bs - processEntireGroups :: [IO _] -> _ -> [[b]] -> IO (Pair [IO _] _) - -- We can't go on if there aren't any more groups - processEntireGroups allocs s [] = do - pure (Pair allocs s) - -- We can't go on if there aren't any more fold initial state allocs - processEntireGroups [] s _ = do - pure (Pair [] s) - processEntireGroups (alloc : allocs) s (bs:bss) = do - !foldState0 <- alloc - !c <- processSingleGroup foldState0 bs -- a single step downstream - !s' <- step s c - if - | stop s' -> do - pure (Pair allocs s') - | otherwise -> do - processEntireGroups allocs s' bss - -- a whole fold is processed here - processSingleGroup :: _ -> [b] -> IO c - processSingleGroup foldState [] = do - foldCoda foldState - processSingleGroup foldState (b:bs) = do - !foldState' <- foldStep foldState b - processSingleGroup foldState' bs - processBeginNextGroup :: _ -> [b] -> IO _ - processBeginNextGroup foldState [] = do - pure foldState - processBeginNextGroup foldState (b:bs) = do - !foldState' <- foldStep foldState b - processBeginNextGroup foldState' bs - initial' = Triple initialSplitterState (RecastState OutsideGroup foldAllocs0) initial - Triple splitterState recastState final <- upstream stop' step' initial' - -- What happens if there's a fold ongoing when we stop? Right now we always close it, which seems to be a reasonable - -- action (because the fold coda might hide a finalizer). - -- - -- Also, when can it happen that we reach this point with an ongoing fold? - -- If I understand correctly: - -- - it can only happen when the upstream closes and leaves the fold open. - -- - it can't (?) happen when the consumer stops early. - let closePendingFold = \case - RecastState OutsideGroup _ -> do - pure () - RecastState (InsideGroup foldState) _ -> do - _ <- foldCoda foldState - pure () - if - | stop final -> do - closePendingFold recastState - pure final - | otherwise -> do - splitResult <- splitterCoda splitterState - -- We discard the "begins next group"; it doesn't make sense in this final step. - Pair recastState' final' <- advanceRecast (splitResult { startOfNewGroup = [] }) recastState final - if | stop final' -> do - -- TODO: - -- should we dealloc here? Maybe there's a fold reaminging... we should close it. See below. - closePendingFold recastState' - pure final' - | otherwise -> do - case recastState' of - RecastState OutsideGroup _ -> do - -- traceIO $ "final! outside group" - pure final' - RecastState (InsideGroup foldState) _ -> do - -- traceIO $ "final! inside group" - c <- foldCoda foldState - final'' <- step final' c - pure final'' - --- | A 'Combiners' value knows how to process a sequence of groups, while --- keeping a (existentially hidden) state for each group. --- --- Very much like a @FoldM IO@ from the --- [foldl](https://hackage.haskell.org/package/foldl-1.4.12/docs/Control-Foldl.html#t:FoldM) --- library, but \"restartable\" with a list of starting states. --- --- For converting one into the other, this function should do the trick: --- --- > \(L.FoldM step allocator coda) -> combiners step coda (Prelude.repeat allocator) -data Combiners a b where - Combiners :: (s -> a -> IO s) -> (s -> IO b) -> [IO s] -> Combiners a b - -deriving stock instance Functor (Combiners a) - --- | Constructor for 'Combiners' values. -combiners :: forall s a b r -- ^ foo - . (s -> a -> IO s) -- ^ Step function that threads the state @s@. - -> (s -> IO b) -- ^ Coda invoked when a group closes. - -> [IO s] -- ^ Actions that produce the initial states @s@ for processing each group. - -> Combiners a b -combiners = Combiners - --- | A simpler version of 'withCombiners' that doen't thread a state; it merely --- allocates and deallocates the resource @h@. -withCombiners_ :: forall h a r - . (h -> a -> IO ()) -- ^ Step function that accesses the resource @h@. - -> (h -> IO ()) -- ^ Finalizer to run after closing each group, and also in the case of an exception. - -> [IO h] -- ^ Actions that allocate a sequence of resources @h@. - -> (Combiners a () -> IO r) -- ^ The 'Combiners' value should be consumed linearly. - -> IO r -withCombiners_ step finalize allocators = do - withCombiners - (\h () a -> step h a) - (\_ () -> pure ()) - finalize - (do allocator <- allocators - pure (allocator, \_ -> pure ())) - --- | 'Combiners' thread a state @s@ while processing each group. Sometimes, in --- addition to that, we want to allocate a resource @h@ when we start --- processing a group, and deallocate it after we finish processing the group --- or an exception is thrown. The typical example is allocating a 'Handle' for --- writing the elements of the group as they arrive. -withCombiners - :: forall h s a b r . - (h -> s -> a -> IO s) -- ^ Step function that accesses the resource @h@ and threads the state @s@. - -> (h -> s -> IO b) -- ^ Coda invoked when a group closes. - -> (h -> IO ()) -- ^ Finalizer to run after each coda, and also in the case of an exception. - -> [(IO h, h -> IO s)] -- ^ Actions that allocate a sequence of resources @h@ and produce initial states @s@ for processing each group. - -> (Combiners a b -> IO r) -- ^ The 'Combiners' value should be consumed linearly. - -> IO r -withCombiners step coda finalize allocators continuation = do - resourceRef <- newEmptyMVar @h - let - step' (Pair h s) a = do - s' <- step h s a - pure (Pair h s') - tryFinalize = do - tryTakeMVar resourceRef >>= \case - Nothing -> pure () - Just resource -> finalize resource - adaptAllocator :: (IO h, h -> IO s) -> IO (Pair h s) - adaptAllocator (allocate, makeInitialState) = do - h <- mask_ do - h <- allocate - putMVar resourceRef h - pure h - s <- makeInitialState h - pure (Pair h s) - coda' :: Pair h s -> IO b - coda' (Pair h s) = do - b <- coda h s - -- this always succeeds, we store the resource at the beginning! - mask_ tryFinalize - pure b - r <- (continuation (combiners step' coda' (adaptAllocator <$> allocators))) - `Control.Exception.finally` - tryFinalize - pure r - --- | Puts the elements of each group into a list that is kept in memory. This breaks streaming within the group. --- --- Useful with 'recast'. -combineIntoLists :: Combiners a [a] -combineIntoLists = combiners - (\s a -> pure (s <> singleton a)) - (pure . closeDList) - (Prelude.repeat (pure mempty)) - --- | Delimits groups in the values yielded by a 'Jet', and can also transform --- those values. -type Splitter a b = MealyIO a (SplitStepResult b) - --- | A [Mealy machine](https://en.wikipedia.org/wiki/Mealy_machine) with an --- existentially hidden state. --- --- Very much like a @FoldM IO@ from the --- [foldl](https://hackage.haskell.org/package/foldl-1.4.12/docs/Control-Foldl.html#t:FoldM) --- library, but it emits an output at each step, not only at the end. -data MealyIO a b where - MealyIO :: (s -> a -> IO (b,s)) -- ^ The step function which threads the state. - -> (s -> IO b) -- ^ The final output, produced from the final state. - -> IO s -- ^ An action that produces the initial state. - -> MealyIO a b - -deriving stock instance Functor (MealyIO a) - --- | For each value coming from upstream, what has the 'Splitter' learned? --- --- * Perhaps we should continue some group we have already started in a previous step. --- --- * Perhaps we have found entire groups that we should emit in one go, groups we know are already complete. --- --- * Perhaps we should start a new group that will continue in the next steps. -data SplitStepResult b = SplitStepResult { - -- | The continued group will be \"closed"\ if in the current step we emit - -- an entire group or we begin a new group. - -- - -- __INVARIANT__: we should only continue a group if we have already - -- opened a \"new one\" with one or more elements in an earlier step. - continuationOfPreviouslyStartedGroup :: [b], - -- | It's ok if the groups we find are empty. - entireGroups :: [[b]], - -- | __INVARIANT__: when we are in the final step, we should not yield elements - -- for the beginning of a new one. - startOfNewGroup :: [b] - } - deriving (Functor, Show) - -instance Semigroup (SplitStepResult b) where - SplitStepResult c1 e1 b1 <> SplitStepResult c2 e2 b2 = - SplitStepResult (c1 <> c2) (e1 <> e2) (b1 <> b2) - -instance Monoid (SplitStepResult b) where - mempty = SplitStepResult [] [] [] - --- TODO: bring back some linear stuff? Perhaps adding a linearFmap ? --- +{-# LANGUAGE BlockArguments #-} +{-# LANGUAGE DerivingStrategies #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE MultiWayIf #-} +{-# LANGUAGE PartialTypeSignatures #-} +{-# LANGUAGE PatternSynonyms #-} +{-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE ViewPatterns #-} +{-# OPTIONS_GHC -Wno-partial-type-signatures #-} + +-- | Tampering with the internals lets you write invalid 'Jet's that don't +-- respect stop signals from consumers, so be careful. +-- +-- Also, the internals expose 'Line' and 'ByteBundle' as thin coats of paint +-- over lazy text and lazy bytestring, respectively. +module Jet.Internal where + +import Control.Applicative +import Control.Concurrent +import Control.Concurrent.Async +import Control.Concurrent.MVar +import Control.Concurrent.STM +import Control.Concurrent.STM.TBMQueue +import Control.Exception +import Control.Monad +import Control.Monad.IO.Class +import Data.Bifunctor +import Data.Bifunctor (first) +import Data.ByteString (ByteString) +import Data.ByteString qualified as B +import Data.ByteString.Lazy qualified as BL +import Data.Foldable qualified +import Data.Function ((&)) +import Data.Functor ((<&>)) +import Data.IORef +import Data.List qualified +import Data.Maybe +import Data.String (IsString (..)) +import Data.Text (Text) +import Data.Text qualified as T +import Data.Text.Encoding qualified as T +import Data.Text.Encoding.Error qualified as T +import Data.Text.IO qualified as T +import Data.Text.Lazy qualified as TL +import Data.Text.Lazy.Encoding qualified as TL +import Data.Traversable qualified +import Data.Typeable +import System.Exit +import System.IO (Handle, IOMode (..), hClose, openBinaryFile) +import System.IO qualified +import System.Process +import Prelude hiding + ( drop, + dropWhile, + filter, + filterM, + fold, + for_, + intersperse, + lines, + take, + takeWhile, + traverse_, + unfold, + unlines, + zip, + zipWith, + ) +import Prelude qualified + +-- import Debug.Trace + +-- $setup +-- +-- >>> :set -XTypeApplications +-- >>> :set -XImportQualifiedPost +-- >>> :set -XScopedTypeVariables +-- >>> :set -XLambdaCase +-- >>> :set -XNumDecimals +-- >>> import Jet (Jet, (&)) +-- >>> import Jet qualified as J +-- >>> import Control.Foldl qualified as L +-- >>> import Control.Concurrent +-- >>> import Data.IORef +-- >>> import Data.Text qualified as T + +-- | A 'Jet' is a sequence of values produced through 'IO' effects. +-- +-- It allows consuming the elements as they are produced and doesn't force them +-- to be present in memory all at the same time, unlike functions like +-- 'Control.Monad.replicateM' from @base@. +newtype Jet a = Jet + { runJet :: forall s. (s -> Bool) -> (s -> a -> IO s) -> s -> IO s + } + +-- | Maps over the yielded elements. '(<&>)' can be used to put the function last. +-- +-- >>> J.each "aa" <&> succ & J.toList +-- "bb" +deriving stock instance Functor Jet + +-- | Go through the elements produced by a 'Jet', while threading an +-- state @s@ and possibly performing some effect. +-- +-- The caller is the one who chooses the type of the state @s@, and must pass +-- an initial value for it. The state is kept in [weak-head normal form](https://en.wikibooks.org/wiki/Haskell/Graph_reduction#Weak_Head_Normal_Form). +-- +-- The caller must also provide a predicate on the state that informs the `Jet` +-- when to stop producing values: whenever the predicate returns +-- @True@. +run :: forall a s. Jet a -> (s -> Bool) -> (s -> a -> IO s) -> s -> IO s +run j = runJet j + +-- | Like 'run', but always goes through all elements produced by the 'Jet'. +-- +-- Equivalent to @run (const False)@. +consume :: forall a s. Jet a -> (s -> a -> IO s) -> s -> IO s +consume j = run j (const False) + +for :: Jet a -> (a -> IO b) -> Jet b +for j k = zipWithIO (\() -> k) (Prelude.repeat (pure ())) j + +for_ :: Jet a -> (a -> IO b) -> IO () +for_ j k = consume j (\() -> void <$> k) () + +-- | Apply an effectful transformation to each element in a 'Jet'. +-- +-- >>> :{ +-- J.each "abc" +-- & J.traverse (\c -> let c' = succ c in putStrLn ([c] ++ " -> " ++ [c']) *> pure c') +-- & J.toList +-- :} +-- a -> b +-- b -> c +-- c -> d +-- "bcd" +traverse :: (a -> IO b) -> Jet a -> Jet b +traverse = flip for + +traverse_ :: (a -> IO b) -> Sink a +traverse_ = flip for_ + +-- | Go through the 'Jet' only for the 'IO' effects, discarding all yielded elements. +drain :: Sink a +drain = traverse_ pure + +-- | Similar to the instance for pure lists, that generates combinations. +-- +-- >>> (,) <$> J.each "ab" <*> J.each "cd" & J.toList +-- [('a','c'),('a','d'),('b','c'),('b','d')] +instance Applicative Jet where + pure i = Jet \stop step initial -> + if + | stop initial -> pure initial + | otherwise -> step initial i + Jet left <*> Jet right = Jet \stop step initial -> + -- Here we assume that the first Jet correctly handles the stop signal. + let step' f s a = step s (f a) + in left stop (\s f -> right stop (step' f) s) initial + +-- | Similar to the instance for pure lists, that does search. +-- +-- >>> :{ +-- do string <- J.each ["ab","cd"] +-- J.each string +-- & +-- J.toList +-- :} +-- "abcd" +instance Monad Jet where + return = pure + Jet m >>= k = Jet \stop step initial -> + m stop (\s a -> runJet (k a) stop step s) initial + +-- | +-- >>> liftIO (putStrLn "foo") <> liftIO (putStrLn "bar") & J.toList +-- foo +-- bar +-- [(),()] +instance MonadIO Jet where + liftIO action = Jet \stop step initial -> + if + | stop initial -> pure initial + | otherwise -> do + a <- action + step initial a + +-- | 'Jet' concatenation. +-- +-- >>> J.each "ab" <> J.each "cd" & J.toList +-- "abcd" +instance Semigroup (Jet a) where + Jet f1 <> Jet f2 = Jet \stop step s0 -> do + -- perhaps some of the stop checks are redundant, the first one in particular? + if + | stop s0 -> + pure s0 + | otherwise -> do + !s1 <- f1 stop step s0 + if + | stop s1 -> + pure s1 + | otherwise -> do + !s2 <- f2 stop step s1 + pure s2 + +-- | 'mempty' is the empty 'Jet'. +-- +-- >>> mempty <> J.each "ab" <> mempty & J.toList +-- "ab" +instance Monoid (Jet a) where + mempty = Jet \_ _ initial -> pure initial + +-- | Same as 'Monoid'. +instance Alternative Jet where + (<|>) = (<>) + empty = mempty + +-- | Same as 'Monoid' +instance MonadPlus Jet where + mzero = mempty + mplus = (<>) + +-- | A failed pattern-match in a do-block produces 'mzero'. +-- +-- >>> :{ +-- do Just c <- J.each [Nothing, Just 'a', Nothing, Just 'b'] +-- pure c +-- & J.toList +-- :} +-- "ab" +instance MonadFail Jet where + fail _ = mzero + +-- | Build a 'Jet' from any 'Foldable' container +-- +-- >>> J.each [True,False] & J.toList +-- [True,False] +each :: forall a f. (Foldable f) => f a -> Jet a +each (Data.Foldable.toList -> seed) = Jet \stop step -> + -- This could be done with Jet.unfold, but let's leave as it is. + let go b s = + if + | stop s -> + pure s + | otherwise -> + case b of + [] -> + pure s + -- see corresponding comment in unfold. + x : xs -> do + !s' <- step s x + go xs s' + in go seed + +-- | +-- +-- >>> J.repeat True & J.take 2 & J.toList +-- [True,True] +repeat :: a -> Jet a +repeat a = repeatIO (pure a) + +-- | +-- +-- >>> J.repeatIO (putStrLn "hi" *> pure True) & J.take 2 & J.toList +-- hi +-- hi +-- [True,True] +repeatIO :: IO a -> Jet a +repeatIO action = untilNothing (fmap Just action) + +-- | +-- +-- >>> J.replicate 2 True & J.toList +-- [True,True] +replicate :: Int -> a -> Jet a +replicate n a = replicateIO n (pure a) + +-- | +-- >>> J.replicateIO 2 (putStrLn "hi" *> pure True) & J.toList +-- hi +-- hi +-- [True,True] +-- +-- Don't confuse this with @Control.Monad.replicateM :: Int -> Jet a -> Jet [a]@ which has a combinatorial behavior. +replicateIO :: Int -> IO a -> Jet a +replicateIO n ioa = take n (repeatIO ioa) + +-- | +-- +-- >>> J.iterate succ (1 :: Int) & J.take 2 & J.toList +-- [1,2] +iterate :: (a -> a) -> a -> Jet a +iterate h = iterateIO (fmap pure h) + +-- | +-- +-- >>> J.iterateIO (\x -> putStrLn "hi" *> pure (succ x)) (1 :: Int) & J.take 2 & J.toList +-- hi +-- [1,2] +iterateIO :: (a -> IO a) -> a -> Jet a +iterateIO h a = pure a <> unfoldIO (fmap (fmap (\x -> Just (x, x))) h) a + +-- | +-- >>> J.unfold (\case [] -> Nothing ; c : cs -> Just (c,cs)) "abc" & J.toList +-- "abc" +unfold :: (b -> Maybe (a, b)) -> b -> Jet a +unfold h = unfoldIO (fmap pure h) + +-- | +-- >>> :{ +-- J.unfoldIO (\x -> do putStrLn "hi" +-- pure $ case x of +-- [] -> Nothing +-- c : cs -> Just (c,cs)) +-- "abc" +-- & J.toList +-- :} +-- hi +-- hi +-- hi +-- hi +-- "abc" +unfoldIO :: (b -> IO (Maybe (a, b))) -> b -> Jet a +unfoldIO h seed = Jet \stop step -> + let go b s = + if + | stop s -> + pure s + | otherwise -> do + next <- h b + case next of + Nothing -> + pure s + -- strictness only on the states. Good idea, or bad? + Just (a, !b') -> do + !s' <- step s a + go b' s' + in go seed + +-- | +-- >>> j = J.untilEOF System.IO.hIsEOF System.IO.hGetLine :: Handle -> Jet String +untilEOF :: (handle -> IO Bool) -> (handle -> IO a) -> handle -> Jet a +untilEOF hIsEOF' hGetLine' handle = untilNothing do + eof <- hIsEOF' handle + if + | eof -> + pure Nothing + | otherwise -> + Just <$> hGetLine' handle + +-- | +-- +-- >>> :{ +-- do ref <- newIORef "abc" +-- let pop = atomicModifyIORef ref (\case [] -> ([], Nothing) +-- x : xs -> (xs, Just x)) +-- J.untilNothing pop & J.toList +-- :} +-- "abc" +untilNothing :: IO (Maybe a) -> Jet a +untilNothing action = unfoldIO (\() -> fmap (fmap (,())) action) () + +-- | Convert to a regular list. This breaks streaming. +-- +-- >>> J.each "abc" & J.toList +-- "abc" +-- +-- Alternatively, we can use 'fold' in combination with 'Control.Foldl.list' form the [foldl](https://hackage.haskell.org/package/foldl) library: +-- +-- >>> L.purely (J.fold (J.each "abc")) L.list +-- "abc" +-- +-- which is more verbose, but more composable. +toList :: Jet a -> IO [a] +toList (Jet f) = do + as <- f (const False) (\xs x -> pure (x : xs)) [] + pure (reverse as) + +-- | Returns the number of elements yielded by the 'Jet', exhausting it in the process. +-- +-- >>> J.each "abc" & J.length +-- 3 +-- +-- Alternatively, we can use 'fold' in combination with 'Control.Foldl.length' form the [foldl](https://hackage.haskell.org/package/foldl) library: +-- +-- >>> L.purely (J.fold (J.each "abc")) L.length +-- 3 +-- +-- which is more verbose, but more composable. +length :: Jet a -> IO Int +length (Jet f) = do + l <- f (const False) (\s _ -> pure (succ s)) 0 + pure l + +data Pair a b = Pair !a !b deriving (Show) + +pairExtract (Pair _ b) = b + +pairEnv (Pair a _) = a + +data Triple a b c = Triple !a !b !c + +tripleExtract (Triple _ _ c) = c + +-- fromTuple :: (a, b) -> Pair a b +-- fromTuple (a, b) -> Pair a b + +-- | >>> J.each "abc" & J.drop 2 & J.toList +-- "c" +drop :: Int -> Jet a -> Jet a +drop limit (Jet f) = Jet \stop step initial -> do + let stop' = stop . pairExtract + step' (Pair count s) a = + if + | count < limit -> do + pure (Pair (succ count) s) + | otherwise -> do + !s' <- step s a + pure (Pair count s') + initial' = Pair 0 initial + Pair _ final <- f stop' step' initial' + pure final + +data DropState = StillDropping | DroppingNoMore + +-- | >>> J.each [1..5] & J.dropWhile (<3) & J.toList +-- [3,4,5] +dropWhile :: (a -> Bool) -> Jet a -> Jet a +dropWhile p = dropWhileIO (fmap pure p) + +dropWhileIO :: (a -> IO Bool) -> Jet a -> Jet a +dropWhileIO p (Jet f) = Jet \stop step initial -> do + let stop' = stop . pairExtract + step' (Pair DroppingNoMore s) a = do + !s' <- step s a + pure (Pair DroppingNoMore s') + step' (Pair StillDropping s) a = do + keepDropping <- p a + if + | keepDropping -> + pure (Pair StillDropping s) + | otherwise -> do + !s' <- step s a + pure (Pair DroppingNoMore s') + initial' = (Pair StillDropping initial) + Pair _ final <- f stop' step' initial' + pure final + +-- | >>> J.each "abc" & J.take 2 & J.toList +-- "ab" +take :: Int -> Jet a -> Jet a +take limit (Jet f) = Jet \stop step initial -> do + let stop' (Pair count s) = + count >= limit || stop s + step' (Pair count s) a = do + !s' <- step s a + pure (Pair (succ count) s') + initial' = Pair 0 initial + Pair _ final <- f stop' step' initial' + pure final + +-- | Synonym for 'take'. +limit :: Int -> Jet a -> Jet a +limit = take + +data TakeState = StillTaking | TakingNoMore + +-- | >>> J.each [1..] & J.takeWhile (<5) & J.toList +-- [1,2,3,4] +takeWhile :: (a -> Bool) -> Jet a -> Jet a +takeWhile p = takeWhileIO (fmap pure p) + +takeWhileIO :: (a -> IO Bool) -> Jet a -> Jet a +takeWhileIO p (Jet f) = Jet \stop step initial -> do + let stop' (Pair TakingNoMore _) = + True + stop' (Pair StillTaking s) = + stop s + step' (Pair internal s) a = do + keepTaking <- p a + if + | keepTaking -> do + !s' <- step s a + pure (Pair internal s') + | otherwise -> + pure (Pair TakingNoMore s) + initial' = Pair StillTaking initial + Pair _ final <- f stop' step' initial' + pure final + +-- | +-- >>> J.each "abc" & J.filter (=='a') & J.toList +-- "a" +filter :: (a -> Bool) -> Jet a -> Jet a +filter p = filterIO (fmap pure p) + +filterIO :: (a -> IO Bool) -> Jet a -> Jet a +filterIO p (Jet f) = Jet \stop step initial -> do + let step' s a = do + shouldPass <- p a + if + | shouldPass -> do + !s' <- step s a + pure s' + | otherwise -> + pure s + f stop step' initial + +-- | Behaves like a combination of 'fmap' and 'foldl'; it applies a function to +-- each element of a structure passing an accumulating parameter from left to right. +-- +-- The resulting 'Jet' has the same number of elements as the original one. +-- +-- Unlike 'Data.Traversable.mapAccumL', it doesn't make the final state available. +-- +-- >>> J.each [1,2,3,4] & J.mapAccum (\a b -> (a + b,a)) 0 & J.toList +-- [0,1,3,6] +mapAccum :: (a -> b -> (a, c)) -> a -> Jet b -> Jet c +mapAccum stepAcc = mapAccumIO (fmap (fmap pure) stepAcc) + +mapAccumIO :: (a -> b -> IO (a, c)) -> a -> Jet b -> Jet c +mapAccumIO stepAcc initialAcc (Jet f) = Jet \stop step initial -> do + let stop' = stop . pairExtract + step' (Pair acc s) b = do + (acc', c) <- stepAcc acc b + !s' <- step s c + pure (Pair acc' s') + initial' = Pair initialAcc initial + Pair _ final <- f stop' step' initial' + pure final + +data Touched + = NotYetTouched + | AlreadyTouched + +-- TODO: there's a bug here!!!! + +-- | +-- >>> J.each "abc" & J.intersperse '-' & J.toList +-- "a-b-c" +intersperse :: a -> Jet a -> Jet a +intersperse intrusion (Jet upstream) = Jet \stop step initial -> do + let stop' = stop . pairExtract + step' (Pair AlreadyTouched s) a = do + !s' <- step s intrusion + if + | stop s' -> + pure (Pair AlreadyTouched s') + | otherwise -> do + !s'' <- step s' a + pure (Pair AlreadyTouched s'') + step' (Pair NotYetTouched s) a = do + !s' <- step s a + pure (Pair AlreadyTouched s') + initial' = Pair NotYetTouched initial + Pair _ final <- upstream stop' step' initial' + pure final + +-- | +-- >>> J.each "abc" & J.zip [1..] & J.toList +-- [(1,'a'),(2,'b'),(3,'c')] +-- +-- >>> J.each [1..] & J.zip "abc" & J.toList +-- [('a',1),('b',2),('c',3)] +zip :: (Foldable f) => f a -> Jet b -> Jet (a, b) +zip = zipWith (,) + +zipWith :: (Foldable f) => (a -> b -> c) -> f a -> Jet b -> Jet c +zipWith zf (Data.Foldable.toList -> as0) = zipWithIO (fmap (fmap pure) zf) (fmap pure as0) + +zipIO :: (Foldable f) => f (IO a) -> Jet b -> Jet (a, b) +zipIO = zipWithIO (\x y -> pure (x, y)) + +-- | +-- Zips a list of 'IO' actions with a 'Jet', where the combining function can also have effects. +-- +-- If the list of actions is exhausted, the 'Jet' stops: +-- +-- >>> J.each [1..] <&> show & zipWithIO (\c1 c2 -> putStrLn (c1 ++ c2)) [pure "a", pure "b"] & J.toList +-- a1 +-- b2 +-- [(),()] +zipWithIO :: (Foldable f) => (a -> b -> IO c) -> f (IO a) -> Jet b -> Jet c +zipWithIO zf (Data.Foldable.toList -> ioas0) (Jet f) = Jet \stop step initial -> do + let stop' (Pair [] _) = True + stop' (Pair _ s) = stop s + step' (Pair (ioa : ioas) s) b = do + a <- ioa + z <- zf a b + !s' <- step s z + pure (Pair ioas s') + step' (Pair [] _) _ = error "never happens" + initial' = Pair ioas0 initial + Pair _ final <- f stop' step' initial' + pure final + +-- | Opens a file and makes the 'Handle' available to all following statements +-- in the do-block. +-- +-- Notice that it's often simpler to use the 'JetSource' (for reading) and +-- 'JetSink' (for writing) instances of 'File'. +withFile :: FilePath -> IOMode -> Jet Handle +withFile path iomode = control @Handle (System.IO.withFile path iomode) + +-- | +-- +-- >>> :{ +-- do r <- J.bracket (putStrLn "allocating" *> pure "foo") (\r -> putStrLn $ "deallocating " ++ r) +-- liftIO $ putStrLn $ "using resource " ++ r +-- & drain +-- :} +-- allocating +-- using resource foo +-- deallocating foo +bracket :: + forall a b. + -- | allocator + IO a -> + -- | finalizer + (a -> IO b) -> + Jet a +bracket allocate free = control @a (Control.Exception.bracket allocate free) + +bracket_ :: + forall a b. + -- | allocator + IO a -> + -- | finalizer + IO b -> + Jet () +bracket_ allocate free = control_ (Control.Exception.bracket_ allocate free) + +bracketOnError :: + forall a b. + -- | allocator + IO a -> + -- | finalizer + (a -> IO b) -> + Jet a +bracketOnError allocate free = control @a (Control.Exception.bracketOnError allocate free) + +-- | +-- +-- Notice how the finalizer runs even when we limit the 'Jet': +-- +-- >>> :{ +-- do J.finally (putStrLn "hi") -- protects statements below +-- liftIO (putStrLn "hey") +-- J.each "abc" +-- & J.limit 2 +-- & J.toList +-- :} +-- hey +-- hi +-- "ab" +-- +-- But if the protected 'Jet' is not consumed at all, the finalizer might not run. +-- +-- >>> :{ +-- do J.finally (putStrLn "hi") -- protects statements below +-- liftIO (putStrLn "hey") +-- J.each "abc" +-- & J.limit 0 +-- & J.toList +-- :} +-- "" +finally :: IO a -> Jet () +finally afterward = + control_ (flip Control.Exception.finally afterward) + +onException :: IO a -> Jet () +onException afterward = + control_ (flip Control.Exception.onException afterward) + +-- | Lift a control operation (like 'Control.Exception.bracket') for which the +-- callback uses the allocated resource. +-- +-- __BEWARE__: the control operation shouldn't do weird things like executing +-- the callback twice. +control :: forall resource. (forall x. (resource -> IO x) -> IO x) -> Jet resource +control f = + Jet \stop step initial -> + if + | stop initial -> + pure initial + | otherwise -> do + f (step initial) + +-- | Lift a control operation (like 'Control.Exception.finally') for which the +-- callback doesn't use the allocated resource. +-- +-- __BEWARE__: the control operation shouldn't do weird things like executing +-- the callback twice. +control_ :: (forall x. IO x -> IO x) -> Jet () +control_ f = + Jet \stop step initial -> + if + | stop initial -> do + pure initial + | otherwise -> do + f (step initial ()) + +-- | +-- +-- >>> L.purely (J.fold (J.each "abc")) ((,) <$> L.list <*> L.length) +-- ("abc",3) +fold :: Jet a -> (s -> a -> s) -> s -> (s -> r) -> IO r +fold (Jet f) step initial coda = do + r <- f (const False) (fmap (fmap pure) step) initial + pure $ coda r + +-- | +-- >>> L.impurely (J.foldIO (J.each "abc")) (L.FoldM (\() c -> putStrLn [c]) (pure ()) pure *> L.generalize L.length) +-- a +-- b +-- c +-- 3 +foldIO :: Jet a -> (s -> a -> IO s) -> IO s -> (s -> IO r) -> IO r +foldIO (Jet f) step initialIO coda = do + initial <- initialIO + r <- f (const False) step initial + coda r + +-- Byte Jets + +-- https://stackoverflow.com/questions/49852060/how-to-choose-chunk-size-when-reading-a-large-file +-- https://askubuntu.com/questions/641900/how-file-system-block-size-works +-- https://stackoverflow.com/questions/1111661/8192-bytes-when-creating-file +data ChunkSize + = DefaultChunkSize + | ChunkSize Int + | ChunkSize1K + | ChunkSize4K + | ChunkSize8K + | ChunkSize16K + | ChunkSize1M + | ChunkSize2M + deriving (Show) + +chunkSize :: ChunkSize -> Int +chunkSize = \case + DefaultChunkSize -> 8192 + ChunkSize c -> c + ChunkSize1K -> 1024 + ChunkSize4K -> 4096 + ChunkSize8K -> 8192 + ChunkSize16K -> 16384 + ChunkSize1M -> 1048576 + ChunkSize2M -> 2097152 + +-- | Helper multi-parameter typeclass for creating 'Jet' values out of a +-- variety of common sources. +-- +-- Because there's no functional dependency, sometimes we need to use +-- @TypeApplications@ to give the compiler a hint about the type of elements +-- we want to produce. For example, here we want 'Line's and not, say, +-- 'ByteString's: +-- +-- >>> action = J.jet @Line (File "foo.txt") & J.sink J.stdout +class JetSource a source where + jet :: source -> Jet a + +bytes :: ChunkSize -> Handle -> Jet ByteString +bytes (chunkSize -> count) handle = + untilEOF System.IO.hIsEOF (flip B.hGetSome count) handle + +instance JetSource ByteString Handle where + jet = bytes DefaultChunkSize + +instance (JetSource a Handle) => JetSource a File where + jet (File path) = do + handle <- withFile path ReadMode + jet handle + +accumByteLengths :: Jet ByteString -> Jet (Int, ByteString) +accumByteLengths = mapAccum (\acc bytes -> let acc' = acc + B.length bytes in (acc', (acc', bytes))) (0 :: Int) + +data AmIContinuing + = Continuing + | NotContinuing + deriving (Show) + +-- | Splits a stream of bytes into groups bounded by maximum byte sizes. When +-- one group \"fills up\", the next one is started. +-- +-- When the list of buckets sizes is exhausted, all incoming bytes are put into +-- the same unbounded group. +-- +-- Useful in combination with 'recast'. +bytesOverBuckets :: [Int] -> Splitter ByteString ByteString +bytesOverBuckets buckets0 = MealyIO step mempty (pure (Pair NotContinuing buckets0)) + where + -- logStep s@(Pair c zzz) a = do + -- putStrLn "foooo!" + -- System.IO.hFlush System.IO.stdout + -- traceIO ("state: " ++ show c) + -- traceIO ("bucket: " ++ show (Prelude.take 2 zzz)) + -- traceIO ("input: " ++ show a) + -- r@(nexts, _) <- step s a + -- traceIO ("output: " ++ show nexts) + -- pure r + step :: Pair AmIContinuing [Int] -> ByteString -> IO (SplitStepResult ByteString, Pair AmIContinuing [Int]) + step splitterState b = do + (continueResult, Pair continuing' buckets', b') <- continue splitterState b + if + | B.null b' -> + pure (continueResult, Pair continuing' buckets') + | otherwise -> do + (entiresResult, splitterState') <- makeEntires mempty b' buckets' + pure (continueResult <> entiresResult, splitterState') + continue :: Pair AmIContinuing [Int] -> ByteString -> IO (SplitStepResult ByteString, Pair AmIContinuing [Int], ByteString) + continue (Pair NotContinuing []) b = pure (nextWith b, Pair NotContinuing [], B.empty) + continue (Pair Continuing []) b = pure (continueWith b, Pair Continuing [], B.empty) + continue (Pair NotContinuing (bucket : buckets)) b = do + let blen = B.length b + -- traceIO ("b = " ++ show b ++ " bucket size= " ++ show bucket) + pure case compare blen bucket of + LT -> (nextWith b, Pair Continuing (bucket - blen : buckets), B.empty) + EQ -> (entireWith (singleton b), Pair NotContinuing buckets, B.empty) + GT -> + let (left, right) = B.splitAt bucket b + in (entireWith (singleton left), Pair NotContinuing buckets, right) + continue (Pair Continuing (bucket : buckets)) b = do + let blen = B.length b + pure case compare blen bucket of + LT -> (continueWith b, Pair Continuing (bucket - blen : buckets), B.empty) + EQ -> (continueWith b, Pair NotContinuing buckets, B.empty) + GT -> + let (left, right) = B.splitAt bucket b + in (continueWith left, Pair NotContinuing buckets, right) + makeEntires :: DList ByteString -> ByteString -> [Int] -> IO (SplitStepResult ByteString, Pair AmIContinuing [Int]) + makeEntires acc b [] = pure (entireWith acc <> nextWith b, Pair Continuing []) + makeEntires acc b (bucket : buckets) = do + let blen = B.length b + case compare blen bucket of + LT -> pure (entireWith acc <> nextWith b, Pair Continuing (bucket - blen : buckets)) + EQ -> pure (entireWith (acc <> singleton b), Pair NotContinuing buckets) + GT -> do + let (left, right) = B.splitAt bucket b + makeEntires (acc <> singleton left) right buckets -- non-terminal + continueWith b = mempty {continuationOfPreviouslyStartedGroup = [b]} + entireWith bdf = mempty {entireGroups = fmap pure (closeDList bdf)} + nextWith b = mempty {startOfNewGroup = [b]} + +-- | A sequence of bytes that we might want to keep together. +newtype ByteBundle = ByteBundle BL.ByteString deriving newtype (Show, Semigroup, Monoid) + +-- | Constructs a 'ByteBundle' out of the bytes of some 'Foldable' container. +bundle :: (Foldable f) => f ByteString -> ByteBundle +bundle = ByteBundle . BL.fromChunks . Data.Foldable.toList + +-- | Length in bytes. +bundleLength :: ByteBundle -> Int +bundleLength (ByteBundle value) = fromIntegral (BL.length value) -- Int64, but unlikely we'll reach the limit + +bundleBytes :: ByteBundle -> Jet ByteString +bundleBytes (ByteBundle value) = each (BL.toChunks value) + +-- | Exception thrown when we try to write too much data in a size-bounded destination. +data BucketOverflow = BucketOverflow + deriving (Show, Typeable) + +instance Exception BucketOverflow + +-- | Splits a stream of 'ByteBundles' into groups bounded by maximum byte +-- sizes. Bytes belonging to the same 'ByteBundle' are always put in the same +-- group. When one group \"fills up\", the next one is started. +-- +-- When the list of buckets sizes is exhausted, all incoming bytes are put into +-- the same unbounded group. +-- +-- Useful in combination with 'recast'. +-- +-- __THROWS__: +-- +-- * 'BucketOverflow' exception if the size bound of a group turns out to be +-- too small for holding even a single 'ByteBundle' value. +byteBundlesOverBuckets :: [Int] -> Splitter ByteBundle ByteString +byteBundlesOverBuckets buckets0 = MealyIO step mempty (pure (Pair NotContinuing buckets0)) + where + step :: Pair AmIContinuing [Int] -> ByteBundle -> IO (SplitStepResult ByteString, Pair AmIContinuing [Int]) + step (Pair splitterState []) (ByteBundle pieces) = + -- We assume [] means "infinite bucket" so once we enter it we'll only be able to continue. + pure + ( case splitterState of + Continuing -> continueWith pieces + NotContinuing -> nextWith pieces, + Pair Continuing [] + ) + step (Pair splitterState (bucket : buckets)) e@(ByteBundle pieces) = do + let elen = bundleLength e + case compare elen bucket of + LT -> + pure + ( case splitterState of + Continuing -> continueWith pieces + NotContinuing -> nextWith pieces, + Pair Continuing (bucket - elen : buckets) + ) + EQ -> + pure + ( case splitterState of + Continuing -> continueWith pieces + NotContinuing -> entireWith pieces, + Pair NotContinuing buckets + ) + -- NB: It's possible to close a bucket and open the next one in the same iteration. + GT -> case splitterState of + Continuing -> step (Pair NotContinuing buckets) e + -- If we are not continuing, that means that the brand-new bucket hasn't + -- enough space to hold a single entity. + NotContinuing -> throwIO BucketOverflow + continueWith bs = mempty {continuationOfPreviouslyStartedGroup = BL.toChunks bs} + entireWith pieces = mempty {entireGroups = [BL.toChunks pieces]} + nextWith bs = mempty {startOfNewGroup = BL.toChunks bs} + +-- | Uses the default system locale. +instance JetSource Line Handle where + jet handle = + textToLine <$> untilEOF System.IO.hIsEOF T.hGetLine handle + +-- +-- +-- Text Jets + +-- | +-- __THROWS__: +-- +-- * 'T.UnicodeException' +decodeUtf8 :: Jet ByteString -> Jet Text +decodeUtf8 (Jet f) = Jet \stop step initial -> do + let stop' = stop . pairExtract + step' (Pair leftovers s) bytes = do + T.Some !text !_ !leftovers' <- pure $ T.streamDecodeUtf8 bytes + !s' <- step s text + pure (Pair leftovers' s') + initial' = Pair leftovers0 initial + Pair leftovers final <- f stop' step' initial' + T.Some !_ !bytes !_ <- pure $ T.streamDecodeUtf8 B.empty + if + | not (B.null bytes) -> + throwIO (T.DecodeError "Unconsumed leftovers at end." Nothing) + | otherwise -> + pure final + where + leftovers0 = + let T.Some _ _ g = T.streamDecodeUtf8 B.empty + in g + +encodeUtf8 :: Jet Text -> Jet ByteString +encodeUtf8 = fmap T.encodeUtf8 + +-- | A line of text. +-- +-- While it is guaranteed that the 'Line's coming out of the 'lines' function +-- do not contain newlines, that invariant is not otherwise enforced. +newtype Line = Line_ TL.Text + deriving newtype (Eq, Ord, Semigroup, Monoid, Show, IsString) + +-- https://ghc.gitlab.haskell.org/ghc/doc/users_guide/exts/pattern_synonyms.html + +-- | Unidirectional pattern that allows converting a 'Line' into a 'Text' +-- during pattern-matching. +pattern Line text <- Line_ (TL.toStrict -> text) + +-- | Converts a 'Line' back to text, without adding the newline. +lineToText :: Line -> Text +lineToText (Line_ text) = TL.toStrict text + +-- | Converts a 'Line' to an utf8-encdoed 'ByteBundle', without adding the newline. +lineToUtf8 :: Line -> ByteBundle +lineToUtf8 (Line_ l) = TL.toChunks l <&> T.encodeUtf8 & bundle + +textToLine :: Text -> Line +textToLine = Line_ . TL.fromStrict + +-- | @Data.Text.singleton '\\n'@ +newline :: Text +newline = T.singleton '\n' + +textToUtf8 :: Text -> ByteBundle +textToUtf8 t = ByteBundle (t & T.encodeUtf8 & BL.fromStrict) + +lineContains :: Text -> Line -> Bool +lineContains t (Line_ l) = TL.isInfixOf (TL.fromStrict t) l + +lineBeginsWith :: Text -> Line -> Bool +lineBeginsWith t (Line_ l) = TL.isPrefixOf (TL.fromStrict t) l + +-- | Adds the 'Text' to the beginning of the 'Line'. +prefixLine :: Text -> Line -> Line +prefixLine t (Line_ l) = Line_ (TL.fromChunks (t : TL.toChunks l)) + +-- textToLine :: Text -> Line +-- textToLine text +-- | Just _ <- T.find (=='\n') text = throw NewlineForbidden +-- | otherwise = Line_ (removeTrailingCarriageReturn text) + +stringToLine :: String -> Line +stringToLine = Line_ . TL.pack + +-- withLineText :: (Text -> r) -> Line -> r +-- withLineText f (Line text) = f text + +isEmptyLine :: Line -> Bool +isEmptyLine (Line_ text) = TL.null text + +emptyLine :: Line +emptyLine = Line_ TL.empty + +-- | Exception thrown when we find newlines in functions which don't accept them. +-- +-- A direct copy of the @NewlineForbidden@ exception from the [turtle](https://hackage.haskell.org/package/turtle) package. +data NewlineForbidden = NewlineForbidden + deriving (Show, Typeable) + +instance Exception NewlineForbidden + +removeTrailingCarriageReturn :: Text -> Text +removeTrailingCarriageReturn text + | T.null text = text + | T.last text == '\r' = T.init text + | otherwise = text + +lines :: Jet Text -> Jet Line +lines (Jet f) = Jet \stop step initial -> do + let stop' = stop . pairExtract + findLinesInCurrentBlock text + | T.null text = + [] + | otherwise = + map (textToLine . removeTrailingCarriageReturn) (T.lines text) + ++ if + | T.last text == '\n' -> + [mempty] + | otherwise -> + [] + step' (Pair lineUnderConstruction s) (findLinesInCurrentBlock -> linesInCurrentBlock) = do + case linesInCurrentBlock of + [] -> do + pure (Pair lineUnderConstruction s) + [l] -> do + pure (Pair (lineUnderConstruction <> singleton l) s) + l : rest@(x : xs) -> do + -- Ineficcient mconcat, better strictify a lazy text here? + let completedLine = mconcat $ runDList lineUnderConstruction [l] + s' <- downstream stop step (completedLine : init rest) s + pure (Pair (singleton (last linesInCurrentBlock)) s') + initial' = Pair mempty initial + Pair (mconcat . closeDList -> lineUnderConstruction) final <- f stop' step' initial' + if + | stop final -> + pure final + | isEmptyLine lineUnderConstruction -> + pure final + | otherwise -> + step final lineUnderConstruction + +unlines :: Jet Line -> Jet Text +unlines j = do + Line text <- j + pure text <> pure (T.singleton '\n') + +downstream :: (s -> Bool) -> (s -> x -> IO s) -> [x] -> s -> IO s +downstream stop step = go + where + go [] s = + pure s + go (x : xs) s + | stop s = + pure s + | otherwise = do + !s' <- step s x + go xs s' + +-- General sinks + +-- | A function that consumes a 'Jet' totally or partially, without returning a result. +type Sink a = Jet a -> IO () + +-- | Helper multi-parameter typeclass for creating 'Jet'-consuming functions +-- out of a variety of common destinations. +-- +-- >>> J.each ["aaa","bbb","ccc"] <&> J.stringToLine & J.sink J.stdout +-- aaa +-- bbb +-- ccc +class JetSink a target where + sink :: target -> Sink a + +instance JetSink ByteString Handle where + sink handle j = for_ j (B.hPut handle) + +instance (JetSink a Handle) => JetSink a File where + sink (File path) j = System.IO.withFile path System.IO.WriteMode \handle -> + sink handle j + +-- | Uses the default system locale. Adds newlines. +instance JetSink Line Handle where + sink handle = traverse_ (T.hPutStrLn handle . lineToText) + +-- | Uses the default system locale. +instance JetSink Text Handle where + sink handle = traverse_ (T.hPutStr handle) + +-- | 'FilePaths' are plain strings. This newtype provides a small measure of +-- safety over them. +newtype File = File {getFilePath :: FilePath} deriving (Show) + +-- | The maximum size in bytes of some destination into which we write the +-- bytes produced by a 'Jet'. +data BoundedSize x = BoundedSize Int x deriving stock (Show, Read) + +instance JetSink ByteBundle Handle where + sink handle j = traverse_ (B.hPut handle) do + s <- j + bundleBytes s + +-- | Distributes incoming bytes through a sequence of files. Once a file is +-- full, we start writing the next one. +instance JetSink ByteString [BoundedSize File] where + sink bucketFiles j = + withCombiners_ + (\handle b -> B.hPut handle b) + hClose + (makeAllocator <$> bucketFiles) + (\combiners -> drain $ recast (bytesOverBuckets bucketSizes) combiners j) + where + bucketSizes = map (\(BoundedSize size _) -> size) bucketFiles + +-- | Distributes incoming bytes through a sequence of files. Once a file is +-- full, we start writing the next one. +-- +-- Each 'ByteBundle' value is garanteed to be written to a single file. If a +-- file turns out to be too small for even a single 'ByteBundle' value, a +-- 'BucketOverflow' exception is thrown. +instance JetSink ByteBundle [BoundedSize File] where + sink bucketFiles j = + withCombiners_ + (\handle b -> B.hPut handle b) + hClose + (makeAllocator <$> bucketFiles) + (\combiners -> drain $ recast (byteBundlesOverBuckets bucketSizes) combiners j) + where + bucketSizes = map (\(BoundedSize size _) -> size) bucketFiles + +makeAllocator :: BoundedSize File -> IO Handle +makeAllocator (BoundedSize _ (File path)) = openBinaryFile path WriteMode + +-- DList helper +newtype DList a = DList {runDList :: [a] -> [a]} + +instance Semigroup (DList a) where + DList a1 <> DList a2 = DList (a1 . a2) + +instance Monoid (DList a) where + mempty = DList id + +makeDList :: [a] -> DList a +makeDList as = DList \xs -> as ++ xs + +closeDList :: DList a -> [a] +closeDList (DList f) = f [] + +singleton :: a -> DList a +singleton a = DList $ (a :) + +-- +-- concurrency + +-- | Process the values yielded by the upstream 'Jet' in a concurrent way, +-- and return the results in the form of another 'Jet' as they are produced. +-- +-- __NB__: this function might scramble the order of the returned values. Right +-- now there isn't a function for unscrambling them. +-- +-- >>> :{ +-- J.each [(3,'a'), (2,'b'), (1,'c')] +-- & J.traverseConcurrently (numberOfWorkers 10) (\(d,c) -> threadDelay (d*1e5) *> pure c) +-- & J.toList +-- :} +-- "cba" +-- +-- What happens if we 'limit' the resulting 'Jet' and we reach that limit, or +-- if we otherwise stop consuming the 'Jet' before it gets exhausted? In those +-- cases, all pending @IO b@ tasks are cancelled. +-- +-- >>> :{ +-- J.each [(9999,'a'), (2,'b'), (1,'c')] +-- & J.traverseConcurrently (numberOfWorkers 10) (\(d,c) -> threadDelay (d*1e5) *> pure c) +-- & J.take 2 +-- & J.toList +-- :} +-- "cb" +traverseConcurrently :: (PoolConf -> PoolConf) -> (a -> IO b) -> Jet a -> Jet b +-- TODO: +-- It would be nice to have 0-lengh channels for which one side blocks until +-- the other side takes the job. +traverseConcurrently adaptConf makeTask upstream = Jet \stop step initial -> do + if + -- If we know we aren't going to do any work, don't bother starting the + -- whole boondoggle. + | stop initial -> + pure initial + | otherwise -> do + -- At this point we know we should do at least one step. + let PoolConf {_inputQueueSize, _numberOfWorkers, _outputQueueSize} = adaptConf defaultPoolConf + input <- newTBMQueueIO _inputQueueSize + inputQueueWriterShouldStop <- newIORef False + aliveWorkers <- newIORef _numberOfWorkers + output <- newTBMQueueIO _outputQueueSize + let -- The inputQueueWriter should *not* be interrupted aynchronously. + -- After each iteration, it reads the IORef to see if it should stop. + -- Once it stops, it closes the input queue. + inputQueueWriter = do + run + upstream + id + ( \_ a -> do + atomically $ writeTBMQueue input (makeTask a) + readIORef inputQueueWriterShouldStop + ) + False + atomically $ closeTBMQueue input + -- Workers *can* be interrupted asynchronously. + worker = do + mtask <- atomically $ readTBMQueue input + case mtask of + Nothing -> do + remaining <- do + atomicModifyIORef' aliveWorkers \count -> + let count' = pred count + in (count', count') + if + | remaining == 0 -> do + atomically $ closeTBMQueue output + | otherwise -> do + pure () + Just task -> do + result <- task + atomically $ writeTBMQueue output result + worker + outputQueueReader s = do + if + | stop s -> do + -- tell the inserter from upstream that it should stop. is this enough? + writeIORef inputQueueWriterShouldStop True + atomically $ closeTBMQueue input -- perhaps unnecessary? + pure s + | otherwise -> do + mresult <- atomically $ readTBMQueue output + case mresult of + Nothing -> do + pure s + Just result -> do + !s' <- step s result + outputQueueReader s' + runConcurrently $ + Concurrently do + inputQueueWriter + *> Concurrently do + finalLeft <- do + runConcurrentlyE $ + -- The worker pool is always killed when the output reader finishes, + -- but for the "happy path" the workers will already be dead. + ConcurrentlyE (Right <$> replicateConcurrently_ _numberOfWorkers worker) + *> + -- This Left is what kills the worker pool. + ConcurrentlyE (Left <$> outputQueueReader initial) + case finalLeft of + Right () -> do + error "never happens, the Left always wins" + Left final -> do + pure final + +-- | Configuration record for the worker pool. +data PoolConf = PoolConf + { _inputQueueSize :: Int, + _numberOfWorkers :: Int, + _outputQueueSize :: Int + } + deriving (Show) + +defaultPoolConf = + PoolConf + { _inputQueueSize = 1, + _numberOfWorkers = 1, + _outputQueueSize = 1 + } + +-- | Size of the waiting queue into the worker pool. The default is @1@. +inputQueueSize :: Int -> PoolConf -> PoolConf +inputQueueSize size poolConf = poolConf {_inputQueueSize = size} + +-- | The size of the worker pool. The default is @1@. +numberOfWorkers :: Int -> PoolConf -> PoolConf +numberOfWorkers number poolConf = poolConf {_numberOfWorkers = number} + +-- | Size of the queue holding results out of the working pool before they +-- are yielded downstream. The default is @1@. +outputQueueSize :: Int -> PoolConf -> PoolConf +outputQueueSize size poolConf = poolConf {_outputQueueSize = size} + +-- | An alias for 'id'. Useful with functions like 'traverseConcurrently' and +-- 'throughProcess', for which it means \"use the default configuration\". +defaults :: a -> a +defaults = id + +-- +-- process invocation + +-- | Feeds the upstream 'Jet' to an external process' @stdin@ and returns the +-- process' @stdout@ as another @Jet@. The feeding and reading of the standard +-- streams is done concurrently in order to avoid deadlocks. +-- +-- What happens if we 'limit' the resulting 'Jet' and we reach that limit, or +-- if we otherwise stop consuming the 'Jet' before it gets exhausted? In those +-- cases, the external process is promptly terminated. +throughProcess :: (ProcConf -> ProcConf) -> CreateProcess -> Jet ByteString -> Jet ByteString +throughProcess adaptConf = throughProcess_ (adaptConf defaultProcConf) + +-- | Like 'throughProcess', but feeding and reading 'Line's using the default +-- system encoding. +-- +-- >>> :{ +-- J.each ["aaa","bbb","ccc"] +-- <&> J.stringToLine +-- & linesThroughProcess defaults (shell "cat") +-- & J.toList +-- :} +-- ["aaa","bbb","ccc"] +-- +-- An example of not reading all the lines from a long-lived process that gets cancelled: +-- +-- >>> :{ +-- mempty +-- & linesThroughProcess defaults (shell "{ printf \"aaa\\nbbb\\nccc\\n\" ; sleep infinity ; }") +-- & J.limit 2 +-- & J.toList +-- :} +-- ["aaa","bbb"] +linesThroughProcess :: (ProcConf -> ProcConf) -> CreateProcess -> Jet Line -> Jet Line +linesThroughProcess adaptConf procSpec = do + let textLinesProcConf = + (adaptConf defaultProcConf) + { _writeToStdIn = T.hPutStrLn, + _readFromStdout = T.hGetLine + } + fmap textToLine . throughProcess_ textLinesProcConf procSpec . fmap lineToText + +-- | Like 'throughProcess', but feeding and reading 'Line's encoded in UTF8. +utf8LinesThroughProcess :: (ProcConf -> ProcConf) -> CreateProcess -> Jet Line -> Jet Line +utf8LinesThroughProcess adaptConf procSpec = do + lines . decodeUtf8 . throughProcess adaptConf procSpec . encodeUtf8 . unlines + +throughProcess_ :: forall a b. ProcConf_ a b -> CreateProcess -> Jet a -> Jet b +throughProcess_ procConf procSpec upstream = Jet \stop step initial -> do + let ProcConf_ {_bufferStdin, _writeToStdIn, _readFromStdout, _readFromStderr, _handleExitCode} = procConf + if + -- If we know we aren't going to do any work, don't bother starting the + -- whole boondoggle. + | stop initial -> + pure initial + | otherwise -> do + let procSpec' = + procSpec + { std_in = CreatePipe, + std_out = CreatePipe, + std_err = CreatePipe + } + input <- newTBMQueueIO @a 1 + inputQueueWriterShouldStop <- newIORef False + -- remember to drain stderr concurrently with stdout... + let inputQueueWriter = do + run + upstream + id + ( \_ a -> do + atomically $ writeTBMQueue input a + readIORef inputQueueWriterShouldStop + ) + False + atomically $ closeTBMQueue input + finalEither <- + runConcurrently $ + Concurrently do + inputQueueWriter + *> Concurrently do + withCreateProcess procSpec' \(Just stdin') (Just stdout') (Just stderr') phandle -> do + when (not _bufferStdin) (System.IO.hSetBuffering stdin' System.IO.NoBuffering) + let stdinWriter = do + ma <- atomically $ readTBMQueue input + case ma of + Nothing -> do + hClose stdin' + Just a -> do + _writeToStdIn stdin' a + stdinWriter + stderrReader = do + untilEOF System.IO.hIsEOF _readFromStdout stderr' & drain + stdoutReader s = do + if + | stop s -> do + writeIORef inputQueueWriterShouldStop True + pure (Left s) + | otherwise -> do + eof <- System.IO.hIsEOF stdout' + if + | eof -> do + writeIORef inputQueueWriterShouldStop True + exitCode <- waitForProcess phandle + _handleExitCode exitCode + pure (Right s) + | otherwise -> do + b <- _readFromStdout stdout' + !s' <- step s b + stdoutReader s' + runConcurrentlyE $ + ConcurrentlyE do Right <$> stdinWriter + *> ConcurrentlyE do Right <$> stderrReader + *> ConcurrentlyE do stdoutReader initial + pure (either id id finalEither) + +-- | Configuration record with some extra options in addition to those in "CreateProcess". +type ProcConf = ProcConf_ ByteString ByteString + +data ProcConf_ a b = ProcConf_ + { _bufferStdin :: Bool, + _writeToStdIn :: Handle -> a -> IO (), + _readFromStdout :: Handle -> IO b, + _readFromStderr :: Handle -> IO (), + _handleExitCode :: ExitCode -> IO () + } + +defaultProcConf :: ProcConf +defaultProcConf = + ProcConf_ + { _bufferStdin = False, + _writeToStdIn = B.hPut, + _readFromStdout = flip B.hGetSome 8192, + _readFromStderr = void . T.hGetLine, + _handleExitCode = \exitCode -> case exitCode of + ExitFailure _ -> throwIO exitCode + ExitSuccess -> pure () + } + +-- | Should we buffer the process' @stdin@? Usually should be 'True' for +-- interactive scenarios. +-- +-- By default, 'False'. +bufferStdin :: Bool -> ProcConf -> ProcConf +bufferStdin doBuffering procConf = procConf {_bufferStdin = doBuffering} + +-- | Sets the function that reads a single line of output from the process +-- @stderr@. It's called repeatedly until @stderr@ is exhausted. The reads are +-- done concurrently with the reads from @stdout@. +-- +-- By default, lines of text are read using the system's default encoding. +-- +-- This is a good place to throw an exception if we don't like what comes out +-- of @stderr@. +readFromStderr :: (Handle -> IO ()) -> ProcConf -> ProcConf +readFromStderr readFunc procConf = procConf {_readFromStderr = readFunc} + +-- | Sets the function that handles the final `ExitCode` of the process. +-- +-- The default behavior is to throw the `ExitCode` as an exception if it's not +-- a success. +handleExitCode :: (ExitCode -> IO ()) -> ProcConf -> ProcConf +handleExitCode handler procConf = procConf {_handleExitCode = handler} + +-- +-- +-- complicated stufff + +data AreWeInsideGroup foldState + = OutsideGroup + | InsideGroup !foldState + +data RecastState foldState = RecastState !(AreWeInsideGroup foldState) [IO foldState] + +-- | This is a complex, unwieldly, yet versatile function. It can be used to +-- define grouping operations, but also for decoding and other purposes. +-- +-- Groups are delimited in the input 'Jet' using the 'Splitter', and the +-- contents of those groups are then combined using 'Combiners'. The result of +-- each combiner is yielded by the return 'Jet'. +-- +-- If the list of combiners is finite and becomes exhausted, we stop splitting +-- and the return 'Jet' stops. +recast :: forall a b c. Splitter a b -> Combiners b c -> Jet a -> Jet c +recast + (MealyIO splitterStep splitterCoda splitterAlloc) + (Combiners foldStep foldCoda foldAllocs0) + (Jet upstream) = Jet \stop step initial -> do + initialSplitterState <- splitterAlloc + let -- When to stop? Either downstream says we need to stop, + -- or we are outside a group and there isn't another group consumer we + -- can use to process the next one. + stop' (Triple _ (RecastState OutsideGroup []) _) = True + stop' (Triple _ _ s) = stop s + + step' (Triple splitterState recastState s) a = do + (splitResult, splitterState') <- splitterStep splitterState a + Pair recastState' s' <- advanceRecast splitResult recastState s + pure (Triple splitterState' recastState' s') + + advanceRecast ssr@(SplitStepResult {continuationOfPreviouslyStartedGroup, entireGroups, startOfNewGroup}) (RecastState areWeInside foldAllocs) s = do + case (areWeInside, entireGroups, startOfNewGroup) of + -- If there aren't any new groups and we don't start an incomplete one, just advance the current fold + (InsideGroup foldState, [], []) -> do + -- traceIO $ "recast inside group just continuing" + foldState' <- advanceGroupWithougClosing foldState continuationOfPreviouslyStartedGroup + pure (Pair (RecastState (InsideGroup foldState') foldAllocs) s) -- main state didn't change + (InsideGroup foldState, _, _) -> do + -- traceIO $ "recast inside group closing" + !c <- processSingleGroup foldState continuationOfPreviouslyStartedGroup + !s' <- step s c + if + | stop s' -> do + -- traceIO $ "recast inside group pure" + pure (Pair (RecastState OutsideGroup foldAllocs) s') + | otherwise -> do + -- traceIO $ "recast inside group advancing" + advanceRecast ssr (RecastState OutsideGroup foldAllocs) s' + -- if we are outside of a group, the "continuationOfPreviouslyStartedGroup" is ignored. + (OutsideGroup, _, _) -> do + -- traceIO $ "recast outside group" + -- doens't return foldState becasue we close the groups + Pair foldAllocs' s' <- processEntireGroups foldAllocs s entireGroups + bail <- pure (Pair (RecastState OutsideGroup foldAllocs') s') + if + | stop s' -> do + pure bail + | otherwise -> do + case startOfNewGroup of + [] -> do + pure bail + (_ : _) -> do + case foldAllocs of + [] -> do + pure bail + alloc : allocs -> do + -- traceIO $ "recast we should be allocating here" + -- there is a next group, so let's begin it + !foldState0 <- alloc + foldState <- processBeginNextGroup foldState0 startOfNewGroup + pure (Pair (RecastState (InsideGroup foldState) allocs) s') + -- foldM ? + advanceGroupWithougClosing :: _ -> [b] -> IO _ + advanceGroupWithougClosing foldState [] = + pure foldState + advanceGroupWithougClosing foldState (b : bs) = do + !foldState' <- foldStep foldState b + advanceGroupWithougClosing foldState' bs + processEntireGroups :: [IO _] -> _ -> [[b]] -> IO (Pair [IO _] _) + -- We can't go on if there aren't any more groups + processEntireGroups allocs s [] = do + pure (Pair allocs s) + -- We can't go on if there aren't any more fold initial state allocs + processEntireGroups [] s _ = do + pure (Pair [] s) + processEntireGroups (alloc : allocs) s (bs : bss) = do + !foldState0 <- alloc + !c <- processSingleGroup foldState0 bs -- a single step downstream + !s' <- step s c + if + | stop s' -> do + pure (Pair allocs s') + | otherwise -> do + processEntireGroups allocs s' bss + -- a whole fold is processed here + processSingleGroup :: _ -> [b] -> IO c + processSingleGroup foldState [] = do + foldCoda foldState + processSingleGroup foldState (b : bs) = do + !foldState' <- foldStep foldState b + processSingleGroup foldState' bs + processBeginNextGroup :: _ -> [b] -> IO _ + processBeginNextGroup foldState [] = do + pure foldState + processBeginNextGroup foldState (b : bs) = do + !foldState' <- foldStep foldState b + processBeginNextGroup foldState' bs + initial' = Triple initialSplitterState (RecastState OutsideGroup foldAllocs0) initial + Triple splitterState recastState final <- upstream stop' step' initial' + -- What happens if there's a fold ongoing when we stop? Right now we always close it, which seems to be a reasonable + -- action (because the fold coda might hide a finalizer). + -- + -- Also, when can it happen that we reach this point with an ongoing fold? + -- If I understand correctly: + -- - it can only happen when the upstream closes and leaves the fold open. + -- - it can't (?) happen when the consumer stops early. + let closePendingFold = \case + RecastState OutsideGroup _ -> do + pure () + RecastState (InsideGroup foldState) _ -> do + _ <- foldCoda foldState + pure () + if + | stop final -> do + closePendingFold recastState + pure final + | otherwise -> do + splitResult <- splitterCoda splitterState + -- We discard the "begins next group"; it doesn't make sense in this final step. + Pair recastState' final' <- advanceRecast (splitResult {startOfNewGroup = []}) recastState final + if + | stop final' -> do + -- TODO: + -- should we dealloc here? Maybe there's a fold reaminging... we should close it. See below. + closePendingFold recastState' + pure final' + | otherwise -> do + case recastState' of + RecastState OutsideGroup _ -> do + -- traceIO $ "final! outside group" + pure final' + RecastState (InsideGroup foldState) _ -> do + -- traceIO $ "final! inside group" + c <- foldCoda foldState + final'' <- step final' c + pure final'' + +-- | A 'Combiners' value knows how to process a sequence of groups, while +-- keeping a (existentially hidden) state for each group. +-- +-- Very much like a @FoldM IO@ from the +-- [foldl](https://hackage.haskell.org/package/foldl-1.4.12/docs/Control-Foldl.html#t:FoldM) +-- library, but \"restartable\" with a list of starting states. +-- +-- For converting one into the other, this function should do the trick: +-- +-- > \(L.FoldM step allocator coda) -> combiners step coda (Prelude.repeat allocator) +data Combiners a b where + Combiners :: (s -> a -> IO s) -> (s -> IO b) -> [IO s] -> Combiners a b + +deriving stock instance Functor (Combiners a) + +-- | Constructor for 'Combiners' values. +combiners :: + forall s a b r. + -- \^ foo + + -- | Step function that threads the state @s@. + (s -> a -> IO s) -> + -- | Coda invoked when a group closes. + (s -> IO b) -> + -- | Actions that produce the initial states @s@ for processing each group. + [IO s] -> + Combiners a b +combiners = Combiners + +-- | A simpler version of 'withCombiners' that doen't thread a state; it merely +-- allocates and deallocates the resource @h@. +withCombiners_ :: + forall h a r. + -- | Step function that accesses the resource @h@. + (h -> a -> IO ()) -> + -- | Finalizer to run after closing each group, and also in the case of an exception. + (h -> IO ()) -> + -- | Actions that allocate a sequence of resources @h@. + [IO h] -> + -- | The 'Combiners' value should be consumed linearly. + (Combiners a () -> IO r) -> + IO r +withCombiners_ step finalize allocators = do + withCombiners + (\h () a -> step h a) + (\_ () -> pure ()) + finalize + ( do + allocator <- allocators + pure (allocator, \_ -> pure ()) + ) + +-- | 'Combiners' thread a state @s@ while processing each group. Sometimes, in +-- addition to that, we want to allocate a resource @h@ when we start +-- processing a group, and deallocate it after we finish processing the group +-- or an exception is thrown. The typical example is allocating a 'Handle' for +-- writing the elements of the group as they arrive. +withCombiners :: + forall h s a b r. + -- | Step function that accesses the resource @h@ and threads the state @s@. + (h -> s -> a -> IO s) -> + -- | Coda invoked when a group closes. + (h -> s -> IO b) -> + -- | Finalizer to run after each coda, and also in the case of an exception. + (h -> IO ()) -> + -- | Actions that allocate a sequence of resources @h@ and produce initial states @s@ for processing each group. + [(IO h, h -> IO s)] -> + -- | The 'Combiners' value should be consumed linearly. + (Combiners a b -> IO r) -> + IO r +withCombiners step coda finalize allocators continuation = do + resourceRef <- newEmptyMVar @h + let step' (Pair h s) a = do + s' <- step h s a + pure (Pair h s') + tryFinalize = do + tryTakeMVar resourceRef >>= \case + Nothing -> pure () + Just resource -> finalize resource + adaptAllocator :: (IO h, h -> IO s) -> IO (Pair h s) + adaptAllocator (allocate, makeInitialState) = do + h <- mask_ do + h <- allocate + putMVar resourceRef h + pure h + s <- makeInitialState h + pure (Pair h s) + coda' :: Pair h s -> IO b + coda' (Pair h s) = do + b <- coda h s + -- this always succeeds, we store the resource at the beginning! + mask_ tryFinalize + pure b + r <- + (continuation (combiners step' coda' (adaptAllocator <$> allocators))) + `Control.Exception.finally` tryFinalize + pure r + +-- | Puts the elements of each group into a list that is kept in memory. This breaks streaming within the group. +-- +-- Useful with 'recast'. +combineIntoLists :: Combiners a [a] +combineIntoLists = + combiners + (\s a -> pure (s <> singleton a)) + (pure . closeDList) + (Prelude.repeat (pure mempty)) + +-- | Delimits groups in the values yielded by a 'Jet', and can also transform +-- those values. +type Splitter a b = MealyIO a (SplitStepResult b) + +-- | A [Mealy machine](https://en.wikipedia.org/wiki/Mealy_machine) with an +-- existentially hidden state. +-- +-- Very much like a @FoldM IO@ from the +-- [foldl](https://hackage.haskell.org/package/foldl-1.4.12/docs/Control-Foldl.html#t:FoldM) +-- library, but it emits an output at each step, not only at the end. +data MealyIO a b where + MealyIO :: + -- | The step function which threads the state. + (s -> a -> IO (b, s)) -> + -- | The final output, produced from the final state. + (s -> IO b) -> + -- | An action that produces the initial state. + IO s -> + MealyIO a b + +deriving stock instance Functor (MealyIO a) + +-- | For each value coming from upstream, what has the 'Splitter' learned? +-- +-- * Perhaps we should continue some group we have already started in a previous step. +-- +-- * Perhaps we have found entire groups that we should emit in one go, groups we know are already complete. +-- +-- * Perhaps we should start a new group that will continue in the next steps. +data SplitStepResult b = SplitStepResult + { -- | The continued group will be \"closed"\ if in the current step we emit + -- an entire group or we begin a new group. + -- + -- __INVARIANT__: we should only continue a group if we have already + -- opened a \"new one\" with one or more elements in an earlier step. + continuationOfPreviouslyStartedGroup :: [b], + -- | It's ok if the groups we find are empty. + entireGroups :: [[b]], + -- | __INVARIANT__: when we are in the final step, we should not yield elements + -- for the beginning of a new one. + startOfNewGroup :: [b] + } + deriving (Functor, Show) + +instance Semigroup (SplitStepResult b) where + SplitStepResult c1 e1 b1 <> SplitStepResult c2 e2 b2 = + SplitStepResult (c1 <> c2) (e1 <> e2) (b1 <> b2) + +instance Monoid (SplitStepResult b) where + mempty = SplitStepResult [] [] [] + +-- TODO: bring back some linear stuff? Perhaps adding a linearFmap ? +-- diff --git a/test/doctests.hs b/test/doctests.hs deleted file mode 100644 --- a/test/doctests.hs +++ /dev/null @@ -1,4 +0,0 @@ -module Main (main) where - -import Test.DocTest -main = doctest ["-ilib", "lib/Jet.hs", "lib/Jet/Internal.hs"] diff --git a/test/tests.hs b/test/tests.hs --- a/test/tests.hs +++ b/test/tests.hs @@ -1,238 +1,238 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE BlockArguments #-} -{-# LANGUAGE DeriveFunctor #-} -{-# LANGUAGE DerivingStrategies #-} -{-# LANGUAGE DerivingVia #-} -{-# LANGUAGE GeneralizedNewtypeDeriving #-} -{-# LANGUAGE ImportQualifiedPost #-} -{-# LANGUAGE MultiWayIf #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE ViewPatterns #-} -{-# LANGUAGE TupleSections #-} -{-# LANGUAGE LinearTypes #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE LambdaCase #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE PatternSynonyms #-} -{-# LANGUAGE UndecidableInstances #-} -{-# LANGUAGE NamedFieldPuns #-} -{-# LANGUAGE GADTSyntax #-} -{-# LANGUAGE ExistentialQuantification #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE StandaloneKindSignatures #-} -{-# LANGUAGE PartialTypeSignatures #-} -{-# LANGUAGE StandaloneDeriving #-} -{-# LANGUAGE ApplicativeDo #-} -{-# LANGUAGE NumDecimals #-} -{-# OPTIONS_GHC -Wno-partial-type-signatures #-} -module Main where - -import Test.Tasty -import Test.Tasty.HUnit -import Data.IORef -import Data.Text (Text) -import Data.Text qualified as T -import Data.Text.IO qualified as T -import Data.Text.Encoding qualified as T -import Data.Text.Encoding.Error qualified as T -import Data.Text.Lazy qualified as TL -import Data.Text.Lazy.Encoding qualified as TL -import Data.ByteString (ByteString) -import Data.ByteString qualified as B -import Data.ByteString.Lazy qualified as BL -import Data.Time.Clock -import Data.Foldable -import Debug.Trace -import Data.Functor.Identity -import Control.Concurrent -import Data.List -import Control.Exception - -import Jet -import Jet qualified as J - -tests :: TestTree -tests = - testGroup - "All" - [ - testGroup "byteSplitter" $ - let tests = do - -- splitSize <- [1] - -- bucketSize <- [2] - splitSize <- [1..7] - bucketSize <- [1..10] - pure $ - testCase ("splitter splitSize=" ++ show splitSize ++ " bucketSize=" ++ show bucketSize) $ - assertBytesCorrectlySplit bucketSize (bytePieces splitSize az) - in tests - , - testGroup "byteBundleSplitter" $ - let tests = do - splitSize <- [1..7] - bucketSize <- [splitSize..13] - pure $ - testCase ("splitter splitSize=" ++ show splitSize ++ " bucketSize=" ++ show bucketSize) $ - assertByteBundlesCorrectlySplit bucketSize (bytePieces splitSize az) - in tests - , - testGroup "lines" $ - let tests = do - fragmentSize <- [1..13] - pure $ - testCase ("fragment size " ++ show fragmentSize) $ - assertLines fragmentSize lineData01 lineExpected01 - in tests - , - testGroup "process" $ - [ - testCase "simple" $ - do resultLines <- - mempty - & linesThroughProcess defaults (shell "echo foo") - & J.toList - assertEqual "input and output lines don't match" (textToLine . T.pack <$> ["foo"]) resultLines - , - testCase "simple 2" $ - do let inputLines = textToLine . T.pack <$> ["aaa","bbb","ccc"] - resultLines <- - J.each inputLines - & linesThroughProcess defaults (shell "cat") - & J.toList - assertEqual "input and output lines don't match" inputLines resultLines - , - testCase "interruption" $ - do let expectedLines = textToLine . T.pack <$> ["aaa","bbb"] - resultLines <- - mempty - & linesThroughProcess defaults (shell "{ printf \"aaa\\nbbb\\nccc\\n\" ; sleep infinity ; }") - & J.limit 2 - & J.toList - assertEqual "unexpected lines at output" expectedLines resultLines - ] - , - testGroup "concurrency" $ - [ - testCase "compare" $ do - let yieldAfter d x = sleep d *> pure x - delay = cents 100 - upstream = J.each "abcde" - (ts, rsequential) <- upstream & J.traverse (yieldAfter delay) & J.toList & time - (t1, rconc1) <- upstream & J.traverseConcurrently (numberOfWorkers 2) (yieldAfter delay) & J.toList & time - (t2, rconc2) <- upstream & J.traverseConcurrently (numberOfWorkers 10) (yieldAfter delay) & J.toList & time - let (rsequential', rconc1', rconc2') = (sort rsequential, sort rconc1, sort rconc2) - assertEqual "sequential != conc" rsequential' rconc1' - assertEqual "conc != conc 2" rconc1' rconc2' - assertBool "conc not faster" (t1 < ts) - assertBool "conc2 not faster" (t2 < t1) - pure () - , testCase "cancelation" $ do - latch <- newEmptyMVar - l <- - J.each "ab" - & traverseConcurrently (numberOfWorkers 2) - (\c -> do let delay 'a' = threadDelay 9999e6 - delay _ = pure () - (delay c *> pure c) `Control.Exception.onException` putMVar latch ()) - & J.limit 1 - & J.length - _ <- takeMVar latch - assertEqual "only 1 element" 1 l - pure () - ] - ] - -az :: ByteString -az = ['a'..'z'] & T.pack & T.encodeUtf8 - -bytePieces :: Int -> ByteString -> [ByteString] -bytePieces size = - let go b = - if B.null b - then [] - else let (left,right) = B.splitAt size b - in left : go right - in go - -assertBytesCorrectlySplit :: Int -> [ByteString] -> IO () -assertBytesCorrectlySplit bucketSize inputs = do - let buckets = Prelude.repeat bucketSize - j = J.recast (J.bytesOverBuckets buckets) combineIntoLists (J.each inputs) - fragmentedGroups <- J.toList j - let groups :: [ByteString] = mconcat <$> fragmentedGroups - concatenatedInput = T.decodeUtf8 $ mconcat inputs - concatenatedOutput = T.decodeUtf8 $ mconcat groups - assertEqual "combined inputs and result" concatenatedInput concatenatedOutput - -- traceIO "--------------------------" - -- traceIO $ "+ original groups = " ++ show fragmentedGroups - -- traceIO $ "+ collected groups = " ++ show groups - -- traceIO $ "* bucket size = " ++ show bucketSize - -- traceIO $ show $ B.length <$> Prelude.init groups - -- traceIO "--------------------------" - assertBool "group sizes are wrong" $ all (\g -> B.length g == bucketSize) (Prelude.init groups) - pure () - -assertByteBundlesCorrectlySplit :: Int -> [ByteString] -> IO () -assertByteBundlesCorrectlySplit bucketSize inputs = do - let buckets = Prelude.repeat bucketSize - j = J.recast (J.byteBundlesOverBuckets buckets) combineIntoLists (bundle . Identity <$> J.each inputs) - fragmentedGroups <- J.toList j - let groups :: [ByteString] = mconcat <$> fragmentedGroups - concatenatedInput = T.decodeUtf8 $ mconcat inputs - concatenatedOutput = T.decodeUtf8 $ mconcat groups - concatenatedOutput' = T.decodeUtf8 $ mconcat $ Data.List.intersperse (T.encodeUtf8 (T.singleton '-')) groups - assertEqual "combined inputs and result" concatenatedInput concatenatedOutput - -- traceIO "--------------------------" - -- traceIO $ "+ inputs = " ++ show inputs - -- traceIO $ "+ original groups = " ++ show fragmentedGroups - -- traceIO $ "+ collected groups = " ++ show groups - -- traceIO $ "* bucket size = " ++ show bucketSize - -- traceIO $ show $ B.length <$> Prelude.init groups - -- traceIO "--------------------------" - assertBool "group sizes are wrong" $ all (\g -> B.length g <= bucketSize) (Prelude.init groups) - Data.Foldable.for_ inputs \i -> - assertBool "bundle divided" $ T.isInfixOf (T.decodeUtf8 i) concatenatedOutput' - pure () - - -lineData01 :: Text -lineData01 = T.pack "aaa\nbb\nccc\ndddd\n\neee\n\n\nfffffffff\ng\niiiii" - -lineExpected01 :: [Line] -lineExpected01 = textToLine . T.pack <$> ["aaa","bb","ccc","dddd","","eee","","","fffffffff", "g", "iiiii"] - -textPieces :: Int -> Text -> [Text] -textPieces size = - let go t = - if T.null t - then [] - else let (left,right) = T.splitAt size t - in left : go right - in go - -assertLines :: Int -> Text -> [Line] -> IO () -assertLines textFragmentSize input expected = do - let pieces = textPieces textFragmentSize input - ls <- J.each pieces & J.lines & J.toList - assertEqual "lines do not match expected" expected ls - -sleep :: Delay -> IO () -sleep (Delay d) = threadDelay d - -newtype Delay = Delay Int - -cents :: Int -> Delay -cents i = Delay $ i * 1e4 - -main :: IO () -main = defaultMain tests - -time :: IO a -> IO (NominalDiffTime, a) -time action = do - start <- getCurrentTime - a <- action - stop <- getCurrentTime - pure (diffUTCTime stop start, a) - +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE BlockArguments #-} +{-# LANGUAGE DeriveFunctor #-} +{-# LANGUAGE DerivingStrategies #-} +{-# LANGUAGE DerivingVia #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE ImportQualifiedPost #-} +{-# LANGUAGE MultiWayIf #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE ViewPatterns #-} +{-# LANGUAGE TupleSections #-} +{-# LANGUAGE LinearTypes #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE PatternSynonyms #-} +{-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE NamedFieldPuns #-} +{-# LANGUAGE GADTSyntax #-} +{-# LANGUAGE ExistentialQuantification #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE StandaloneKindSignatures #-} +{-# LANGUAGE PartialTypeSignatures #-} +{-# LANGUAGE StandaloneDeriving #-} +{-# LANGUAGE ApplicativeDo #-} +{-# LANGUAGE NumDecimals #-} +{-# OPTIONS_GHC -Wno-partial-type-signatures #-} +module Main where + +import Test.Tasty +import Test.Tasty.HUnit +import Data.IORef +import Data.Text (Text) +import Data.Text qualified as T +import Data.Text.IO qualified as T +import Data.Text.Encoding qualified as T +import Data.Text.Encoding.Error qualified as T +import Data.Text.Lazy qualified as TL +import Data.Text.Lazy.Encoding qualified as TL +import Data.ByteString (ByteString) +import Data.ByteString qualified as B +import Data.ByteString.Lazy qualified as BL +import Data.Time.Clock +import Data.Foldable +import Debug.Trace +import Data.Functor.Identity +import Control.Concurrent +import Data.List +import Control.Exception + +import Jet +import Jet qualified as J + +tests :: TestTree +tests = + testGroup + "All" + [ + testGroup "byteSplitter" $ + let tests = do + -- splitSize <- [1] + -- bucketSize <- [2] + splitSize <- [1..7] + bucketSize <- [1..10] + pure $ + testCase ("splitter splitSize=" ++ show splitSize ++ " bucketSize=" ++ show bucketSize) $ + assertBytesCorrectlySplit bucketSize (bytePieces splitSize az) + in tests + , + testGroup "byteBundleSplitter" $ + let tests = do + splitSize <- [1..7] + bucketSize <- [splitSize..13] + pure $ + testCase ("splitter splitSize=" ++ show splitSize ++ " bucketSize=" ++ show bucketSize) $ + assertByteBundlesCorrectlySplit bucketSize (bytePieces splitSize az) + in tests + , + testGroup "lines" $ + let tests = do + fragmentSize <- [1..13] + pure $ + testCase ("fragment size " ++ show fragmentSize) $ + assertLines fragmentSize lineData01 lineExpected01 + in tests + , + testGroup "process" $ + [ + testCase "simple" $ + do resultLines <- + mempty + & linesThroughProcess defaults (shell "echo foo") + & J.toList + assertEqual "input and output lines don't match" (textToLine . T.pack <$> ["foo"]) resultLines + , + testCase "simple 2" $ + do let inputLines = textToLine . T.pack <$> ["aaa","bbb","ccc"] + resultLines <- + J.each inputLines + & linesThroughProcess defaults (shell "cat") + & J.toList + assertEqual "input and output lines don't match" inputLines resultLines + , + testCase "interruption" $ + do let expectedLines = textToLine . T.pack <$> ["aaa","bbb"] + resultLines <- + mempty + & linesThroughProcess defaults (shell "{ printf \"aaa\\nbbb\\nccc\\n\" ; sleep infinity ; }") + & J.limit 2 + & J.toList + assertEqual "unexpected lines at output" expectedLines resultLines + ] + , + testGroup "concurrency" $ + [ + testCase "compare" $ do + let yieldAfter d x = sleep d *> pure x + delay = cents 100 + upstream = J.each "abcde" + (ts, rsequential) <- upstream & J.traverse (yieldAfter delay) & J.toList & time + (t1, rconc1) <- upstream & J.traverseConcurrently (numberOfWorkers 2) (yieldAfter delay) & J.toList & time + (t2, rconc2) <- upstream & J.traverseConcurrently (numberOfWorkers 10) (yieldAfter delay) & J.toList & time + let (rsequential', rconc1', rconc2') = (sort rsequential, sort rconc1, sort rconc2) + assertEqual "sequential != conc" rsequential' rconc1' + assertEqual "conc != conc 2" rconc1' rconc2' + assertBool "conc not faster" (t1 < ts) + assertBool "conc2 not faster" (t2 < t1) + pure () + , testCase "cancelation" $ do + latch <- newEmptyMVar + l <- + J.each "ab" + & traverseConcurrently (numberOfWorkers 2) + (\c -> do let delay 'a' = threadDelay 9999e6 + delay _ = pure () + (delay c *> pure c) `Control.Exception.onException` putMVar latch ()) + & J.limit 1 + & J.length + _ <- takeMVar latch + assertEqual "only 1 element" 1 l + pure () + ] + ] + +az :: ByteString +az = ['a'..'z'] & T.pack & T.encodeUtf8 + +bytePieces :: Int -> ByteString -> [ByteString] +bytePieces size = + let go b = + if B.null b + then [] + else let (left,right) = B.splitAt size b + in left : go right + in go + +assertBytesCorrectlySplit :: Int -> [ByteString] -> IO () +assertBytesCorrectlySplit bucketSize inputs = do + let buckets = Prelude.repeat bucketSize + j = J.recast (J.bytesOverBuckets buckets) combineIntoLists (J.each inputs) + fragmentedGroups <- J.toList j + let groups :: [ByteString] = mconcat <$> fragmentedGroups + concatenatedInput = T.decodeUtf8 $ mconcat inputs + concatenatedOutput = T.decodeUtf8 $ mconcat groups + assertEqual "combined inputs and result" concatenatedInput concatenatedOutput + -- traceIO "--------------------------" + -- traceIO $ "+ original groups = " ++ show fragmentedGroups + -- traceIO $ "+ collected groups = " ++ show groups + -- traceIO $ "* bucket size = " ++ show bucketSize + -- traceIO $ show $ B.length <$> Prelude.init groups + -- traceIO "--------------------------" + assertBool "group sizes are wrong" $ all (\g -> B.length g == bucketSize) (Prelude.init groups) + pure () + +assertByteBundlesCorrectlySplit :: Int -> [ByteString] -> IO () +assertByteBundlesCorrectlySplit bucketSize inputs = do + let buckets = Prelude.repeat bucketSize + j = J.recast (J.byteBundlesOverBuckets buckets) combineIntoLists (bundle . Identity <$> J.each inputs) + fragmentedGroups <- J.toList j + let groups :: [ByteString] = mconcat <$> fragmentedGroups + concatenatedInput = T.decodeUtf8 $ mconcat inputs + concatenatedOutput = T.decodeUtf8 $ mconcat groups + concatenatedOutput' = T.decodeUtf8 $ mconcat $ Data.List.intersperse (T.encodeUtf8 (T.singleton '-')) groups + assertEqual "combined inputs and result" concatenatedInput concatenatedOutput + -- traceIO "--------------------------" + -- traceIO $ "+ inputs = " ++ show inputs + -- traceIO $ "+ original groups = " ++ show fragmentedGroups + -- traceIO $ "+ collected groups = " ++ show groups + -- traceIO $ "* bucket size = " ++ show bucketSize + -- traceIO $ show $ B.length <$> Prelude.init groups + -- traceIO "--------------------------" + assertBool "group sizes are wrong" $ all (\g -> B.length g <= bucketSize) (Prelude.init groups) + Data.Foldable.for_ inputs \i -> + assertBool "bundle divided" $ T.isInfixOf (T.decodeUtf8 i) concatenatedOutput' + pure () + + +lineData01 :: Text +lineData01 = T.pack "aaa\nbb\nccc\ndddd\n\neee\n\n\nfffffffff\ng\niiiii" + +lineExpected01 :: [Line] +lineExpected01 = textToLine . T.pack <$> ["aaa","bb","ccc","dddd","","eee","","","fffffffff", "g", "iiiii"] + +textPieces :: Int -> Text -> [Text] +textPieces size = + let go t = + if T.null t + then [] + else let (left,right) = T.splitAt size t + in left : go right + in go + +assertLines :: Int -> Text -> [Line] -> IO () +assertLines textFragmentSize input expected = do + let pieces = textPieces textFragmentSize input + ls <- J.each pieces & J.lines & J.toList + assertEqual "lines do not match expected" expected ls + +sleep :: Delay -> IO () +sleep (Delay d) = threadDelay d + +newtype Delay = Delay Int + +cents :: Int -> Delay +cents i = Delay $ i * 1e4 + +main :: IO () +main = defaultMain tests + +time :: IO a -> IO (NominalDiffTime, a) +time action = do + start <- getCurrentTime + a <- action + stop <- getCurrentTime + pure (diffUTCTime stop start, a) +