Added splitAt functions (contributed by Bas van Dijk)
[darcs-mirrors/vector.git] / Data / Vector / Generic.hs
index bd316cd..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
@@ -130,7 +130,7 @@ module Data.Vector.Generic (
   convert,
 
   -- ** Mutable vectors
-  thaw, thawMany, copy, unsafeCopy,
+  freeze, thaw, copy, unsafeFreeze, unsafeThaw, unsafeCopy,
 
   -- * Fusion support
 
@@ -170,7 +170,7 @@ import qualified Data.List as List
 import Prelude hiding ( length, null,
                         replicate, (++), concat,
                         head, last,
-                        init, tail, take, drop, reverse,
+                        init, tail, take, drop, splitAt, reverse,
                         map, concat, concatMap,
                         zipWith, zipWith3, zip, zip3, unzip, unzip3,
                         filter, takeWhile, dropWhile, span, break,
@@ -182,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
 -- ------------------
 
@@ -222,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 #-}
@@ -375,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 #-}
@@ -391,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
@@ -792,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
 -- ------------------------
@@ -1566,6 +1612,20 @@ 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 #-}
@@ -1574,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)
@@ -1594,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.
@@ -1619,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
 
@@ -1658,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
 
@@ -1749,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))