Skip to content

Commit 640b4c6

Browse files
doliodolio
authored
Merge pull request #41 from Shimuuar/iterate-unfold
Add monadic variants for iterateN, unfoldr, unfoldrN
2 parents 8061d1f + 4662eb4 commit 640b4c6

File tree

7 files changed

+167
-18
lines changed

7 files changed

+167
-18
lines changed

‎Data/Vector.hs

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -56,10 +56,11 @@ module Data.Vector (
5656
empty, singleton, replicate, generate, iterateN,
5757

5858
-- ** Monadic initialisation
59-
replicateM, generateM, create, createT,
59+
replicateM, generateM, iterateNM, create, createT,
6060

6161
-- ** Unfolding
6262
unfoldr, unfoldrN,
63+
unfoldrM, unfoldrNM,
6364
constructN, constructrN,
6465

6566
-- ** Enumeration
@@ -652,6 +653,22 @@ unfoldrN :: Int -> (b -> Maybe (a, b)) -> b -> Vector a
652653
{-# INLINE unfoldrN #-}
653654
unfoldrN = G.unfoldrN
654655

656+
-- | /O(n)/ Construct a vector by repeatedly applying the monadic
657+
-- generator function to a seed. The generator function yields 'Just'
658+
-- the next element and the new seed or 'Nothing' if there are no more
659+
-- elements.
660+
unfoldrM :: (Monad m) => (b -> m (Maybe (a, b))) -> b -> m (Vector a)
661+
{-# INLINE unfoldrM #-}
662+
unfoldrM = G.unfoldrM
663+
664+
-- | /O(n)/ Construct a vector by repeatedly applying the monadic
665+
-- generator function to a seed. The generator function yields 'Just'
666+
-- the next element and the new seed or 'Nothing' if there are no more
667+
-- elements.
668+
unfoldrNM :: (Monad m) => Int -> (b -> m (Maybe (a, b))) -> b -> m (Vector a)
669+
{-# INLINE unfoldrNM #-}
670+
unfoldrNM = G.unfoldrNM
671+
655672
-- | /O(n)/ Construct a vector with @n@ elements by repeatedly applying the
656673
-- generator function to the already constructed part of the vector.
657674
--
@@ -745,6 +762,11 @@ generateM :: Monad m => Int -> (Int -> m a) -> m (Vector a)
745762
{-# INLINE generateM #-}
746763
generateM = G.generateM
747764

765+
-- | /O(n)/ Apply monadic function n times to value. Zeroth element is original value.
766+
iterateNM :: Monad m => Int -> (a -> m a) -> a -> m (Vector a)
767+
{-# INLINE iterateNM #-}
768+
iterateNM = G.iterateNM
769+
748770
-- | Execute the monadic action and freeze the resulting vector.
749771
--
750772
-- @

‎Data/Vector/Generic.hs

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -39,10 +39,11 @@ module Data.Vector.Generic (
3939
empty, singleton, replicate, generate, iterateN,
4040

4141
-- ** Monadic initialisation
42-
replicateM, generateM, create, createT,
42+
replicateM, generateM, iterateNM, create, createT,
4343

4444
-- ** Unfolding
4545
unfoldr, unfoldrN,
46+
unfoldrM, unfoldrNM,
4647
constructN, constructrN,
4748

4849
-- ** Enumeration
@@ -551,6 +552,22 @@ unfoldrN :: Vector v a => Int -> (b -> Maybe (a, b)) -> b -> v a
551552
{-# INLINE unfoldrN #-}
552553
unfoldrN n f = unstream . Bundle.unfoldrN n f
553554

555+
-- | /O(n)/ Construct a vector by repeatedly applying the monadic
556+
-- generator function to a seed. The generator function yields 'Just'
557+
-- the next element and the new seed or 'Nothing' if there are no more
558+
-- elements.
559+
unfoldrM :: (Monad m, Vector v a) => (b -> m (Maybe (a, b))) -> b -> m (v a)
560+
{-# INLINE unfoldrM #-}
561+
unfoldrM f = unstreamM . MBundle.unfoldrM f
562+
563+
-- | /O(n)/ Construct a vector by repeatedly applying the monadic
564+
-- generator function to a seed. The generator function yields 'Just'
565+
-- the next element and the new seed or 'Nothing' if there are no more
566+
-- elements.
567+
unfoldrNM :: (Monad m, Vector v a) => Int -> (b -> m (Maybe (a, b))) -> b -> m (v a)
568+
{-# INLINE unfoldrNM #-}
569+
unfoldrNM n f = unstreamM . MBundle.unfoldrNM n f
570+
554571
-- | /O(n)/ Construct a vector with @n@ elements by repeatedly applying the
555572
-- generator function to the already constructed part of the vector.
556573
--
@@ -711,6 +728,11 @@ generateM :: (Monad m, Vector v a) => Int -> (Int -> m a) -> m (v a)
711728
{-# INLINE generateM #-}
712729
generateM n f = unstreamM (MBundle.generateM n f)
713730

731+
-- | /O(n)/ Apply monadic function n times to value. Zeroth element is original value.
732+
iterateNM :: (Monad m, Vector v a) => Int -> (a -> m a) -> a -> m (v a)
733+
{-# INLINE iterateNM #-}
734+
iterateNM n f x = unstreamM (MBundle.iterateNM n f x)
735+
714736
-- | Execute the monadic action and freeze the resulting vector.
715737
--
716738
-- @

‎Data/Vector/Primitive.hs

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -42,10 +42,11 @@ module Data.Vector.Primitive (
4242
empty, singleton, replicate, generate, iterateN,
4343

4444
-- ** Monadic initialisation
45-
replicateM, generateM, create, createT,
45+
replicateM, generateM, iterateNM, create, createT,
4646

4747
-- ** Unfolding
4848
unfoldr, unfoldrN,
49+
unfoldrM, unfoldrNM,
4950
constructN, constructrN,
5051

5152
-- ** Enumeration
@@ -516,6 +517,22 @@ unfoldrN :: Prim a => Int -> (b -> Maybe (a, b)) -> b -> Vector a
516517
{-# INLINE unfoldrN #-}
517518
unfoldrN = G.unfoldrN
518519

520+
-- | /O(n)/ Construct a vector by repeatedly applying the monadic
521+
-- generator function to a seed. The generator function yields 'Just'
522+
-- the next element and the new seed or 'Nothing' if there are no more
523+
-- elements.
524+
unfoldrM :: (Monad m, Prim a) => (b -> m (Maybe (a, b))) -> b -> m (Vector a)
525+
{-# INLINE unfoldrM #-}
526+
unfoldrM = G.unfoldrM
527+
528+
-- | /O(n)/ Construct a vector by repeatedly applying the monadic
529+
-- generator function to a seed. The generator function yields 'Just'
530+
-- the next element and the new seed or 'Nothing' if there are no more
531+
-- elements.
532+
unfoldrNM :: (Monad m, Prim a) => Int -> (b -> m (Maybe (a, b))) -> b -> m (Vector a)
533+
{-# INLINE unfoldrNM #-}
534+
unfoldrNM = G.unfoldrNM
535+
519536
-- | /O(n)/ Construct a vector with @n@ elements by repeatedly applying the
520537
-- generator function to the already constructed part of the vector.
521538
--
@@ -609,6 +626,11 @@ generateM :: (Monad m, Prim a) => Int -> (Int -> m a) -> m (Vector a)
609626
{-# INLINE generateM #-}
610627
generateM = G.generateM
611628

629+
-- | /O(n)/ Apply monadic function n times to value. Zeroth element is original value.
630+
iterateNM :: (Monad m, Prim a) => Int -> (a -> m a) -> a -> m (Vector a)
631+
{-# INLINE iterateNM #-}
632+
iterateNM = G.iterateNM
633+
612634
-- | Execute the monadic action and freeze the resulting vector.
613635
--
614636
-- @

‎Data/Vector/Storable.hs

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -39,10 +39,11 @@ module Data.Vector.Storable (
3939
empty, singleton, replicate, generate, iterateN,
4040

4141
-- ** Monadic initialisation
42-
replicateM, generateM, create, createT,
42+
replicateM, generateM, iterateNM, create, createT,
4343

4444
-- ** Unfolding
4545
unfoldr, unfoldrN,
46+
unfoldrM, unfoldrNM,
4647
constructN, constructrN,
4748

4849
-- ** Enumeration
@@ -526,6 +527,22 @@ unfoldrN :: Storable a => Int -> (b -> Maybe (a, b)) -> b -> Vector a
526527
{-# INLINE unfoldrN #-}
527528
unfoldrN = G.unfoldrN
528529

530+
-- | /O(n)/ Construct a vector by repeatedly applying the monadic
531+
-- generator function to a seed. The generator function yields 'Just'
532+
-- the next element and the new seed or 'Nothing' if there are no more
533+
-- elements.
534+
unfoldrM :: (Monad m, Storable a) => (b -> m (Maybe (a, b))) -> b -> m (Vector a)
535+
{-# INLINE unfoldrM #-}
536+
unfoldrM = G.unfoldrM
537+
538+
-- | /O(n)/ Construct a vector by repeatedly applying the monadic
539+
-- generator function to a seed. The generator function yields 'Just'
540+
-- the next element and the new seed or 'Nothing' if there are no more
541+
-- elements.
542+
unfoldrNM :: (Monad m, Storable a) => Int -> (b -> m (Maybe (a, b))) -> b -> m (Vector a)
543+
{-# INLINE unfoldrNM #-}
544+
unfoldrNM = G.unfoldrNM
545+
529546
-- | /O(n)/ Construct a vector with @n@ elements by repeatedly applying the
530547
-- generator function to the already constructed part of the vector.
531548
--
@@ -619,6 +636,11 @@ generateM :: (Monad m, Storable a) => Int -> (Int -> m a) -> m (Vector a)
619636
{-# INLINE generateM #-}
620637
generateM = G.generateM
621638

639+
-- | /O(n)/ Apply monadic function n times to value. Zeroth element is original value.
640+
iterateNM :: (Monad m, Storable a) => Int -> (a -> m a) -> a -> m (Vector a)
641+
{-# INLINE iterateNM #-}
642+
iterateNM = G.iterateNM
643+
622644
-- | Execute the monadic action and freeze the resulting vector.
623645
--
624646
-- @

‎Data/Vector/Unboxed.hs

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -62,10 +62,11 @@ module Data.Vector.Unboxed (
6262
empty, singleton, replicate, generate, iterateN,
6363

6464
-- ** Monadic initialisation
65-
replicateM, generateM, create, createT,
65+
replicateM, generateM, iterateNM, create, createT,
6666

6767
-- ** Unfolding
6868
unfoldr, unfoldrN,
69+
unfoldrM, unfoldrNM,
6970
constructN, constructrN,
7071

7172
-- ** Enumeration
@@ -495,6 +496,22 @@ unfoldrN :: Unbox a => Int -> (b -> Maybe (a, b)) -> b -> Vector a
495496
{-# INLINE unfoldrN #-}
496497
unfoldrN = G.unfoldrN
497498

499+
-- | /O(n)/ Construct a vector by repeatedly applying the monadic
500+
-- generator function to a seed. The generator function yields 'Just'
501+
-- the next element and the new seed or 'Nothing' if there are no more
502+
-- elements.
503+
unfoldrM :: (Monad m, Unbox a) => (b -> m (Maybe (a, b))) -> b -> m (Vector a)
504+
{-# INLINE unfoldrM #-}
505+
unfoldrM = G.unfoldrM
506+
507+
-- | /O(n)/ Construct a vector by repeatedly applying the monadic
508+
-- generator function to a seed. The generator function yields 'Just'
509+
-- the next element and the new seed or 'Nothing' if there are no more
510+
-- elements.
511+
unfoldrNM :: (Monad m, Unbox a) => Int -> (b -> m (Maybe (a, b))) -> b -> m (Vector a)
512+
{-# INLINE unfoldrNM #-}
513+
unfoldrNM = G.unfoldrNM
514+
498515
-- | /O(n)/ Construct a vector with @n@ elements by repeatedly applying the
499516
-- generator function to the already constructed part of the vector.
500517
--
@@ -588,6 +605,11 @@ generateM :: (Monad m, Unbox a) => Int -> (Int -> m a) -> m (Vector a)
588605
{-# INLINE generateM #-}
589606
generateM = G.generateM
590607

608+
-- | /O(n)/ Apply monadic function n times to value. Zeroth element is original value.
609+
iterateNM :: (Monad m, Unbox a) => Int -> (a -> m a) -> a -> m (Vector a)
610+
{-# INLINE iterateNM #-}
611+
iterateNM = G.iterateNM
612+
591613
-- | Execute the monadic action and freeze the resulting vector.
592614
--
593615
-- @

‎tests/Tests/Vector.hs

Lines changed: 23 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ module Tests.Vector (tests) where
33
import Boilerplater
44
import Utilities as Util
55

6+
import Data.Functor.Identity
67
import qualified Data.Traversable as T (Traversable(..))
78
import Data.Foldable (Foldable(foldMap))
89

@@ -117,14 +118,14 @@ testPolymorphicFunctions _ = $(testProperties [
117118

118119
-- Initialisation (FIXME)
119120
'prop_empty, 'prop_singleton, 'prop_replicate,
120-
'prop_generate, 'prop_iterateN,
121+
'prop_generate, 'prop_iterateN, 'prop_iterateNM,
121122

122123
-- Monadic initialisation (FIXME)
123124
'prop_createT,
124125
{- 'prop_replicateM, 'prop_generateM, 'prop_create, -}
125126

126-
-- Unfolding (FIXME)
127-
{- 'prop_unfoldr, prop_unfoldrN, -}
127+
-- Unfolding
128+
'prop_unfoldr, 'prop_unfoldrN, 'prop_unfoldrM, 'prop_unfoldrNM,
128129
'prop_constructN, 'prop_constructrN,
129130

130131
-- Enumeration? (FIXME?)
@@ -235,7 +236,8 @@ testPolymorphicFunctions _ = $(testProperties [
235236
= (\n _ -> n < 1000) ===> V.generate `eq` Util.generate
236237
prop_iterateN :: P (Int -> (a -> a) -> a -> v a)
237238
= (\n _ _ -> n < 1000) ===> V.iterateN `eq` (\n f -> take n . iterate f)
238-
239+
prop_iterateNM :: P (Int -> (a -> Writer [Int] a) -> a -> Writer [Int] (v a))
240+
= (\n _ _ -> n < 1000) ===> V.iterateNM `eq` Util.iterateNM
239241
prop_createT :: P ((a, v a) -> (a, v a))
240242
prop_createT = (\v -> V.createT (T.mapM V.thaw v)) `eq` id
241243

@@ -434,12 +436,26 @@ testPolymorphicFunctions _ = $(testProperties [
434436

435437
-- Because the vectors are strict, we need to be totally sure that the unfold eventually terminates. This
436438
-- is achieved by injecting our own bit of state into the unfold - the maximum number of unfolds allowed.
437-
limitUnfolds f (theirs, ours) | ours >= 0
438-
, Just (out, theirs') <- f theirs = Just (out, (theirs', ours - 1))
439-
| otherwise = Nothing
439+
limitUnfolds f (theirs, ours)
440+
| ours > 0
441+
, Just (out, theirs') <- f theirs = Just (out, (theirs', ours - 1))
442+
| otherwise = Nothing
443+
limitUnfoldsM f (theirs, ours)
444+
| ours > 0 = do r <- f theirs
445+
return $ (\(a,b) -> (a,(b,ours - 1))) `fmap` r
446+
| otherwise = return Nothing
447+
448+
440449
prop_unfoldr :: P (Int -> (Int -> Maybe (a,Int)) -> Int -> v a)
441450
= (\n f a -> V.unfoldr (limitUnfolds f) (a, n))
442451
`eq` (\n f a -> unfoldr (limitUnfolds f) (a, n))
452+
prop_unfoldrN :: P (Int -> (Int -> Maybe (a,Int)) -> Int -> v a)
453+
= V.unfoldrN `eq` (\n f a -> unfoldr (limitUnfolds f) (a, n))
454+
prop_unfoldrM :: P (Int -> (Int -> Writer [Int] (Maybe (a,Int))) -> Int -> Writer [Int] (v a))
455+
= (\n f a -> V.unfoldrM (limitUnfoldsM f) (a,n))
456+
`eq` (\n f a -> Util.unfoldrM (limitUnfoldsM f) (a, n))
457+
prop_unfoldrNM :: P (Int -> (Int -> Writer [Int] (Maybe (a,Int))) -> Int -> Writer [Int] (v a))
458+
= V.unfoldrNM `eq` (\n f a -> Util.unfoldrM (limitUnfoldsM f) (a, n))
443459

444460
prop_constructN = \f -> forAll (choose (0,20)) $ \n -> unP prop n f
445461
where

‎tests/Utilities.hs

Lines changed: 29 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -52,8 +52,10 @@ instance Arbitrary a => Arbitrary (S.Bundle v a) where
5252
instance CoArbitrary a => CoArbitrary (S.Bundle v a) where
5353
coarbitrary = coarbitrary . S.toList
5454

55-
instance Arbitrary a => Arbitrary (Writer a ()) where
56-
arbitrary = fmap (writer . ((,) ())) arbitrary
55+
instance (Arbitrary a, Arbitrary b) => Arbitrary (Writer a b) where
56+
arbitrary = do b <- arbitrary
57+
a <- arbitrary
58+
return $ writer (b,a)
5759

5860
instance CoArbitrary a => CoArbitrary (Writer a ()) where
5961
coarbitrary = coarbitrary . runWriter
@@ -148,13 +150,13 @@ instance (Eq a, TestData a) => TestData (Identity a) where
148150
type EqTest (Identity a) = Property
149151
equal = (property .) . on (==) runIdentity
150152

151-
instance (Eq a, TestData a, Monoid a) => TestData (Writer a ()) where
152-
type Model (Writer a ()) = Writer (Model a) ()
153+
instance (Eq a, TestData a, Eq b, TestData b, Monoid a) => TestData (Writer a b) where
154+
type Model (Writer a b) = Writer (Model a) (Model b)
153155
model = mapWriter model
154156
unmodel = mapWriter unmodel
155157

156-
type EqTest (Writer a ()) = Property
157-
equal = (property .) . on (==) execWriter
158+
type EqTest (Writer a b) = Property
159+
equal = (property .) . on (==) runWriter
158160

159161
instance (Eq a, Eq b, TestData a, TestData b) => TestData (a,b) where
160162
type Model (a,b) = (Model a, Model b)
@@ -325,3 +327,24 @@ maxIndex = fst . foldr1 imax . zip [0..]
325327
imax (i,x) (j,y) | x >= y = (i,x)
326328
| otherwise = (j,y)
327329

330+
iterateNM :: Monad m => Int -> (a -> m a) -> a -> m [a]
331+
iterateNM n f x
332+
| n <= 0 = return []
333+
| n == 1 = return [x]
334+
| otherwise = do x' <- f x
335+
xs <- iterateNM (n-1) f x'
336+
return (x : xs)
337+
338+
unfoldrM :: Monad m => (b -> m (Maybe (a,b))) -> b -> m [a]
339+
unfoldrM step b0 = do
340+
r <- step b0
341+
case r of
342+
Nothing -> return []
343+
Just (a,b) -> do as <- unfoldrM step b
344+
return (a : as)
345+
346+
347+
limitUnfolds f (theirs, ours)
348+
| ours >= 0
349+
, Just (out, theirs') <- f theirs = Just (out, (theirs', ours - 1))
350+
| otherwise = Nothing

0 commit comments

Comments
 (0)