c5b2ad3 · streamly · git.iko.soy

git.iko.soy/streamly

on master help sign in

Refactor/add cross product APIs

c5b2ad3Harendra Kumar6 years ago

src/Streamly/Data/Unfold.hs+3−1

compare with master

unifiedsplit

@@ -97,8 +97,10 @@ module Streamly.Data.Unfold

97 , zipWithM

98 , zipWith

100+ -- ** Cross Product

101+ , crossWith

102+

100 -- ** Nesting

103 -- ** Nesting

101− , cross

102 , many

104 , many

103

105

104 -- ** Exceptions

106 -- ** Exceptions

src/Streamly/Internal/Data/Unfold.hs+5−41

@@ -161,14 +161,17 @@ module Streamly.Internal.Data.Unfold

161 , zipWith

162 , teeZipWith

163

164− -- ** Nesting

164+ -- ** Cross product

165+ , crossWithM

166+ , crossWith

165 , cross

167 , cross

166 , apply

168 , apply

169+

170+ -- ** Nesting

167 , ConcatState (..)

171 , ConcatState (..)

168 , many

172 , many

169 , concatMapM

173 , concatMapM

170 , bind

174 , bind

171− , outerProduct

172

175

173 -- ** Exceptions

176 -- ** Exceptions

174 , gbracket_

177 , gbracket_

@@ -896,45 +899,6 @@ teeZipWith :: Monad m

896 => (a -> b -> c) -> Unfold m x a -> Unfold m x b -> Unfold m x c

899 => (a -> b -> c) -> Unfold m x a -> Unfold m x b -> Unfold m x c

897teeZipWith f unf1 unf2 = lmap (\x -> (x,x)) $ zipWith f unf1 unf2

900teeZipWith f unf1 unf2 = lmap (\x -> (x,x)) $ zipWith f unf1 unf2

898

901

899−-------------------------------------------------------------------------------

900−-- Nested

901−-------------------------------------------------------------------------------

902−

903−data OuterProductState s1 s2 sy x y =

904− OuterProductOuter s1 y | OuterProductInner s1 sy s2 x

905−

906−-- XXX this can be written in terms of "cross".

907−-- XXX Remove this in favor of cross?

908−--

909−-- | Create an outer product (vector product or cartesian product) of the

910−-- output streams of two unfolds.

911−--

912−{-# INLINE_NORMAL outerProduct #-}

913−outerProduct :: Monad m

914− => Unfold m a b -> Unfold m c d -> Unfold m (a, c) (b, d)

915−outerProduct (Unfold step1 inject1) (Unfold step2 inject2) = Unfold step inject

916− where

917− inject (x, y) = do

918− s1 <- inject1 x

919− return $ OuterProductOuter s1 y

920−

921− {-# INLINE_LATE step #-}

922− step (OuterProductOuter st1 sy) = do

923− r <- step1 st1

924− case r of

925− Yield x s -> do

926− s2 <- inject2 sy

927− return $ Skip (OuterProductInner s sy s2 x)

928− Skip s -> return $ Skip (OuterProductOuter s sy)

929− Stop -> return Stop

930−

931− step (OuterProductInner ost sy ist x) = do

932− r <- step2 ist

933− return $ case r of

934− Yield y s -> Yield (x, y) (OuterProductInner ost sy s x)

935− Skip s -> Skip (OuterProductInner ost sy s x)

936− Stop -> Skip (OuterProductOuter ost sy)

937−

938------------------------------------------------------------------------------

902------------------------------------------------------------------------------

939-- Exceptions

903-- Exceptions

940------------------------------------------------------------------------------

904------------------------------------------------------------------------------

src/Streamly/Internal/Data/Unfold/Type.hs+41−9

@@ -24,14 +24,18 @@ module Streamly.Internal.Data.Unfold.Type

24 , ConcatState (..)

25 , many

27+ -- Applicative

27 , apSequence

28 , apSequence

28 , apDiscardSnd

29 , apDiscardSnd

30+ , crossWithM

31+ , crossWith

29 , cross

32 , cross

30 , apply

33 , apply

31− , bind

35+ -- Monad

33 , concatMapM

36 , concatMapM

34 , concatMap

37 , concatMap

38+ , bind

36 , zipWithM

40 , zipWithM

37 , zipWith

41 , zipWith

@@ -144,11 +148,13 @@ apDiscardSnd (Unfold _step1 _inject1) (Unfold _step2 _inject2) = undefined

144data Cross a s1 b s2 = CrossOuter a s1 | CrossInner a s1 b s2

148data Cross a s1 b s2 = CrossOuter a s1 | CrossInner a s1 b s2

145

149

146-- | Create a cross product (vector product or cartesian product) of the

150-- | Create a cross product (vector product or cartesian product) of the

147−-- output streams of two unfolds.

151+-- output streams of two unfolds using a monadic combining function.

148--

152--

153+-- /Pre-release/

149−{-# INLINE_NORMAL cross #-}

154+{-# INLINE_NORMAL crossWithM #-}

155+crossWithM :: Monad m =>

150−cross :: Monad m => Unfold m a b -> Unfold m a c -> Unfold m a (b, c)

156+ (b -> c -> m d) -> Unfold m a b -> Unfold m a c -> Unfold m a d

151−cross (Unfold step1 inject1) (Unfold step2 inject2) = Unfold step inject

157+crossWithM f (Unfold step1 inject1) (Unfold step2 inject2) = Unfold step inject

152

158

153 where

159 where

154

160

@@ -168,10 +174,36 @@ cross (Unfold step1 inject1) (Unfold step2 inject2) = Unfold step inject

168

174

169 step (CrossInner a s1 b s2) = do

175 step (CrossInner a s1 b s2) = do

170 r <- step2 s2

176 r <- step2 s2

171− return $ case r of

177+ case r of

172− Yield c s -> Yield (b, c) (CrossInner a s1 b s)

178+ Yield c s -> f b c >>= \d -> return $ Yield d (CrossInner a s1 b s)

173− Skip s -> Skip (CrossInner a s1 b s)

179+ Skip s -> return $ Skip (CrossInner a s1 b s)

174− Stop -> Skip (CrossOuter a s1)

180+ Stop -> return $ Skip (CrossOuter a s1)

181+

182+-- | Like 'crossWithM' but uses a pure combining function.

183+--

184+-- > crossWith f = crossWithM (\b c -> return $ f b c)

185+--

186+-- /Pre-release/

187+{-# INLINE crossWith #-}

188+crossWith :: Monad m =>

189+ (b -> c -> d) -> Unfold m a b -> Unfold m a c -> Unfold m a d

190+crossWith f = crossWithM (\b c -> return $ f b c)

191+

192+-- | See 'crossWith'.

193+--

194+-- > cross = crossWith (,)

195+--

196+-- To cross the streams from a tuple we can write:

197+--

198+-- @

199+-- crossProduct :: Monad m => Unfold m a b -> Unfold m c d -> Unfold m (a, c) (b, d)

200+-- crossProduct u1 u2 = cross (lmap fst u1) (lmap snd u2)

201+-- @

202+--

203+-- /Pre-release/

204+{-# INLINE_NORMAL cross #-}

205+cross :: Monad m => Unfold m a b -> Unfold m a c -> Unfold m a (b, c)

206+cross = crossWith (,)

175

207

176apply :: Monad m => Unfold m a (b -> c) -> Unfold m a b -> Unfold m a c

208apply :: Monad m => Unfold m a (b -> c) -> Unfold m a b -> Unfold m a c

177apply u1 u2 = fmap (\(a, b) -> a b) (cross u1 u2)

209apply u1 u2 = fmap (\(a, b) -> a b) (cross u1 u2)

test/Streamly/Test/Data/Unfold.hs+5−1

@@ -335,10 +335,14 @@ outerProduct :: Bool

335outerProduct =

336 let unf1 = UF.enumerateFromToIntegral 10

337 unf2 = UF.enumerateFromToIntegral 20

338− unf = UF.outerProduct unf1 unf2

338+ unf = crossProduct unf1 unf2

339 lst = [(a, b) :: (Int, Int) | a <- [0 .. 10], b <- [0 .. 20]]

340 in testUnfold unf ((0, 0) :: (Int, Int)) lst

341

342+ where

343+

344+ crossProduct u1 u2 = UF.cross (UF.lmap fst u1) (UF.lmap snd u2)

345+

342concatMapM :: Bool

346concatMapM :: Bool

343concatMapM =

347concatMapM =

344 let inner b =

348 let inner b =