Add INLINE pragma
[darcs-mirrors/vector.git] / Data / Vector / IVector.hs
index ce8fd99..a8a257c 100644 (file)
@@ -5,7 +5,7 @@
 -- Copyright   : (c) Roman Leshchinskiy 2008
 -- License     : BSD-style
 --
--- Maintainer  : rl@cse.unsw.edu.au
+-- Maintainer  : Roman Leshchinskiy <rl@cse.unsw.edu.au>
 -- Stability   : experimental
 -- Portability : non-portable
 -- 
@@ -19,34 +19,49 @@ module Data.Vector.IVector (
   IVector,
 
   -- * Length information
-  length,
+  length, null,
 
   -- * 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,
+  accum, (//), update, backpermute, reverse,
 
-  -- * Mapping and zipping
-  map, zipWith, zip,
+  -- * Mapping
+  map, concatMap,
+
+  -- * Zipping and unzipping
+  zipWith, zipWith3, zip, zip3, unzip, unzip3,
 
   -- * Comparisons
   eq, cmp,
 
   -- * Filtering
-  filter, takeWhileSlice, takeWhile, dropWhileSlice, dropWhile,
+  filter, takeWhile, dropWhile,
 
   -- * Searching
   elem, notElem, find, findIndex,
 
   -- * Folding
   foldl, foldl1, foldl', foldl1', foldr, foldr1,
+  -- * Specialised folds
+  and, or, sum, product, maximum, minimum,
+  
+  -- * Enumeration
+  enumFromTo, enumFromThenTo,
+
+  -- * Unfolding
+  unfoldr,
+
+  -- * Scans
+  prescanl, prescanl',
 
   -- * Conversion to/from lists
   toList, fromList,
@@ -71,21 +86,23 @@ 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 )
-import qualified Data.Vector.Fusion.MStream as MStream
-import           Data.Vector.Fusion.MStream ( MStream )
+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 )
 
-import Prelude hiding ( length,
+import Prelude hiding ( length, null,
                         replicate, (++),
                         head, last,
-                        init, tail, take, drop,
-                        map, zipWith, zip,
+                        init, tail, take, drop, reverse,
+                        map, concatMap,
+                        zipWith, zipWith3, zip, zip3, unzip, unzip3,
                         filter, takeWhile, dropWhile,
                         elem, notElem,
-                        foldl, foldl1, foldr, foldr1 )
+                        foldl, foldl1, foldr, foldr1,
+                        and, or, sum, product, maximum, minimum,
+                        enumFromTo, enumFromThenTo )
 
 -- | Class of immutable vectors.
 --
@@ -125,13 +142,21 @@ class IVector v a where
 
 -- | 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)
+{-# INLINE new #-}
+new m = new' undefined m
+
+-- | Same as 'new' but with a dummy argument necessary for correctly typing
+-- the rule @uninplace@.
+--
+-- See http://hackage.haskell.org/trac/ghc/ticket/2600
+new' :: IVector v a => 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 #-}
-stream v = v `seq` (Stream.unfold get 0 `Stream.sized` Exact n)
+stream v = v `seq` (Stream.unfoldr get 0 `Stream.sized` Exact n)
   where
     n = length v
 
@@ -146,30 +171,34 @@ unstream s = new (New.unstream s)
 
 {-# RULES
 
-"stream/unstream [IVector]" forall s.
-  stream (new (New.unstream s)) = s
+"stream/unstream [IVector]" forall s.
+  stream (new' v (New.unstream s)) = s
 
-"New.unstream/stream/new [IVector]" forall p.
-  New.unstream (stream (new p)) = p
+"New.unstream/stream/new [IVector]" forall p.
+  New.unstream (stream (new' v p)) = p
 
  #-}
 
-inplace :: (Stream a -> Stream a)
-        -> (forall m. Monad m => MStream m a -> MStream m a)
+inplace :: (forall m. Monad m => MStream m a -> MStream m a)
         -> Stream a -> Stream a
 {-# INLINE_STREAM inplace #-}
-inplace f s = f s
+inplace f s = f s
 
 {-# RULES
 
 "inplace [IVector]"
-  forall f (mf :: forall m. Monad m => MStream m a -> MStream m a) m.
-  New.unstream (inplace f mf (stream (new m))) = New.inplace mf m
+  forall (f :: forall m. Monad m => MStream m a -> MStream m a) v m.
+  New.unstream (inplace f (stream (new' v m))) = New.transform f m
+
+"uninplace [IVector]"
+  forall (f :: forall m. Monad m => MStream m a -> MStream m a) v m.
+  stream (new' v (New.transform f m)) = inplace f (stream (new' v m))
 
 "inplace/inplace [IVector]"
-  forall f (mf :: forall m. Monad m => MStream m a -> MStream m a)
-         g (mg :: forall m. Monad m => MStream m a -> MStream m a) s.
-  inplace f mf (inplace g mg s) = inplace (f . g) (mf . mg) s
+  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
 
  #-}
 
@@ -182,8 +211,19 @@ length v = vlength v
 
 {-# RULES
 
-"length/unstream [IVector]" forall s.
-  length (new (New.unstream s)) = Stream.length s
+"length/unstream [IVector]" forall v s.
+  length (new' v (New.unstream s)) = Stream.length s
+
+  #-}
+
+null :: IVector v a => v a -> Bool
+{-# INLINE_STREAM null #-}
+null v = vlength v == 0
+
+{-# RULES
+
+"null/unstream [IVector]" forall v s.
+  null (new' v (New.unstream s)) = Stream.null s
 
   #-}
 
@@ -221,6 +261,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
 -- -----------------------------
 
@@ -242,72 +294,81 @@ last v = v ! (length v - 1)
 
 {-# RULES
 
-"(!)/unstream [IVector]" forall i s.
-  new (New.unstream s) ! i = s Stream.!! i
+"(!)/unstream [IVector]" forall i s.
+  new' v (New.unstream s) ! i = s Stream.!! i
 
-"head/unstream [IVector]" forall s.
-  head (new (New.unstream s)) = Stream.head s
+"head/unstream [IVector]" forall s.
+  head (new' v (New.unstream s)) = Stream.head s
 
-"last/unstream [IVector]" forall s.
-  last (new (New.unstream s)) = Stream.last s
+"last/unstream [IVector]" forall s.
+  last (new' v (New.unstream s)) = Stream.last s
 
  #-}
 
 -- Subarrays
 -- ---------
 
+-- FIXME: slicing doesn't work with the inplace stuff at the moment
+
 -- | Yield a part of the vector without copying it. Safer version of
 -- 'unsafeSlice'.
 slice :: IVector v a => v a -> Int   -- ^ starting index
                             -> Int   -- ^ length
                             -> v a
-{-# INLINE slice #-}
+{-# INLINE_STREAM slice #-}
 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))
+  where n' = max n 0
 
 -- | 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 (min n' len) (max 0 (len - n'))
+  where n' = max n 0
+        len = length v
 
 {-# RULES
 
-"slice/extract [IVector]" forall i n s.
-  slice (new (New.unstream s)) i n = extract (new (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)
 
-"takeSlice/unstream [IVector]" forall n s.
-  takeSlice n (new (New.unstream s)) = take n (new (New.unstream s))
+"take/new [IVector]" forall n v p.
+  take n (new' v p) = new' v (New.take n p)
 
-"dropSlice/unstream [IVector]" forall n s.
-  dropSlice n (new (New.unstream s)) = drop n (new (New.unstream s))
+"drop/new [IVector]" forall n v p.
+  drop n (new' v p) = new' v (New.drop n p)
 
   #-}
 
 -- Permutations
 -- ------------
 
+accum :: IVector v a => (a -> b -> a) -> v a -> [(Int,b)] -> v a
+{-# INLINE accum #-}
+accum f v us = new (New.accum f (New.unstream (stream v))
+                                (Stream.fromList us))
+
 (//) :: IVector v a => v a -> [(Int, a)] -> v a
 {-# INLINE (//) #-}
 v // us = new (New.update (New.unstream (stream v))
@@ -317,12 +378,16 @@ 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
+backpermute :: (IVector v a, IVector v Int) => v a -> v Int -> v a
+{-# INLINE backpermute #-}
+backpermute v is = v `seq` map (v!) is
+
+reverse :: (IVector v a) => v a -> v a
+{-# INLINE reverse #-}
+reverse = new . New.reverse . New.unstream . stream
 
--- Mapping/zipping
--- ---------------
+-- Mapping
+-- -------
 
 -- | Map a function over a vector
 map :: (IVector v a, IVector v b) => (a -> b) -> v a -> v b
@@ -331,7 +396,7 @@ 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 (Stream.map f) (MStream.map f) . stream
+inplace_map f = unstream . inplace (MStream.map f) . stream
 
 {-# RULES
 
@@ -339,15 +404,39 @@ inplace_map f = unstream . inplace (Stream.map f) (MStream.map f) . stream
 
  #-}
 
+concatMap :: (IVector v a, IVector v b) => (a -> v b) -> v a -> v b
+{-# INLINE concatMap #-}
+concatMap f = unstream . Stream.concatMap (stream . f) . stream
+
+-- Zipping/unzipping
+-- -----------------
+
 -- | 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 three vectors with the given function.
+zipWith3 :: (IVector v a, IVector v b, IVector v c, IVector v d) => (a -> b -> c -> d) -> v a -> v b -> v c -> v d
+{-# INLINE zipWith3 #-}
+zipWith3 f xs ys zs = unstream (Stream.zipWith3 f (stream xs) (stream ys) (stream zs))
+
 zip :: (IVector v a, IVector v b, IVector v (a,b)) => v a -> v b -> v (a, b)
 {-# INLINE zip #-}
 zip = zipWith (,)
 
+zip3 :: (IVector v a, IVector v b, IVector v c, IVector v (a, b, c)) => v a -> v b -> v c -> v (a, b, c)
+{-# INLINE zip3 #-}
+zip3 = zipWith3 (,,)
+
+unzip :: (IVector v a, IVector v b, IVector v (a,b)) => v (a, b) -> (v a, v b)
+{-# INLINE unzip #-}
+unzip xs = (map fst xs, map snd xs)
+
+unzip3 :: (IVector v a, IVector v b, IVector v c, IVector v (a, b, c)) => v (a, b, c) -> (v a, v b, v c)
+{-# INLINE unzip3 #-}
+unzip3 xs = (map (\(a, b, c) -> a) xs, map (\(a, b, c) -> b) xs, map (\(a, b, c) -> c) xs)
+
 -- Comparisons
 -- -----------
 
@@ -365,46 +454,18 @@ cmp xs ys = compare (stream xs) (stream ys)
 -- | Drop elements which do not satisfy the predicate
 filter :: IVector v a => (a -> Bool) -> v a -> v a
 {-# INLINE filter #-}
-filter f = unstream . inplace (Stream.filter f) (MStream.filter f) . stream
-
--- | Yield the longest prefix of elements satisfying the predicate without
--- copying.
-takeWhileSlice :: IVector v a => (a -> Bool) -> v a -> v a
-{-# INLINE takeWhileSlice #-}
-takeWhileSlice f v = case findIndex (not . f) v of
-                       Just n  -> takeSlice n v
-                       Nothing -> v
-
--- | Copy the longest prefix of elements satisfying the predicate to a new
--- vector
+filter f = unstream . inplace (MStream.filter f) . stream
+
+-- | Yield the longest prefix of elements satisfying the predicate.
 takeWhile :: IVector v a => (a -> Bool) -> v a -> v a
 {-# INLINE takeWhile #-}
 takeWhile f = unstream . Stream.takeWhile f . stream
 
--- | Drop the longest prefix of elements that satisfy the predicate without
--- copying
-dropWhileSlice :: IVector v a => (a -> Bool) -> v a -> v a
-{-# INLINE dropWhileSlice #-}
-dropWhileSlice f v = case findIndex (not . f) v of
-                       Just n  -> dropSlice n v
-                       Nothing -> v
-
--- | Drop the longest prefix of elements that satisfy the predicate and copy
--- the rest to a new vector.
+-- | Drop the longest prefix of elements that satisfy the predicate.
 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
 -- ---------
 
@@ -465,6 +526,95 @@ foldr1 :: IVector v a => (a -> a -> a) -> v a -> a
 {-# INLINE foldr1 #-}
 foldr1 f = Stream.foldr1 f . stream
 
+-- Specialised folds
+-- -----------------
+
+and :: IVector v Bool => v Bool -> Bool
+{-# INLINE and #-}
+and = Stream.and . stream
+
+or :: IVector v Bool => v Bool -> Bool
+{-# INLINE or #-}
+or = Stream.or . stream
+
+sum :: (IVector v a, Num a) => v a -> a
+{-# INLINE sum #-}
+sum = Stream.foldl' (+) 0 . stream
+
+product :: (IVector v a, Num a) => v a -> a
+{-# INLINE product #-}
+product = Stream.foldl' (*) 1 . stream
+
+maximum :: (IVector v a, Ord a) => v a -> a
+{-# INLINE maximum #-}
+maximum = Stream.foldl1' max . stream
+
+minimum :: (IVector v a, Ord a) => v a -> a
+{-# INLINE minimum #-}
+minimum = Stream.foldl1' min . stream
+
+-- Enumeration
+-- -----------
+
+enumFromTo :: (IVector v a, Enum a) => a -> a -> v a
+{-# INLINE enumFromTo #-}
+enumFromTo from to = from `seq` to `seq` unfoldr enumFromTo_go (fromEnum from)
+  where
+    to_i = fromEnum to
+    enumFromTo_go i | i <= to_i = Just (toEnum i, i + 1)
+                    | otherwise = Nothing
+
+enumFromThenTo :: (IVector v a, Enum a) => a -> a -> a -> v a
+{-# INLINE enumFromThenTo #-}
+enumFromThenTo from next to = from `seq` next `seq` to `seq` unfoldr enumFromThenTo_go from_i
+  where
+    from_i = fromEnum from
+    to_i = fromEnum to
+    step_i = fromEnum next - from_i
+    enumFromThenTo_go i | i <= to_i = Just (toEnum i, i + step_i)
+                        | otherwise = Nothing
+
+-- Unfolding
+-- ---------
+
+unfoldr :: IVector v a => (b -> Maybe (a, b)) -> b -> v a
+{-# INLINE unfoldr #-}
+unfoldr f = unstream . Stream.unfoldr f
+
+-- Scans
+-- -----
+
+-- | Prefix scan
+prescanl :: (IVector v a, IVector v b) => (a -> b -> a) -> a -> v b -> v a
+{-# INLINE prescanl #-}
+prescanl f z = unstream . Stream.prescanl f z . stream
+
+inplace_prescanl :: IVector v a => (a -> a -> a) -> a -> v a -> v a
+{-# INLINE inplace_prescanl #-}
+inplace_prescanl f z = unstream . inplace (MStream.prescanl f z) . stream
+
+{-# RULES
+
+"prescanl -> inplace_prescanl [IVector]" prescanl = inplace_prescanl
+
+ #-}
+
+-- | Prefix scan with strict accumulator
+prescanl' :: (IVector v a, IVector v b) => (a -> b -> a) -> a -> v b -> v a
+{-# INLINE prescanl' #-}
+prescanl' f z = unstream . Stream.prescanl' f z . stream
+
+inplace_prescanl' :: IVector v a => (a -> a -> a) -> a -> v a -> v a
+{-# INLINE inplace_prescanl' #-}
+inplace_prescanl' f z = unstream . inplace (MStream.prescanl' f z) . stream
+
+{-# RULES
+
+"prescanl' -> inplace_prescanl' [IVector]" prescanl' = inplace_prescanl'
+
+ #-}
+
+
 -- | Convert a vector to a list
 toList :: IVector v a => v a -> [a]
 {-# INLINE toList #-}