Simplify IVector.inplace
[darcs-mirrors/vector.git] / Data / Vector / IVector.hs
index 0f01cbc..fa0ed7b 100644 (file)
@@ -1,4 +1,5 @@
-{-# LANGUAGE Rank2Types, MultiParamTypeClasses #-}
+{-# LANGUAGE Rank2Types, MultiParamTypeClasses, FlexibleContexts,
+             ScopedTypeVariables #-}
 -- |
 -- Module      : Data.Vector.IVector
 -- Copyright   : (c) Roman Leshchinskiy 2008
@@ -27,10 +28,16 @@ module Data.Vector.IVector (
   (!), head, last,
 
   -- * Subvectors
-  slice, takeSlice, take, dropSlice, drop,
+  slice, extract, takeSlice, take, dropSlice, drop,
+
+  -- * Permutations
+  (//), update, bpermute,
 
   -- * Mapping and zipping
-  map, zipWith,
+  map, zipWith, zip,
+
+  -- * Comparisons
+  eq, cmp,
 
   -- * Filtering
   filter, takeWhileSlice, takeWhile, dropWhileSlice, dropWhile,
@@ -48,18 +55,25 @@ module Data.Vector.IVector (
   stream, unstream,
 
   -- * MVector-based initialisation
-  create,
+  new,
 
   -- * Unsafe functions
-  unsafeSlice, unsafeIndex
+  unsafeSlice, unsafeIndex,
+
+  -- * Utility functions
+  vlength, vnew
 ) where
 
 import qualified Data.Vector.MVector as MVector
 import           Data.Vector.MVector ( MVector )
 
-import qualified Data.Vector.Stream as Stream
-import           Data.Vector.Stream ( Stream )
-import           Data.Vector.Stream.Size
+import qualified Data.Vector.MVector.New as New
+import           Data.Vector.MVector.New ( New )
+
+import qualified Data.Vector.Fusion.Stream as Stream
+import           Data.Vector.Fusion.Stream ( Stream, MStream )
+import qualified Data.Vector.Fusion.Stream.Monadic as MStream
+import           Data.Vector.Fusion.Stream.Size
 
 import Control.Exception ( assert )
 
@@ -67,20 +81,19 @@ import Prelude hiding ( length,
                         replicate, (++),
                         head, last,
                         init, tail, take, drop,
-                        map, zipWith,
+                        map, zipWith, zip,
                         filter, takeWhile, dropWhile,
                         elem, notElem,
                         foldl, foldl1, foldr, foldr1 )
 
--- | Class of immutable vectors. Just like with 'MVector', the type of the
--- elements can be restricted by using GADTs.
+-- | Class of immutable vectors.
 --
 class IVector v a where
-  -- | Construct a pure vector from a monadic initialiser.
-  create       :: (forall mv m. MVector mv m a => m (mv m a)) -> v a
+  -- | Construct a pure vector from a monadic initialiser (not fusible!)
+  vnew         :: (forall mv m. MVector mv m a => m (mv a)) -> v a
 
-  -- | Length of the vector
-  length       :: v a -> Int
+  -- | Length of the vector (not fusible!)
+  vlength      :: v a -> Int
 
   -- | Yield a part of the vector without copying it. No range checks!
   unsafeSlice  :: v a -> Int -> Int -> v a
@@ -106,6 +119,14 @@ class IVector v a where
   --
   unsafeIndex  :: v a -> Int -> (a -> b) -> b
 
+-- Fusion
+-- ------
+
+-- | Construct a pure vector from a monadic initialiser 
+new :: IVector v a => New a -> v a
+{-# INLINE_STREAM new #-}
+new m = vnew (New.run m)
+
 -- | Convert a vector to a 'Stream'
 stream :: IVector v a => v a -> Stream a
 {-# INLINE_STREAM stream #-}
@@ -119,16 +140,52 @@ stream v = v `seq` (Stream.unfold get 0 `Stream.sized` Exact n)
 
 -- | Create a vector from a 'Stream'
 unstream :: IVector v a => Stream a -> v a
-{-# INLINE_STREAM unstream #-}
-unstream s = create (MVector.unstream s)
+{-# INLINE unstream #-}
+unstream s = new (New.unstream s)
 
 {-# RULES
 
 "stream/unstream [IVector]" forall s.
-  stream (unstream s) = s
+  stream (new (New.unstream s)) = s
+
+"New.unstream/stream/new [IVector]" forall p.
+  New.unstream (stream (new p)) = p
+
+ #-}
+
+inplace :: (forall m. Monad m => MStream m a -> MStream m a)
+        -> Stream a -> Stream a
+{-# INLINE_STREAM inplace #-}
+inplace f s = f s
+
+{-# RULES
+
+"inplace [IVector]"
+  forall (f :: forall m. Monad m => MStream m a -> MStream m a) m.
+  New.unstream (inplace f (stream (new m))) = New.transform f m
+
+"inplace/inplace [IVector]"
+  forall (f :: forall m. Monad m => MStream m a -> MStream m a)
+         (g :: forall m. Monad m => MStream m a -> MStream m a)
+         s.
+  inplace f (inplace g s) = inplace (f . g) s
 
  #-}
 
+-- Length
+-- ------
+
+length :: IVector v a => v a -> Int
+{-# INLINE_STREAM length #-}
+length v = vlength v
+
+{-# RULES
+
+"length/unstream [IVector]" forall s.
+  length (new (New.unstream s)) = Stream.length s
+
+  #-}
+
 -- Construction
 -- ------------
 
@@ -168,20 +225,33 @@ v ++ w = unstream (stream v Stream.++ stream w)
 
 -- | Indexing
 (!) :: IVector v a => v a -> Int -> a
-{-# INLINE (!) #-}
+{-# INLINE_STREAM (!) #-}
 v ! i = assert (i >= 0 && i < length v)
       $ unsafeIndex v i id
 
 -- | First element
 head :: IVector v a => v a -> a
-{-# INLINE head #-}
+{-# INLINE_STREAM head #-}
 head v = v ! 0
 
 -- | Last element
 last :: IVector v a => v a -> a
-{-# INLINE last #-}
+{-# INLINE_STREAM last #-}
 last v = v ! (length v - 1)
 
+{-# RULES
+
+"(!)/unstream [IVector]" forall i s.
+  new (New.unstream s) ! i = s Stream.!! i
+
+"head/unstream [IVector]" forall s.
+  head (new (New.unstream s)) = Stream.head s
+
+"last/unstream [IVector]" forall s.
+  last (new (New.unstream s)) = Stream.last s
+
+ #-}
+
 -- Subarrays
 -- ---------
 
@@ -194,6 +264,13 @@ 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 the first @n@ elements without copying.
 takeSlice :: IVector v a => Int -> v a -> v a
 {-# INLINE takeSlice #-}
@@ -214,6 +291,35 @@ drop :: IVector v a => Int -> v a -> v a
 {-# INLINE drop #-}
 drop n = unstream . Stream.drop n . stream
 
+{-# RULES
+
+"slice/extract [IVector]" forall i n s.
+  slice (new (New.unstream s)) i n = extract (new (New.unstream s)) i n
+
+"takeSlice/unstream [IVector]" forall n s.
+  takeSlice n (new (New.unstream s)) = take n (new (New.unstream s))
+
+"dropSlice/unstream [IVector]" forall n s.
+  dropSlice n (new (New.unstream s)) = drop n (new (New.unstream s))
+
+  #-}
+
+-- Permutations
+-- ------------
+
+(//) :: IVector v a => v a -> [(Int, a)] -> v a
+{-# INLINE (//) #-}
+v // us = new (New.update (New.unstream (stream v))
+                          (Stream.fromList us))
+
+update :: (IVector v a, IVector v (Int, a)) => v a -> v (Int, a) -> v a
+{-# INLINE update #-}
+update v w = new (New.update (New.unstream (stream v)) (stream w))
+
+bpermute :: (IVector v a, IVector v Int) => v a -> v Int -> v a
+{-# INLINE bpermute #-}
+bpermute v is = is `seq` map (v!) is
+
 -- Mapping/zipping
 -- ---------------
 
@@ -222,18 +328,43 @@ map :: (IVector v a, IVector v b) => (a -> b) -> v a -> v b
 {-# INLINE map #-}
 map f = unstream . Stream.map f . stream
 
+inplace_map :: IVector v a => (a -> a) -> v a -> v a
+{-# INLINE inplace_map #-}
+inplace_map f = unstream . inplace (MStream.map f) . stream
+
+{-# RULES
+
+"map->inplace_map [IVector]" map = inplace_map
+
+ #-}
+
 -- | Zip two vectors with the given function.
 zipWith :: (IVector v a, IVector v b, IVector v c) => (a -> b -> c) -> v a -> v b -> v c
 {-# INLINE zipWith #-}
 zipWith f xs ys = unstream (Stream.zipWith f (stream xs) (stream ys))
 
+zip :: (IVector v a, IVector v b, IVector v (a,b)) => v a -> v b -> v (a, b)
+{-# INLINE zip #-}
+zip = zipWith (,)
+
+-- Comparisons
+-- -----------
+
+eq :: (IVector v a, Eq a) => v a -> v a -> Bool
+{-# INLINE eq #-}
+xs `eq` ys = stream xs == stream ys
+
+cmp :: (IVector v a, Ord a) => v a -> v a -> Ordering
+{-# INLINE cmp #-}
+cmp xs ys = compare (stream xs) (stream ys)
+
 -- Filtering
 -- ---------
 
 -- | Drop elements which do not satisfy the predicate
 filter :: IVector v a => (a -> Bool) -> v a -> v a
 {-# INLINE filter #-}
-filter f = unstream . Stream.filter f . stream
+filter f = unstream . inplace (MStream.filter f) . stream
 
 -- | Yield the longest prefix of elements satisfying the predicate without
 -- copying.
@@ -263,6 +394,16 @@ dropWhile :: IVector v a => (a -> Bool) -> v a -> v a
 {-# INLINE dropWhile #-}
 dropWhile f = unstream . Stream.dropWhile f . stream
 
+{-# RULES
+
+"takeWhileSlice/unstream" forall f s.
+  takeWhileSlice f (new (New.unstream s)) = takeWhile f (new (New.unstream s))
+
+"dropWhileSlice/unstream" forall f s.
+  dropWhileSlice f (new (New.unstream s)) = dropWhile f (new (New.unstream s))
+
+ #-}
+
 -- Searching
 -- ---------