Added splitAt functions (contributed by Bas van Dijk)
[darcs-mirrors/vector.git] / Data / Vector / Generic.hs
index 9c9fbe0..3a7836e 100644 (file)
@@ -22,7 +22,7 @@ module Data.Vector.Generic (
   length, null,
 
   -- ** Indexing
-  (!), head, last,
+  (!), (!?), head, last,
   unsafeIndex, unsafeHead, unsafeLast,
 
   -- ** Monadic indexing
@@ -30,7 +30,7 @@ module Data.Vector.Generic (
   unsafeIndexM, unsafeHeadM, unsafeLastM,
 
   -- ** Extracting subvectors (slicing)
-  slice, init, tail, take, drop,
+  slice, init, tail, take, drop, splitAt,
   unsafeSlice, unsafeInit, unsafeTail, unsafeTake, unsafeDrop,
 
   -- * Construction
@@ -48,7 +48,7 @@ module Data.Vector.Generic (
   enumFromN, enumFromStepN, enumFromTo, enumFromThenTo,
 
   -- ** Concatenation
-  cons, snoc, (++),
+  cons, snoc, (++), concat,
 
   -- ** Restricting memory usage
   force,
@@ -126,8 +126,11 @@ module Data.Vector.Generic (
   -- ** Lists
   toList, fromList, fromListN,
 
+  -- ** Different vector types
+  convert,
+
   -- ** Mutable vectors
-  thaw, thawMany, copy, unsafeCopy,
+  freeze, thaw, copy, unsafeFreeze, unsafeThaw, unsafeCopy,
 
   -- * Fusion support
 
@@ -165,10 +168,10 @@ import Control.Monad.Primitive
 import qualified Control.Monad as Monad
 import qualified Data.List as List
 import Prelude hiding ( length, null,
-                        replicate, (++),
+                        replicate, (++), concat,
                         head, last,
-                        init, tail, take, drop, reverse,
-                        map, concatMap,
+                        init, tail, take, drop, splitAt, reverse,
+                        map, concat, concatMap,
                         zipWith, zipWith3, zip, zip3, unzip, unzip3,
                         filter, takeWhile, dropWhile, span, break,
                         elem, notElem,
@@ -179,10 +182,18 @@ import Prelude hiding ( length, null,
                         mapM, mapM_ )
 
 import Data.Typeable ( Typeable1, gcast1 )
-import Data.Data ( Data, DataType, mkNorepType )
 
 #include "vector.h"
 
+import Data.Data ( Data, DataType )
+#if MIN_VERSION_base(4,2,0)
+import Data.Data ( mkNoRepType )
+#else
+import Data.Data ( mkNorepType )
+mkNoRepType :: String -> DataType
+mkNoRepType = mkNorepType
+#endif
+
 -- Length information
 -- ------------------
 
@@ -219,6 +230,12 @@ null v = basicLength v == 0
 v ! i = BOUNDS_CHECK(checkIndex) "(!)" i (length v)
       $ unId (basicUnsafeIndexM v i)
 
+-- | O(1) Safe indexing
+(!?) :: Vector v a => v a -> Int -> Maybe a
+{-# INLINE_STREAM (!?) #-}
+v !? i | i < 0 || i >= length v = Nothing
+       | otherwise              = Just $ unsafeIndex v i
+
 -- | /O(1)/ First element
 head :: Vector v a => v a -> a
 {-# INLINE_STREAM head #-}
@@ -372,7 +389,7 @@ tail :: Vector v a => v a -> v a
 {-# INLINE_STREAM tail #-}
 tail v = slice 1 (length v - 1) v
 
--- | /O(1)/ Yield at the first @n@ elements without copying. The vector may
+-- | /O(1)/ Yield the first @n@ elements without copying. The vector may
 -- contain less than @n@ elements in which case it is returned unchanged.
 take :: Vector v a => Int -> v a -> v a
 {-# INLINE_STREAM take #-}
@@ -388,6 +405,20 @@ drop n v = unsafeSlice (delay_inline min n' len)
   where n' = max n 0
         len = length v
 
+-- | /O(1)/ Yield the first @n@ elements paired with the remainder without copying.
+--
+-- Note that @'splitAt' n v@ is equivalent to @('take' n v, 'drop' n v)@
+-- but slightly more efficient.
+{-# INLINE_STREAM splitAt #-}
+splitAt :: Vector v a => Int -> v a -> (v a, v a)
+splitAt n v = ( unsafeSlice 0 m v
+              , unsafeSlice m (delay_inline max 0 (len - n')) v
+              )
+    where
+      m   = delay_inline min n' len
+      n'  = max n 0
+      len = length v
+
 -- | /O(1)/ Yield a slice of the vector without copying. The vector must
 -- contain at least @i+n@ elements but this is not checked.
 unsafeSlice :: Vector v a => Int   -- ^ @i@ starting index
@@ -561,6 +592,25 @@ infixr 5 ++
 {-# INLINE (++) #-}
 v ++ w = unstream (stream v Stream.++ stream w)
 
+-- | /O(n)/ Concatenate all vectors in the list
+concat :: Vector v a => [v a] -> v a
+{-# INLINE concat #-}
+-- concat vs = create (thawMany vs)
+concat vs = unstream (Stream.flatten mk step (Exact n) (Stream.fromList vs))
+  where
+    n = List.foldl' (\k v -> k + length v) 0 vs
+
+    {-# INLINE_INNER step #-}
+    step (v,i,k)
+      | i < k = case unsafeIndexM v i of
+                  Box x -> Stream.Yield x (v,i+1,k)
+      | otherwise = Stream.Done
+
+    {-# INLINE mk #-}
+    mk v = let k = length v
+           in
+           k `seq` (v,0,k)
+
 -- Monadic initialisation
 -- ----------------------
 
@@ -770,19 +820,37 @@ backpermute :: (Vector v a, Vector v Int)
 -- does not retain references to the original one even if it is lazy in its
 -- elements. This would not be the case if we simply used map (v!)
 backpermute v is = seq v
+                 $ seq n
                  $ unstream
                  $ Stream.unbox
-                 $ Stream.map (indexM v)
+                 $ Stream.map index
                  $ stream is
+  where
+    n = length v
+
+    {-# INLINE index #-}
+    -- NOTE: we do it this way to avoid triggering LiberateCase on n in
+    -- polymorphic code
+    index i = BOUNDS_CHECK(checkIndex) "backpermute" i n
+            $ basicUnsafeIndexM v i
 
 -- | Same as 'backpermute' but without bounds checking.
 unsafeBackpermute :: (Vector v a, Vector v Int) => v a -> v Int -> v a
 {-# INLINE unsafeBackpermute #-}
 unsafeBackpermute v is = seq v
+                       $ seq n
                        $ unstream
                        $ Stream.unbox
-                       $ Stream.map (unsafeIndexM v)
+                       $ Stream.map index
                        $ stream is
+  where
+    n = length v
+
+    {-# INLINE index #-}
+    -- NOTE: we do it this way to avoid triggering LiberateCase on n in
+    -- polymorphic code
+    index i = UNSAFE_CHECK(checkIndex) "unsafeBackpermute" i n
+            $ basicUnsafeIndexM v i
 
 -- Safe destructive updates
 -- ------------------------
@@ -823,7 +891,9 @@ imap f = unstream . inplace (MStream.map (uncurry f) . MStream.indexed)
 -- | Map a function over a vector and concatenate the results.
 concatMap :: (Vector v a, Vector v b) => (a -> v b) -> v a -> v b
 {-# INLINE concatMap #-}
-concatMap f = unstream . Stream.concatMap (stream . f) . stream
+-- NOTE: We can't fuse concatMap anyway so don't pretend we do.
+-- concatMap f = unstream . Stream.concatMap (stream . f) . stream
+concatMap f = concat . Stream.toList . Stream.map f . stream
 
 -- Monadic mapping
 -- ---------------
@@ -1531,9 +1601,31 @@ fromListN :: Vector v a => Int -> [a] -> v a
 {-# INLINE fromListN #-}
 fromListN n = unstream . Stream.fromListN n
 
+-- Conversions - Immutable vectors
+-- -------------------------------
+
+-- | /O(n)/ Convert different vector types
+convert :: (Vector v a, Vector w a) => v a -> w a
+{-# INLINE convert #-}
+convert = unstream . stream
+
 -- Conversions - Mutable vectors
 -- -----------------------------
 
+-- | /O(1)/ Unsafe convert a mutable vector to an immutable one without
+-- copying. The mutable vector may not be used after this operation.
+unsafeFreeze
+  :: (PrimMonad m, Vector v a) => Mutable v (PrimState m) a -> m (v a)
+{-# INLINE unsafeFreeze #-}
+unsafeFreeze = basicUnsafeFreeze
+
+
+-- | /O(1)/ Unsafely convert an immutable vector to a mutable one without
+-- copying. The immutable vector may not be used after this operation.
+unsafeThaw :: (PrimMonad m, Vector v a) => v a -> m (Mutable v (PrimState m) a)
+{-# INLINE unsafeThaw #-}
+unsafeThaw = basicUnsafeThaw
+
 -- | /O(n)/ Yield a mutable copy of the immutable vector.
 thaw :: (PrimMonad m, Vector v a) => v a -> m (Mutable v (PrimState m) a)
 {-# INLINE_STREAM thaw #-}
@@ -1542,6 +1634,12 @@ thaw v = do
            unsafeCopy mv v
            return mv
 
+-- | /O(n)/ Yield an immutable copy of the mutable vector.
+freeze :: (PrimMonad m, Vector v a) => Mutable v (PrimState m) a -> m (v a)
+{-# INLINE freeze #-}
+freeze mv = unsafeFreeze =<< M.clone mv
+
+{-
 -- | /O(n)/ Yield a mutable vector containing copies of each vector in the
 -- list.
 thawMany :: (PrimMonad m, Vector v a) => [v a] -> m (Mutable v (PrimState m) a)
@@ -1562,6 +1660,7 @@ thawMany vs = do
           let n = length v
           unsafeCopy (M.unsafeTake n mv) v
           thaw_loop (M.unsafeDrop n mv) vs
+-}
 
 -- | /O(n)/ Copy an immutable vector into a mutable one. The two vectors must
 -- have the same length.
@@ -1587,7 +1686,7 @@ unsafeCopy dst src = UNSAFE_CHECK(check) "unsafeCopy" "length mismatch"
 -- | /O(1)/ Convert a vector to a 'Stream'
 stream :: Vector v a => v a -> Stream a
 {-# INLINE_STREAM stream #-}
-stream v = v `seq` (Stream.unfoldr get 0 `Stream.sized` Exact n)
+stream v = v `seq` n `seq` (Stream.unfoldr get 0 `Stream.sized` Exact n)
   where
     n = length v
 
@@ -1626,7 +1725,7 @@ unstream s = new (New.unstream s)
 -- | /O(1)/ Convert a vector to a 'Stream', proceeding from right to left
 streamR :: Vector v a => v a -> Stream a
 {-# INLINE_STREAM streamR #-}
-streamR v = v `seq` (Stream.unfoldr get n `Stream.sized` Exact n)
+streamR v = v `seq` n `seq` (Stream.unfoldr get n `Stream.sized` Exact n)
   where
     n = length v
 
@@ -1717,7 +1816,7 @@ gfoldl f z v = z fromList `f` toList v
 
 mkType :: String -> DataType
 {-# INLINE mkType #-}
-mkType = mkNorepType
+mkType = mkNoRepType
 
 dataCast :: (Vector v a, Data a, Typeable1 v, Typeable1 t)
          => (forall d. Data  d => c (t d)) -> Maybe  (c (v a))