Make slicing work with the in-place stuff
authorRoman Leshchinskiy <rl@cse.unsw.edu.au>
Fri, 3 Oct 2008 04:43:18 +0000 (04:43 +0000)
committerRoman Leshchinskiy <rl@cse.unsw.edu.au>
Fri, 3 Oct 2008 04:43:18 +0000 (04:43 +0000)
Data/Vector/IVector.hs
Data/Vector/MVector/New.hs

index af65016..e8d8412 100644 (file)
@@ -22,13 +22,13 @@ module Data.Vector.IVector (
   length,
 
   -- * Construction
-  empty, singleton, cons, snoc, replicate, (++),
+  empty, singleton, cons, snoc, replicate, (++), copy,
 
   -- * Accessing individual elements
   (!), head, last,
 
   -- * Subvectors
-  slice, extract, takeSlice, take, dropSlice, drop,
+  slice, init, tail, take, drop,
 
   -- * Permutations
   (//), update, bpermute,
@@ -235,6 +235,18 @@ infixr 5 ++
 {-# INLINE (++) #-}
 v ++ w = unstream (stream v Stream.++ stream w)
 
+-- | Create a copy of a vector. Useful when dealing with slices.
+copy :: IVector v a => v a -> v a
+{-# INLINE_STREAM copy #-}
+copy = unstream . stream
+
+{-# RULES
+
+"copy/unstream [IVector]" forall v s.
+  copy (new' v (New.unstream s)) = new' v (New.unstream s)
+
+ #-}
+
 -- Accessing individual elements
 -- -----------------------------
 
@@ -281,37 +293,42 @@ slice :: IVector v a => v a -> Int   -- ^ starting index
 slice v i n = assert (i >= 0 && n >= 0  && i+n <= length v)
             $ unsafeSlice v i n
 
--- | Copy @n@ elements starting at the given position to a new vector.
-extract :: IVector v a => v a -> Int  -- ^ starting index
-                              -> Int  -- ^ length
-                              -> v a
-{-# INLINE extract #-}
-extract v i n = unstream (Stream.extract (stream v) i n)
+-- | Yield all but the last element without copying.
+init :: IVector v a => v a -> v a
+{-# INLINE_STREAM init #-}
+init v = slice v 0 (length v - 1)
 
--- | Yield the first @n@ elements without copying.
-takeSlice :: IVector v a => Int -> v a -> v a
-{-# INLINE takeSlice #-}
-takeSlice n v = slice v 0 n
+-- | All but the first element (without copying).
+tail :: IVector v a => v a -> v a
+{-# INLINE_STREAM tail #-}
+tail v = slice v 1 (length v - 1)
 
--- | Copy the first @n@ elements to a new vector.
+-- | Yield the first @n@ elements without copying.
 take :: IVector v a => Int -> v a -> v a
-{-# INLINE take #-}
-take n = unstream . Stream.take n . stream
+{-# INLINE_STREAM take #-}
+take n v = slice v 0 (min n (length v))
 
 -- | Yield all but the first @n@ elements without copying.
-dropSlice :: IVector v a => Int -> v a -> v a
-{-# INLINE dropSlice #-}
-dropSlice n v = slice v n (length v - n)
-
--- | Copy all but the first @n@ elements to a new vectors.
 drop :: IVector v a => Int -> v a -> v a
-{-# INLINE drop #-}
-drop n = unstream . Stream.drop n . stream
+{-# INLINE_STREAM drop #-}
+drop n v = slice v n (max 0 (length v - n))
 
 {-# RULES
 
-"slice/unstream [IVector]" forall v i n s.
-  slice (new' v (New.unstream s)) i n = extract (new' v (New.unstream s)) i n
+"slice/new [IVector]" forall v p i n.
+  slice (new' v p) i n = new' v (New.slice p i n)
+
+"init/new [IVector]" forall v p.
+  init (new' v p) = new' v (New.init p)
+
+"tail/new [IVector]" forall v p.
+  tail (new' v p) = new' v (New.tail p)
+
+"take/new [IVector]" forall n v p.
+  take n (new' v p) = new' v (New.take n p)
+
+"drop/new [IVector]" forall n v p.
+  drop n (new' v p) = new' v (New.drop n p)
 
   #-}
 
@@ -382,7 +399,7 @@ filter f = unstream . inplace (MStream.filter f) . stream
 takeWhileSlice :: IVector v a => (a -> Bool) -> v a -> v a
 {-# INLINE_STREAM takeWhileSlice #-}
 takeWhileSlice f v = case findIndex (not . f) v of
-                       Just n  -> takeSlice n v
+                       Just n  -> take n v
                        Nothing -> v
 
 -- | Copy the longest prefix of elements satisfying the predicate to a new
@@ -396,7 +413,7 @@ takeWhile f = unstream . Stream.takeWhile f . stream
 dropWhileSlice :: IVector v a => (a -> Bool) -> v a -> v a
 {-# INLINE_STREAM dropWhileSlice #-}
 dropWhileSlice f v = case findIndex (not . f) v of
-                       Just n  -> dropSlice n v
+                       Just n  -> drop n v
                        Nothing -> v
 
 -- | Drop the longest prefix of elements that satisfy the predicate and copy
index 22b9238..9aa2117 100644 (file)
@@ -3,11 +3,12 @@
 #include "phases.h"
 
 module Data.Vector.MVector.New (
-  New(..), run, unstream, transform, update, reverse
+  New(..), run, unstream, transform, update, reverse,
+  slice, init, tail, take, drop
 ) where
 
 import qualified Data.Vector.MVector as MVector
-import           Data.Vector.MVector ( MVector )
+import           Data.Vector.MVector ( MVector, MVectorPure )
 
 import           Data.Vector.Fusion.Stream ( Stream, MStream )
 import qualified Data.Vector.Fusion.Stream as Stream
@@ -15,7 +16,7 @@ import qualified Data.Vector.Fusion.Stream as Stream
 import qualified Data.Vector.Fusion.Stream.Monadic as MStream
 
 import Control.Monad  ( liftM )
-import Prelude hiding ( reverse, map, filter )
+import Prelude hiding ( init, tail, take, drop, reverse, map, filter )
 
 newtype New a = New (forall m mv. MVector mv m a => m (mv a))
 
@@ -23,6 +24,10 @@ run :: MVector mv m a => New a -> m (mv a)
 {-# INLINE run #-}
 run (New p) = p
 
+apply :: (forall mv a. MVectorPure mv a => mv a -> mv a) -> New a -> New a
+{-# INLINE apply #-}
+apply f (New p) = New (liftM f p)
+
 modify :: New a -> (forall m mv. MVector mv m a => mv a -> m ()) -> New a
 {-# INLINE modify #-}
 modify (New p) q = New (do { v <- p; q v; return v })
@@ -45,6 +50,26 @@ transform f (New p) = New (MVector.transform f =<< p)
 
  #-}
 
+slice :: New a -> Int -> Int -> New a
+{-# INLINE_STREAM slice #-}
+slice m i n = apply (\v -> MVector.slice v i n) m
+
+init :: New a -> New a
+{-# INLINE_STREAM init #-}
+init m = apply (\v -> MVector.slice v 0 (MVector.length v - 1)) m
+
+tail :: New a -> New a
+{-# INLINE_STREAM tail #-}
+tail m = apply (\v -> MVector.slice v 1 (MVector.length v - 1)) m
+
+take :: Int -> New a -> New a
+{-# INLINE_STREAM take #-}
+take n m = apply (\v -> MVector.slice v 0 (min n (MVector.length v))) m
+
+drop :: Int -> New a -> New a
+{-# INLINE_STREAM drop #-}
+drop n m = apply (\v -> MVector.slice v n (max 0 (MVector.length v - n))) m
+
 update :: New a -> Stream (Int, a) -> New a
 {-# INLINE_STREAM update #-}
 update m s = modify m (\v -> MVector.update v s)