-
Notifications
You must be signed in to change notification settings - Fork 2
/
Copy pathProducer.hs
96 lines (86 loc) · 2.88 KB
/
Producer.hs
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
{-# LANGUAGE DeriveFoldable #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RankNTypes #-}
module Stutter.Producer where
import Control.Monad.IO.Class
import Control.Monad.State
import Control.Monad.Trans.Resource (MonadResource)
import Data.Conduit
import Data.Conduit.Internal (zipSources)
import Data.Monoid
import System.IO (stdin)
import qualified Data.ByteString.Lazy as BL
import qualified Data.Conduit.Binary as CB
import qualified Data.Conduit.Combinators as CL
import qualified Data.Text as T
data Range
= IntRange (Int, Int)
| CharRange (Char, Char)
deriving (Eq, Show)
data ProducerGroup_ a
= PSum (ProducerGroup_ a) (ProducerGroup_ a)
| PProduct (ProducerGroup_ a) (ProducerGroup_ a)
| PZip (ProducerGroup_ a) (ProducerGroup_ a)
| PRepeat (ProducerGroup_ a)
| PRanges [Range]
| PFile FilePath
| PStdin a
| PText T.Text
deriving (Eq, Show, Functor, Foldable, Traversable)
type ProducerGroup = ProducerGroup_ BL.ByteString
prepareStdin :: ProducerGroup_ () -> IO (ProducerGroup_ BL.ByteString)
prepareStdin p = evalStateT (traverse f p) Nothing
where
f () = get >>= \case
Just bs -> return bs
Nothing -> do
bs <- liftIO $ BL.hGetContents stdin
modify (const $ Just bs)
return bs
cardinality :: ProducerGroup_ a -> Maybe Int
cardinality (PSum p p') = (+) <$> cardinality p <*> cardinality p'
cardinality (PProduct p p') = (*) <$> cardinality p <*> cardinality p'
cardinality (PZip p p') = min <$> cardinality p <*> cardinality p'
cardinality (PRepeat _) = Nothing
cardinality (PRanges rs) = pure $ sum $ map rangeCardinality rs
where
rangeCardinality (IntRange (a,z)) = length [a..z]
rangeCardinality (CharRange (a,z)) = length [a..z]
cardinality PFile{} = Nothing
cardinality PStdin{} = Nothing
cardinality PText{} = pure 1
produceRanges :: (Monad m) => [Range] -> Producer m T.Text
produceRanges = CL.yieldMany
. concat
. map rangeToList
where
rangeToList (IntRange (a,z)) = tshow <$> [a..z]
rangeToList (CharRange (a,z)) = T.pack . (:[]) <$> [a..z]
tshow = T.pack . show
produceGroup
:: (MonadIO m, MonadResource m)
=> ProducerGroup
-> Source m T.Text
produceGroup (PRanges rs) = produceRanges rs
produceGroup (PText t) = yield t
produceGroup (PProduct g g') =
produceGroup g
.| awaitForever ( \t -> forever (yield ())
.| produceGroup g'
.| awaitForever (\t' -> yield (t <> t')))
produceGroup (PSum g g') = produceGroup g >> produceGroup g'
produceGroup (PZip g g') =
zipSources (produceGroup g) (produceGroup g')
.| CL.map (uncurry (<>))
produceGroup (PRepeat g) = forever $ produceGroup g
produceGroup (PFile f) =
CB.sourceFile f
.| CB.lines
.| CL.decodeUtf8
produceGroup (PStdin bs) =
CB.sourceLbs bs
.| CB.lines
.| CL.decodeUtf8