Clean up handling of New.slice and friends
authorRoman Leshchinskiy <rl@cse.unsw.edu.au>
Wed, 9 Dec 2009 14:20:53 +0000 (14:20 +0000)
committerRoman Leshchinskiy <rl@cse.unsw.edu.au>
Wed, 9 Dec 2009 14:20:53 +0000 (14:20 +0000)
Data/Vector/Generic/Mutable.hs
Data/Vector/Generic/New.hs

index 2b88796..a757242 100644 (file)
@@ -16,10 +16,13 @@ module Data.Vector.Generic.Mutable (
   MVector(..),
 
   -- * Operations on mutable vectors
-  length, overlaps, slice, new, newWith, read, write, clear, set, copy, grow,
+  length, overlaps, new, newWith, read, write, clear, set, copy, grow,
+
+  slice, take, drop, init, tail,
+  unsafeSlice, unsafeInit, unsafeTail,
 
   -- * Unsafe operations
-  unsafeSlice, unsafeNew, unsafeNewWith, unsafeRead, unsafeWrite,
+  unsafeNew, unsafeNewWith, unsafeRead, unsafeWrite,
   unsafeCopy, unsafeGrow,
 
   -- * Internal operations
@@ -38,7 +41,8 @@ import GHC.Float (
     double2Int, int2Double
   )
 
-import Prelude hiding ( length, reverse, map, read )
+import Prelude hiding ( length, reverse, map, read,
+                        take, drop, init, tail )
 
 #include "vector.h"
 
@@ -139,19 +143,6 @@ class MVector v a where
     where
       n = length v
 
-
-
--- | Yield a part of the mutable vector without copying it. No bounds checks
--- are performed.
-unsafeSlice :: MVector v a => Int  -- ^ starting index
-                           -> Int  -- ^ length of the slice
-                           -> v s a
-                           -> v s a
-{-# INLINE unsafeSlice #-}
-unsafeSlice i n v = UNSAFE_CHECK(checkSlice) "unsafeSlice" i n (length v)
-                  $ basicUnsafeSlice i n v
-
-
 -- | Create a mutable vector of the given length. The length is not checked.
 unsafeNew :: (PrimMonad m, MVector v a) => Int -> m (v (PrimState m) a)
 {-# INLINE unsafeNew #-}
@@ -209,12 +200,6 @@ overlaps :: MVector v a => v s a -> v s a -> Bool
 {-# INLINE overlaps #-}
 overlaps = basicOverlaps
 
--- | Yield a part of the mutable vector without copying it.
-slice :: MVector v a => Int -> Int -> v s a -> v s a
-{-# INLINE slice #-}
-slice i n v = BOUNDS_CHECK(checkSlice) "slice" i n (length v)
-            $ unsafeSlice i n v
-
 -- | Create a mutable vector of the given length.
 new :: (PrimMonad m, MVector v a) => Int -> m (v (PrimState m) a)
 {-# INLINE new #-}
@@ -279,6 +264,50 @@ enlarge v = unsafeGrow v
           $ double2Int
           $ int2Double (length v) * gROWTH_FACTOR
 
+
+-- | Yield a part of the mutable vector without copying it.
+slice :: MVector v a => Int -> Int -> v s a -> v s a
+{-# INLINE slice #-}
+slice i n v = BOUNDS_CHECK(checkSlice) "slice" i n (length v)
+            $ unsafeSlice i n v
+
+take :: MVector v a => Int -> v s a -> v s a
+{-# INLINE take #-}
+take n v = unsafeSlice 0 (min (max n 0) (length v)) v
+
+drop :: MVector v a => Int -> v s a -> v s a
+{-# INLINE drop #-}
+drop n v = unsafeSlice (min m n') (max 0 (m - n')) v
+  where
+    n' = max n 0
+    m  = length v
+
+init :: MVector v a => v s a -> v s a
+{-# INLINE init #-}
+init v = slice 0 (length v - 1) v
+
+tail :: MVector v a => v s a -> v s a
+{-# INLINE tail #-}
+tail v = slice 1 (length v - 1) v
+
+-- | Yield a part of the mutable vector without copying it. No bounds checks
+-- are performed.
+unsafeSlice :: MVector v a => Int  -- ^ starting index
+                           -> Int  -- ^ length of the slice
+                           -> v s a
+                           -> v s a
+{-# INLINE unsafeSlice #-}
+unsafeSlice i n v = UNSAFE_CHECK(checkSlice) "unsafeSlice" i n (length v)
+                  $ basicUnsafeSlice i n v
+
+unsafeInit :: MVector v a => v s a -> v s a
+{-# INLINE unsafeInit #-}
+unsafeInit v = unsafeSlice 0 (length v - 1) v
+
+unsafeTail :: MVector v a => v s a -> v s a
+{-# INLINE unsafeTail #-}
+unsafeTail v = unsafeSlice 1 (length v - 1) v
+
 unsafeAppend1 :: (PrimMonad m, MVector v a)
         => v (PrimState m) a -> Int -> a -> m (v (PrimState m) a)
 {-# INLINE_INNER unsafeAppend1 #-}
index 56645e6..de601d5 100644 (file)
@@ -15,7 +15,8 @@
 module Data.Vector.Generic.New (
   New(..), run, unstream, transform, accum, update, reverse,
   slice, init, tail, take, drop,
-  unsafeSlice, unsafeAccum, unsafeUpdate
+  unsafeSlice, unsafeInit, unsafeTail,
+  unsafeAccum, unsafeUpdate
 ) where
 
 import qualified Data.Vector.Generic.Mutable as MVector
@@ -71,37 +72,39 @@ slice :: Int -> Int -> New a -> New a
 {-# INLINE_STREAM slice #-}
 slice i n m = apply (MVector.slice i n) m
 
-unsafeSlice :: Int -> Int -> New a -> New a
-{-# INLINE_STREAM unsafeSlice #-}
-unsafeSlice i n m = apply (MVector.unsafeSlice i n) m
-
 init :: New a -> New a
 {-# INLINE_STREAM init #-}
-init m = apply (\v -> MVector.slice 0 (MVector.length v - 1) v) m
+init m = apply MVector.init m
 
 tail :: New a -> New a
 {-# INLINE_STREAM tail #-}
-tail m = apply (\v -> MVector.slice 1 (MVector.length v - 1) v) m
+tail m = apply MVector.tail m
 
 take :: Int -> New a -> New a
 {-# INLINE_STREAM take #-}
-take n m = apply (\v -> MVector.unsafeSlice 0
-                                   (min (max 0 n) (MVector.length v)) v) m
+take n m = apply (MVector.take n) m
 
 drop :: Int -> New a -> New a
 {-# INLINE_STREAM drop #-}
-drop n m = apply (\v -> MVector.unsafeSlice
-                                (min (max 0 n) (MVector.length v))
-                                (max 0 (MVector.length v - n)) v) m
+drop n m = apply (MVector.drop n) m
+
+unsafeSlice :: Int -> Int -> New a -> New a
+{-# INLINE_STREAM unsafeSlice #-}
+unsafeSlice i n m = apply (MVector.unsafeSlice i n) m
+
+unsafeInit :: New a -> New a
+{-# INLINE_STREAM unsafeInit #-}
+unsafeInit m = apply MVector.unsafeInit m
+
+unsafeTail :: New a -> New a
+{-# INLINE_STREAM unsafeTail #-}
+unsafeTail m = apply MVector.unsafeTail m
 
 {-# RULES
 
 "slice/unstream [New]" forall i n s.
   slice i n (unstream s) = unstream (Stream.slice i n s)
 
-"unsafeSlice/unstream [New]" forall i n s.
-  unsafeSlice i n (unstream s) = unstream (Stream.slice i n s)
-
 "init/unstream [New]" forall s.
   init (unstream s) = unstream (Stream.init s)
 
@@ -114,6 +117,15 @@ drop n m = apply (\v -> MVector.unsafeSlice
 "drop/unstream [New]" forall n s.
   drop n (unstream s) = unstream (Stream.drop n s)
 
+"unsafeSlice/unstream [New]" forall i n s.
+  unsafeSlice i n (unstream s) = unstream (Stream.slice i n s)
+
+"unsafeInit/unstream [New]" forall s.
+  unsafeInit (unstream s) = unstream (Stream.init s)
+
+"unsafeTail/unstream [New]" forall s.
+  unsafeTail (unstream s) = unstream (Stream.tail s)
+
   #-}
 
 unsafeAccum :: (a -> b -> a) -> New a -> Stream (Int, b) -> New a