Implement unfoldrExactN & unfoldrExactNM

Fixes haskell#140

Data/Vector.hs

@@ -59,8 +59,8 @@ module Data.Vector (
   replicateM, generateM, iterateNM, create, createT,
 
   -- ** Unfolding
-  unfoldr, unfoldrN,
-  unfoldrM, unfoldrNM,
+  unfoldr, unfoldrN, unfoldrExactN,
+  unfoldrM, unfoldrNM, unfoldrExactNM,
   constructN, constructrN,
 
   -- ** Enumeration
@@ -691,6 +691,15 @@ unfoldrN :: Int -> (b -> Maybe (a, b)) -> b -> Vector a
 {-# INLINE unfoldrN #-}
 unfoldrN = G.unfoldrN
 
+-- | /O(n)/ Construct a vector with exactly @n@ elements by repeatedly applying
+-- the generator function to a seed. The generator function yields the
+-- next element and the new seed.
+--
+-- > unfoldrExactN 3 (\n -> (n,n-1)) 10 = <10,9,8>
+unfoldrExactN  :: Int -> (b -> (a, b)) -> b -> Vector a
+{-# INLINE unfoldrExactN #-}
+unfoldrExactN = G.unfoldrExactN
+
 -- | /O(n)/ Construct a vector by repeatedly applying the monadic
 -- generator function to a seed. The generator function yields 'Just'
 -- the next element and the new seed or 'Nothing' if there are no more
@@ -707,6 +716,13 @@ unfoldrNM :: (Monad m) => Int -> (b -> m (Maybe (a, b))) -> b -> m (Vector a)
 {-# INLINE unfoldrNM #-}
 unfoldrNM = G.unfoldrNM
 
+-- | /O(n)/ Construct a vector with exactly @n@ elements by repeatedly
+-- applying the monadic generator function to a seed. The generator
+-- function yields the next element and the new seed.
+unfoldrExactNM :: (Monad m) => Int -> (b -> m (a, b)) -> b -> m (Vector a)
+{-# INLINE unfoldrExactNM #-}
+unfoldrExactNM = G.unfoldrExactNM
+
 -- | /O(n)/ Construct a vector with @n@ elements by repeatedly applying the
 -- generator function to the already constructed part of the vector.
 --

Data/Vector/Fusion/Bundle.hs

@@ -55,7 +55,7 @@ module Data.Vector.Fusion.Bundle (
   and, or,
 
   -- * Unfolding
-  unfoldr, unfoldrN, iterateN,
+  unfoldr, unfoldrN, unfoldrExactN, iterateN,
 
   -- * Scans
   prescanl, prescanl',
@@ -437,6 +437,11 @@ unfoldrN :: Int -> (s -> Maybe (a, s)) -> s -> Bundle v a
 {-# INLINE unfoldrN #-}
 unfoldrN = M.unfoldrN
 
+-- | Unfold exactly @n@ elements
+unfoldrExactN :: Int -> (s -> (a, s)) -> s -> Bundle v a
+{-# INLINE unfoldrExactN #-}
+unfoldrExactN = M.unfoldrExactN
+
 -- | Apply function n-1 times to value. Zeroth element is original value.
 iterateN :: Int -> (a -> a) -> a -> Bundle v a
 {-# INLINE iterateN #-}

Data/Vector/Fusion/Bundle/Monadic.hs

@@ -59,6 +59,7 @@ module Data.Vector.Fusion.Bundle.Monadic (
   -- * Unfolding
   unfoldr, unfoldrM,
   unfoldrN, unfoldrNM,
+  unfoldrExactN, unfoldrExactNM,
   iterateN, iterateNM,
 
   -- * Scans
@@ -645,6 +646,16 @@ unfoldrNM :: Monad m => Int -> (s -> m (Maybe (a, s))) -> s -> Bundle m u a
 {-# INLINE_FUSED unfoldrNM #-}
 unfoldrNM n f s = fromStream (S.unfoldrNM n f s) (Max (delay_inline max n 0))
 
+-- | Unfold exactly @n@ elements
+unfoldrExactN :: Monad m => Int -> (s -> (a, s)) -> s -> Bundle m u a
+{-# INLINE_FUSED unfoldrExactN #-}
+unfoldrExactN n f = unfoldrExactNM n (return . f)
+
+-- | Unfold exactly @n@ elements with a monadic functions
+unfoldrExactNM :: Monad m => Int -> (s -> m (a, s)) -> s -> Bundle m u a
+{-# INLINE_FUSED unfoldrExactNM #-}
+unfoldrExactNM n f s = fromStream (S.unfoldrExactNM n f s) (Max (delay_inline max n 0))
+
 -- | Apply monadic function n times to value. Zeroth element is original value.
 iterateNM :: Monad m => Int -> (a -> m a) -> a -> Bundle m u a
 {-# INLINE_FUSED iterateNM #-}

Data/Vector/Fusion/Stream/Monadic.hs

@@ -56,6 +56,7 @@ module Data.Vector.Fusion.Stream.Monadic (
   -- * Unfolding
   unfoldr, unfoldrM,
   unfoldrN, unfoldrNM,
+  unfoldrExactN, unfoldrExactNM,
   iterateN, iterateNM,
 
   -- * Scans
@@ -1117,6 +1118,21 @@ unfoldrNM m f t = Stream step (t,m)
                                  Nothing     -> Done
                              ) (f s)
 
+-- | Unfold exactly @n@ elements
+unfoldrExactN :: Monad m => Int -> (s -> (a, s)) -> s -> Stream m a
+{-# INLINE_FUSED unfoldrExactN #-}
+unfoldrExactN n f = unfoldrExactNM n (return . f)
+
+-- | Unfold exactly @n@ elements with a monadic functions
+unfoldrExactNM :: Monad m => Int -> (s -> m (a, s)) -> s -> Stream m a
+{-# INLINE_FUSED unfoldrExactNM #-}
+unfoldrExactNM m f t = Stream step (t,m)
+  where
+    {-# INLINE_INNER step #-}
+    step (s,n) | n <= 0    = return Done
+               | otherwise = do (x,s') <- f s
+                                return $ Yield x (s',n-1)
+
 -- | Apply monadic function n times to value. Zeroth element is original value.
 iterateNM :: Monad m => Int -> (a -> m a) -> a -> Stream m a
 {-# INLINE_FUSED iterateNM #-}

Data/Vector/Generic.hs

@@ -42,8 +42,8 @@ module Data.Vector.Generic (
   replicateM, generateM, iterateNM, create, createT,
 
   -- ** Unfolding
-  unfoldr, unfoldrN,
-  unfoldrM, unfoldrNM,
+  unfoldr, unfoldrN, unfoldrExactN,
+  unfoldrM, unfoldrNM, unfoldrExactNM,
   constructN, constructrN,
 
   -- ** Enumeration
@@ -550,6 +550,15 @@ unfoldrN  :: Vector v a => Int -> (b -> Maybe (a, b)) -> b -> v a
 {-# INLINE unfoldrN #-}
 unfoldrN n f = unstream . Bundle.unfoldrN n f
 
+-- | /O(n)/ Construct a vector with exactly @n@ elements by repeatedly applying
+-- the generator function to a seed. The generator function yields the
+-- next element and the new seed.
+--
+-- > unfoldrExactN 3 (\n -> (n,n-1)) 10 = <10,9,8>
+unfoldrExactN  :: Vector v a => Int -> (b -> (a, b)) -> b -> v a
+{-# INLINE unfoldrExactN #-}
+unfoldrExactN n f = unstream . Bundle.unfoldrExactN n f
+
 -- | /O(n)/ Construct a vector by repeatedly applying the monadic
 -- generator function to a seed. The generator function yields 'Just'
 -- the next element and the new seed or 'Nothing' if there are no more
@@ -566,6 +575,13 @@ unfoldrNM :: (Monad m, Vector v a) => Int -> (b -> m (Maybe (a, b))) -> b -> m (
 {-# INLINE unfoldrNM #-}
 unfoldrNM n f = unstreamM . MBundle.unfoldrNM n f
 
+-- | /O(n)/ Construct a vector with exactly @n@ elements by repeatedly
+-- applying the monadic generator function to a seed. The generator
+-- function yields the next element and the new seed.
+unfoldrExactNM :: (Monad m, Vector v a) => Int -> (b -> m (a, b)) -> b -> m (v a)
+{-# INLINE unfoldrExactNM #-}
+unfoldrExactNM n f = unstreamM . MBundle.unfoldrExactNM n f
+
 -- | /O(n)/ Construct a vector with @n@ elements by repeatedly applying the
 -- generator function to the already constructed part of the vector.
 --

Data/Vector/Primitive.hs

@@ -49,8 +49,8 @@ module Data.Vector.Primitive (
   replicateM, generateM, iterateNM, create, createT,
 
   -- ** Unfolding
-  unfoldr, unfoldrN,
-  unfoldrM, unfoldrNM,
+  unfoldr, unfoldrN, unfoldrExactN,
+  unfoldrM, unfoldrNM, unfoldrExactNM,
   constructN, constructrN,
 
   -- ** Enumeration
@@ -534,6 +534,15 @@ unfoldrN :: Prim a => Int -> (b -> Maybe (a, b)) -> b -> Vector a
 {-# INLINE unfoldrN #-}
 unfoldrN = G.unfoldrN
 
+-- | /O(n)/ Construct a vector with exactly @n@ elements by repeatedly applying
+-- the generator function to a seed. The generator function yields the
+-- next element and the new seed.
+--
+-- > unfoldrExactN 3 (\n -> (n,n-1)) 10 = <10,9,8>
+unfoldrExactN :: (Prim a) => Int -> (b -> (a, b)) -> b -> Vector a
+{-# INLINE unfoldrExactN #-}
+unfoldrExactN = G.unfoldrExactN
+
 -- | /O(n)/ Construct a vector by repeatedly applying the monadic
 -- generator function to a seed. The generator function yields 'Just'
 -- the next element and the new seed or 'Nothing' if there are no more
@@ -550,6 +559,13 @@ unfoldrNM :: (Monad m, Prim a) => Int -> (b -> m (Maybe (a, b))) -> b -> m (Vect
 {-# INLINE unfoldrNM #-}
 unfoldrNM = G.unfoldrNM
 
+-- | /O(n)/ Construct a vector with exactly @n@ elements by repeatedly
+-- applying the monadic generator function to a seed. The generator
+-- function yields the next element and the new seed.
+unfoldrExactNM :: (Monad m, Prim a) => Int -> (b -> m (a, b)) -> b -> m (Vector a)
+{-# INLINE unfoldrExactNM #-}
+unfoldrExactNM = G.unfoldrExactNM
+
 -- | /O(n)/ Construct a vector with @n@ elements by repeatedly applying the
 -- generator function to the already constructed part of the vector.
 --

Data/Vector/Storable.hs

@@ -46,8 +46,8 @@ module Data.Vector.Storable (
   replicateM, generateM, iterateNM, create, createT,
 
   -- ** Unfolding
-  unfoldr, unfoldrN,
-  unfoldrM, unfoldrNM,
+  unfoldr, unfoldrN, unfoldrExactN,
+  unfoldrM, unfoldrNM, unfoldrExactNM,
   constructN, constructrN,
 
   -- ** Enumeration
@@ -561,6 +561,15 @@ unfoldrN :: Storable a => Int -> (b -> Maybe (a, b)) -> b -> Vector a
 {-# INLINE unfoldrN #-}
 unfoldrN = G.unfoldrN
 
+-- | /O(n)/ Construct a vector with exactly @n@ elements by repeatedly applying
+-- the generator function to a seed. The generator function yields the
+-- next element and the new seed.
+--
+-- > unfoldrExactN 3 (\n -> (n,n-1)) 10 = <10,9,8>
+unfoldrExactN :: (Storable a) => Int -> (b -> (a, b)) -> b -> Vector a
+{-# INLINE unfoldrExactN #-}
+unfoldrExactN = G.unfoldrExactN
+
 -- | /O(n)/ Construct a vector by repeatedly applying the monadic
 -- generator function to a seed. The generator function yields 'Just'
 -- the next element and the new seed or 'Nothing' if there are no more
@@ -577,6 +586,13 @@ unfoldrNM :: (Monad m, Storable a) => Int -> (b -> m (Maybe (a, b))) -> b -> m (
 {-# INLINE unfoldrNM #-}
 unfoldrNM = G.unfoldrNM
 
+-- | /O(n)/ Construct a vector with exactly @n@ elements by repeatedly
+-- applying the monadic generator function to a seed. The generator
+-- function yields the next element and the new seed.
+unfoldrExactNM :: (Monad m, Storable a) => Int -> (b -> m (a, b)) -> b -> m (Vector a)
+{-# INLINE unfoldrExactNM #-}
+unfoldrExactNM = G.unfoldrExactNM
+
 -- | /O(n)/ Construct a vector with @n@ elements by repeatedly applying the
 -- generator function to the already constructed part of the vector.
 --

Data/Vector/Unboxed.hs

@@ -65,8 +65,8 @@ module Data.Vector.Unboxed (
   replicateM, generateM, iterateNM, create, createT,
 
   -- ** Unfolding
-  unfoldr, unfoldrN,
-  unfoldrM, unfoldrNM,
+  unfoldr, unfoldrN, unfoldrExactN,
+  unfoldrM, unfoldrNM, unfoldrExactNM,
   constructN, constructrN,
 
   -- ** Enumeration
@@ -496,6 +496,15 @@ unfoldrN :: Unbox a => Int -> (b -> Maybe (a, b)) -> b -> Vector a
 {-# INLINE unfoldrN #-}
 unfoldrN = G.unfoldrN
 
+-- | /O(n)/ Construct a vector with exactly @n@ elements by repeatedly applying
+-- the generator function to a seed. The generator function yields the
+-- next element and the new seed.
+--
+-- > unfoldrExactN 3 (\n -> (n,n-1)) 10 = <10,9,8>
+unfoldrExactN  :: Unbox a => Int -> (b -> (a, b)) -> b -> Vector a
+{-# INLINE unfoldrExactN #-}
+unfoldrExactN = G.unfoldrExactN
+
 -- | /O(n)/ Construct a vector by repeatedly applying the monadic
 -- generator function to a seed. The generator function yields 'Just'
 -- the next element and the new seed or 'Nothing' if there are no more
@@ -512,6 +521,13 @@ unfoldrNM :: (Monad m, Unbox a) => Int -> (b -> m (Maybe (a, b))) -> b -> m (Vec
 {-# INLINE unfoldrNM #-}
 unfoldrNM = G.unfoldrNM
 
+-- | /O(n)/ Construct a vector with exactly @n@ elements by repeatedly
+-- applying the monadic generator function to a seed. The generator
+-- function yields the next element and the new seed.
+unfoldrExactNM :: (Monad m, Unbox a) => Int -> (b -> m (a, b)) -> b -> m (Vector a)
+{-# INLINE unfoldrExactNM #-}
+unfoldrExactNM = G.unfoldrExactNM
+
 -- | /O(n)/ Construct a vector with @n@ elements by repeatedly applying the
 -- generator function to the already constructed part of the vector.
 --

tests/Tests/Vector/Property.hs

@@ -25,6 +25,7 @@ module Tests.Vector.Property
 import Boilerplater
 import Utilities as Util hiding (limitUnfolds)
 
+import Control.Monad
 import Data.Functor.Identity
 import qualified Data.Traversable as T (Traversable(..))
 import Data.Foldable (Foldable(foldMap))
@@ -136,7 +137,8 @@ testPolymorphicFunctions _ = $(testProperties [
         {- 'prop_replicateM, 'prop_generateM, 'prop_create, -}
 
         -- Unfolding
-        'prop_unfoldr, 'prop_unfoldrN, 'prop_unfoldrM, 'prop_unfoldrNM,
+        'prop_unfoldr, 'prop_unfoldrN, 'prop_unfoldrExactN,
+        'prop_unfoldrM, 'prop_unfoldrNM, 'prop_unfoldrExactNM,
         'prop_constructN, 'prop_constructrN,
 
         -- Enumeration? (FIXME?)
@@ -468,11 +470,15 @@ testPolymorphicFunctions _ = $(testProperties [
            `eq` (\n f a -> unfoldr (limitUnfolds f) (a, n))
     prop_unfoldrN :: P (Int -> (Int -> Maybe (a,Int)) -> Int -> v a)
          = V.unfoldrN `eq` (\n f a -> unfoldr (limitUnfolds f) (a, n))
+    prop_unfoldrExactN :: P (Int -> (Int -> (a,Int)) -> Int -> v a)
+         = V.unfoldrExactN `eq` (\n f a -> unfoldr (limitUnfolds (Just . f)) (a, n))
     prop_unfoldrM :: P (Int -> (Int -> Writer [Int] (Maybe (a,Int))) -> Int -> Writer [Int] (v a))
          = (\n f a -> V.unfoldrM (limitUnfoldsM f) (a,n))
            `eq` (\n f a -> Util.unfoldrM (limitUnfoldsM f) (a, n))
     prop_unfoldrNM :: P (Int -> (Int -> Writer [Int] (Maybe (a,Int))) -> Int -> Writer [Int] (v a))
          = V.unfoldrNM `eq` (\n f a -> Util.unfoldrM (limitUnfoldsM f) (a, n))
+    prop_unfoldrExactNM :: P (Int -> (Int -> Writer [Int] (a,Int)) -> Int -> Writer [Int] (v a))
+         = V.unfoldrExactNM `eq` (\n f a -> Util.unfoldrM (limitUnfoldsM (liftM Just . f)) (a, n))
 
     prop_constructN  = \f -> forAll (choose (0,20)) $ \n -> unP prop n f
       where