Add unfoldrN
authorRoman Leshchinskiy <rl@cse.unsw.edu.au>
Sun, 21 Feb 2010 02:28:11 +0000 (02:28 +0000)
committerRoman Leshchinskiy <rl@cse.unsw.edu.au>
Sun, 21 Feb 2010 02:28:11 +0000 (02:28 +0000)
Data/Vector.hs
Data/Vector/Fusion/Stream.hs
Data/Vector/Fusion/Stream/Monadic.hs
Data/Vector/Generic.hs
Data/Vector/Primitive.hs
Data/Vector/Storable.hs
Data/Vector/Unboxed.hs

index 87c8e2c..49349a6 100644 (file)
@@ -100,7 +100,7 @@ module Data.Vector (
   minIndex, minIndexBy, maxIndex, maxIndexBy,
 
   -- * Unfolding
-  unfoldr,
+  unfoldr, unfoldrN,
 
   -- * Scans
   prescanl, prescanl',
@@ -848,6 +848,11 @@ unfoldr :: (b -> Maybe (a, b)) -> b -> Vector a
 {-# INLINE unfoldr #-}
 unfoldr = G.unfoldr
 
+-- | Unfold at most @n@ elements
+unfoldrN :: Int -> (b -> Maybe (a, b)) -> b -> Vector a
+{-# INLINE unfoldrN #-}
+unfoldrN = G.unfoldrN
+
 -- Scans
 -- -----
 
index f45baa6..db25db4 100644 (file)
@@ -55,7 +55,7 @@ module Data.Vector.Fusion.Stream (
   and, or,
 
   -- * Unfolding
-  unfoldr,
+  unfoldr, unfoldrN,
 
   -- * Scans
   prescanl, prescanl',
@@ -411,6 +411,11 @@ unfoldr :: (s -> Maybe (a, s)) -> s -> Stream a
 {-# INLINE unfoldr #-}
 unfoldr = M.unfoldr
 
+-- | Unfold at most @n@ elements
+unfoldrN :: Int -> (s -> Maybe (a, s)) -> s -> Stream a
+{-# INLINE unfoldrN #-}
+unfoldrN = M.unfoldrN
+
 -- Scans
 -- -----
 
index 4a95e99..b19eed4 100644 (file)
@@ -55,6 +55,7 @@ module Data.Vector.Fusion.Stream.Monadic (
 
   -- * Unfolding
   unfoldr, unfoldrM,
+  unfoldrN, unfoldrNM,
 
   -- * Scans
   prescanl, prescanlM, prescanl', prescanlM',
@@ -937,6 +938,24 @@ unfoldrM f s = Stream step s Unknown
                  Nothing      -> Done
              ) (f s)
 
+-- | Unfold at most @n@ elements
+unfoldrN :: Monad m => Int -> (s -> Maybe (a, s)) -> s -> Stream m a
+{-# INLINE_STREAM unfoldrN #-}
+unfoldrN n f = unfoldrNM n (return . f)
+
+-- | Unfold at most @n@ elements with a monadic functions
+unfoldrNM :: Monad m => Int -> (s -> m (Maybe (a, s))) -> s -> Stream m a
+{-# INLINE_STREAM unfoldrNM #-}
+unfoldrNM n f s = Stream step (s,n) (Max (delay_inline max n 0))
+  where
+    {-# INLINE_INNER step #-}
+    step (s,n) | n <= 0    = return Done
+               | otherwise = liftM (\r ->
+                               case r of
+                                 Just (x,s') -> Yield x (s',n-1)
+                                 Nothing     -> Done
+                             ) (f s)
+
 -- Scans
 -- -----
 
index 09e5149..0ba2c71 100644 (file)
@@ -69,7 +69,7 @@ module Data.Vector.Generic (
   minIndex, minIndexBy, maxIndex, maxIndexBy,
 
   -- * Unfolding
-  unfoldr,
+  unfoldr, unfoldrN,
 
   -- * Scans
   prescanl, prescanl',
@@ -1136,10 +1136,16 @@ minIndexBy cmp = fst . Stream.foldl1' imin . Stream.indexed . stream
 -- Unfolding
 -- ---------
 
+-- | Unfold
 unfoldr :: Vector v a => (b -> Maybe (a, b)) -> b -> v a
 {-# INLINE unfoldr #-}
 unfoldr f = unstream . Stream.unfoldr f
 
+-- | Unfoldr at most @n@ elements.
+unfoldrN  :: Vector v a => Int -> (b -> Maybe (a, b)) -> b -> v a
+{-# INLINE unfoldrN #-}
+unfoldrN n f = unstream . Stream.unfoldrN n f
+
 -- Scans
 -- -----
 
index f0741f9..e7cbf04 100644 (file)
@@ -61,7 +61,7 @@ module Data.Vector.Primitive (
   minIndex, minIndexBy, maxIndex, maxIndexBy,
 
   -- * Unfolding
-  unfoldr,
+  unfoldr, unfoldrN,
 
   -- * Scans
   prescanl, prescanl',
@@ -675,10 +675,27 @@ minIndexBy = G.minIndexBy
 -- Unfolding
 -- ---------
 
+-- | The 'unfoldr' function is a \`dual\' to 'foldr': while 'foldr'
+-- reduces a vector to a summary value, 'unfoldr' builds a list from
+-- a seed value.  The function takes the element and returns 'Nothing'
+-- if it is done generating the vector or returns 'Just' @(a,b)@, in which
+-- case, @a@ is a prepended to the vector and @b@ is used as the next
+-- element in a recursive call.
+--
+-- A simple use of unfoldr:
+--
+-- > unfoldr (\b -> if b == 0 then Nothing else Just (b, b-1)) 10
+-- >  [10,9,8,7,6,5,4,3,2,1]
+--
 unfoldr :: Prim a => (b -> Maybe (a, b)) -> b -> Vector a
 {-# INLINE unfoldr #-}
 unfoldr = G.unfoldr
 
+-- | Unfold at most @n@ elements
+unfoldrN :: Prim a => Int -> (b -> Maybe (a, b)) -> b -> Vector a
+{-# INLINE unfoldrN #-}
+unfoldrN = G.unfoldrN
+
 -- Scans
 -- -----
 
index c2d9d89..ca3b595 100644 (file)
@@ -61,7 +61,7 @@ module Data.Vector.Storable (
   minIndex, minIndexBy, maxIndex, maxIndexBy,
 
   -- * Unfolding
-  unfoldr,
+  unfoldr, unfoldrN,
 
   -- * Scans
   prescanl, prescanl',
@@ -723,10 +723,27 @@ minIndexBy = G.minIndexBy
 -- Unfolding
 -- ---------
 
+-- | The 'unfoldr' function is a \`dual\' to 'foldr': while 'foldr'
+-- reduces a vector to a summary value, 'unfoldr' builds a list from
+-- a seed value.  The function takes the element and returns 'Nothing'
+-- if it is done generating the vector or returns 'Just' @(a,b)@, in which
+-- case, @a@ is a prepended to the vector and @b@ is used as the next
+-- element in a recursive call.
+--
+-- A simple use of unfoldr:
+--
+-- > unfoldr (\b -> if b == 0 then Nothing else Just (b, b-1)) 10
+-- >  [10,9,8,7,6,5,4,3,2,1]
+--
 unfoldr :: Storable a => (b -> Maybe (a, b)) -> b -> Vector a
 {-# INLINE unfoldr #-}
 unfoldr = G.unfoldr
 
+-- | Unfold at most @n@ elements
+unfoldrN :: Storable a => Int -> (b -> Maybe (a, b)) -> b -> Vector a
+{-# INLINE unfoldrN #-}
+unfoldrN = G.unfoldrN
+
 -- Scans
 -- -----
 
index 9e89f42..069e1b0 100644 (file)
@@ -63,7 +63,7 @@ module Data.Vector.Unboxed (
   minIndex, minIndexBy, maxIndex, maxIndexBy,
 
   -- * Unfolding
-  unfoldr,
+  unfoldr, unfoldrN,
 
   -- * Scans
   prescanl, prescanl',
@@ -676,10 +676,27 @@ minIndexBy = G.minIndexBy
 -- Unfolding
 -- ---------
 
+-- | The 'unfoldr' function is a \`dual\' to 'foldr': while 'foldr'
+-- reduces a vector to a summary value, 'unfoldr' builds a list from
+-- a seed value.  The function takes the element and returns 'Nothing'
+-- if it is done generating the vector or returns 'Just' @(a,b)@, in which
+-- case, @a@ is a prepended to the vector and @b@ is used as the next
+-- element in a recursive call.
+--
+-- A simple use of unfoldr:
+--
+-- > unfoldr (\b -> if b == 0 then Nothing else Just (b, b-1)) 10
+-- >  [10,9,8,7,6,5,4,3,2,1]
+--
 unfoldr :: Unbox a => (b -> Maybe (a, b)) -> b -> Vector a
 {-# INLINE unfoldr #-}
 unfoldr = G.unfoldr
 
+-- | Unfold at most @n@ elements
+unfoldrN :: Unbox a => Int -> (b -> Maybe (a, b)) -> b -> Vector a
+{-# INLINE unfoldrN #-}
+unfoldrN = G.unfoldrN
+
 -- Scans
 -- -----