Add constructN and constructrN
authorRoman Leshchinskiy <rl@cse.unsw.edu.au>
Fri, 19 Aug 2011 23:40:52 +0000 (23:40 +0000)
committerRoman Leshchinskiy <rl@cse.unsw.edu.au>
Fri, 19 Aug 2011 23:40:52 +0000 (23:40 +0000)
Data/Vector.hs
Data/Vector/Generic.hs
Data/Vector/Primitive.hs
Data/Vector/Storable.hs
Data/Vector/Unboxed.hs

index 5f0d0b5..6b3cc98 100644 (file)
@@ -53,6 +53,7 @@ module Data.Vector (
 
   -- ** Unfolding
   unfoldr, unfoldrN,
+  constructN, constructrN,
 
   -- ** Enumeration
   enumFromN, enumFromStepN, enumFromTo, enumFromThenTo,
@@ -541,6 +542,25 @@ unfoldrN :: Int -> (b -> Maybe (a, b)) -> b -> Vector a
 {-# INLINE unfoldrN #-}
 unfoldrN = G.unfoldrN
 
+-- | /O(n)/ Construct a vector with @n@ elements by repeatedly applying the
+-- generator function to the already constructed part of the vector.
+--
+-- > constructN 3 f = let a = f <> ; b = f <a> ; c = f <a,b> in f <a,b,c>
+--
+constructN :: Int -> (Vector a -> a) -> Vector a
+{-# INLINE constructN #-}
+constructN = G.constructN
+
+-- | /O(n)/ Construct a vector with @n@ elements from right to left by
+-- repeatedly applying the generator function to the already constructed part
+-- of the vector.
+--
+-- > constructrN 3 f = let a = f <> ; b = f<a> ; c = f <b,a> in f <c,b,a>
+--
+constructrN :: Int -> (Vector a -> a) -> Vector a
+{-# INLINE constructrN #-}
+constructrN = G.constructrN
+
 -- Enumeration
 -- -----------
 
index 8d559a3..67746cc 100644 (file)
@@ -1,5 +1,5 @@
 {-# LANGUAGE Rank2Types, MultiParamTypeClasses, FlexibleContexts,
-             TypeFamilies, ScopedTypeVariables #-}
+             TypeFamilies, ScopedTypeVariables, BangPatterns #-}
 -- |
 -- Module      : Data.Vector.Generic
 -- Copyright   : (c) Roman Leshchinskiy 2008-2010
@@ -43,6 +43,7 @@ module Data.Vector.Generic (
 
   -- ** Unfolding
   unfoldr, unfoldrN,
+  constructN, constructrN,
 
   -- ** Enumeration
   enumFromN, enumFromStepN, enumFromTo, enumFromThenTo,
@@ -235,8 +236,8 @@ infixl 9 !
 -- | O(1) Indexing
 (!) :: Vector v a => v a -> Int -> a
 {-# INLINE_STREAM (!) #-}
-v ! i = BOUNDS_CHECK(checkIndex) "(!)" i (length v)
-      $ unId (basicUnsafeIndexM v i)
+(!) v i = BOUNDS_CHECK(checkIndex) "(!)" i (length v)
+        $ unId (basicUnsafeIndexM v i)
 
 infixl 9 !?
 -- | O(1) Safe indexing
@@ -547,6 +548,64 @@ unfoldrN  :: Vector v a => Int -> (b -> Maybe (a, b)) -> b -> v a
 {-# INLINE unfoldrN #-}
 unfoldrN n f = unstream . Stream.unfoldrN n f
 
+-- | /O(n)/ Construct a vector with @n@ elements by repeatedly applying the
+-- generator function to the already constructed part of the vector.
+--
+-- > constructN 3 f = let a = f <> ; b = f <a> ; c = f <a,b> in f <a,b,c>
+--
+constructN :: forall v a. Vector v a => Int -> (v a -> a) -> v a
+{-# INLINE constructN #-}
+-- NOTE: We *CANNOT* wrap this in New and then fuse because the elements
+-- might contain references to the immutable vector!
+constructN !n f = runST (
+  do
+    v  <- M.new n
+    v' <- unsafeFreeze v
+    fill v' 0
+  )
+  where
+    fill :: forall s. v a -> Int -> ST s (v a)
+    fill !v i | i < n = let x = f (unsafeTake i v)
+                        in
+                        elemseq v x $
+                        do
+                          v'  <- unsafeThaw v
+                          M.unsafeWrite v' i x
+                          v'' <- unsafeFreeze v'
+                          fill v'' (i+1)
+
+    fill v i = return v
+
+-- | /O(n)/ Construct a vector with @n@ elements from right to left by
+-- repeatedly applying the generator function to the already constructed part
+-- of the vector.
+--
+-- > constructrN 3 f = let a = f <> ; b = f<a> ; c = f <b,a> in f <c,b,a>
+--
+constructrN :: forall v a. Vector v a => Int -> (v a -> a) -> v a
+{-# INLINE constructrN #-}
+-- NOTE: We *CANNOT* wrap this in New and then fuse because the elements
+-- might contain references to the immutable vector!
+constructrN !n f = runST (
+  do
+    v  <- n `seq` M.new n
+    v' <- unsafeFreeze v
+    fill v' 0
+  )
+  where
+    fill :: forall s. v a -> Int -> ST s (v a)
+    fill !v i | i < n = let x = f (unsafeSlice (n-i) i v)
+                        in
+                        elemseq v x $
+                        do
+                          v'  <- unsafeThaw v
+                          M.unsafeWrite v' (n-i-1) x
+                          v'' <- unsafeFreeze v'
+                          fill v'' (i+1)
+
+    fill v i = return v
+
+
 -- Enumeration
 -- -----------
 
index 6a42da8..56ae9c4 100644 (file)
@@ -46,6 +46,7 @@ module Data.Vector.Primitive (
 
   -- ** Unfolding
   unfoldr, unfoldrN,
+  constructN, constructrN,
 
   -- ** Enumeration
   enumFromN, enumFromStepN, enumFromTo, enumFromThenTo,
@@ -481,6 +482,25 @@ unfoldrN :: Prim a => Int -> (b -> Maybe (a, b)) -> b -> Vector a
 {-# INLINE unfoldrN #-}
 unfoldrN = G.unfoldrN
 
+-- | /O(n)/ Construct a vector with @n@ elements by repeatedly applying the
+-- generator function to the already constructed part of the vector.
+--
+-- > constructN 3 f = let a = f <> ; b = f <a> ; c = f <a,b> in f <a,b,c>
+--
+constructN :: Prim a => Int -> (Vector a -> a) -> Vector a
+{-# INLINE constructN #-}
+constructN = G.constructN
+
+-- | /O(n)/ Construct a vector with @n@ elements from right to left by
+-- repeatedly applying the generator function to the already constructed part
+-- of the vector.
+--
+-- > constructrN 3 f = let a = f <> ; b = f<a> ; c = f <b,a> in f <c,b,a>
+--
+constructrN :: Prim a => Int -> (Vector a -> a) -> Vector a
+{-# INLINE constructrN #-}
+constructrN = G.constructrN
+
 -- Enumeration
 -- -----------
 
index bd04166..a0985cc 100644 (file)
@@ -43,6 +43,7 @@ module Data.Vector.Storable (
 
   -- ** Unfolding
   unfoldr, unfoldrN,
+  constructN, constructrN,
 
   -- ** Enumeration
   enumFromN, enumFromStepN, enumFromTo, enumFromThenTo,
@@ -489,6 +490,25 @@ unfoldrN :: Storable a => Int -> (b -> Maybe (a, b)) -> b -> Vector a
 {-# INLINE unfoldrN #-}
 unfoldrN = G.unfoldrN
 
+-- | /O(n)/ Construct a vector with @n@ elements by repeatedly applying the
+-- generator function to the already constructed part of the vector.
+--
+-- > constructN 3 f = let a = f <> ; b = f <a> ; c = f <a,b> in f <a,b,c>
+--
+constructN :: Storable a => Int -> (Vector a -> a) -> Vector a
+{-# INLINE constructN #-}
+constructN = G.constructN
+
+-- | /O(n)/ Construct a vector with @n@ elements from right to left by
+-- repeatedly applying the generator function to the already constructed part
+-- of the vector.
+--
+-- > constructrN 3 f = let a = f <> ; b = f<a> ; c = f <b,a> in f <c,b,a>
+--
+constructrN :: Storable a => Int -> (Vector a -> a) -> Vector a
+{-# INLINE constructrN #-}
+constructrN = G.constructrN
+
 -- Enumeration
 -- -----------
 
index aca9109..3d98bed 100644 (file)
@@ -66,6 +66,7 @@ module Data.Vector.Unboxed (
 
   -- ** Unfolding
   unfoldr, unfoldrN,
+  constructN, constructrN,
 
   -- ** Enumeration
   enumFromN, enumFromStepN, enumFromTo, enumFromThenTo,
@@ -460,6 +461,25 @@ unfoldrN :: Unbox a => Int -> (b -> Maybe (a, b)) -> b -> Vector a
 {-# INLINE unfoldrN #-}
 unfoldrN = G.unfoldrN
 
+-- | /O(n)/ Construct a vector with @n@ elements by repeatedly applying the
+-- generator function to the already constructed part of the vector.
+--
+-- > constructN 3 f = let a = f <> ; b = f <a> ; c = f <a,b> in f <a,b,c>
+--
+constructN :: Unbox a => Int -> (Vector a -> a) -> Vector a
+{-# INLINE constructN #-}
+constructN = G.constructN
+
+-- | /O(n)/ Construct a vector with @n@ elements from right to left by
+-- repeatedly applying the generator function to the already constructed part
+-- of the vector.
+--
+-- > constructrN 3 f = let a = f <> ; b = f<a> ; c = f <b,a> in f <c,b,a>
+--
+constructrN :: Unbox a => Int -> (Vector a -> a) -> Vector a
+{-# INLINE constructrN #-}
+constructrN = G.constructrN
+
 -- Enumeration
 -- -----------