Add D.V.IVector.{indexM,headM,lastM}
authorRoman Leshchinskiy <rl@cse.unsw.edu.au>
Thu, 10 Sep 2009 10:20:04 +0000 (10:20 +0000)
committerRoman Leshchinskiy <rl@cse.unsw.edu.au>
Thu, 10 Sep 2009 10:20:04 +0000 (10:20 +0000)
Data/Vector/IVector.hs

index a8a257c..52dcfd3 100644 (file)
@@ -25,7 +25,7 @@ module Data.Vector.IVector (
   empty, singleton, cons, snoc, replicate, (++), copy,
 
   -- * Accessing individual elements
   empty, singleton, cons, snoc, replicate, (++), copy,
 
   -- * Accessing individual elements
-  (!), head, last,
+  (!), head, last, indexM, headM, lastM,
 
   -- * Subvectors
   slice, init, tail, take, drop,
 
   -- * Subvectors
   slice, init, tail, take, drop,
@@ -305,6 +305,34 @@ last v = v ! (length v - 1)
 
  #-}
 
 
  #-}
 
+-- | Monadic indexing which can be strict in the array while remaining lazy in
+-- the element.
+indexM :: (IVector v a, Monad m) => v a -> Int -> m a
+{-# INLINE_STREAM indexM #-}
+indexM v i = assert (i >= 0 && i < length v)
+           $ unsafeIndex v i return
+
+headM :: (IVector v a, Monad m) => v a -> m a
+{-# INLINE_STREAM headM #-}
+headM v = indexM v 0
+
+lastM :: (IVector v a, Monad m) => v a -> m a
+{-# INLINE_STREAM lastM #-}
+lastM v = indexM v (length v - 1)
+
+{-# RULES
+
+"indexM/unstream [IVector]" forall v i s.
+  indexM (new' v (New.unstream s)) i = return (s Stream.!! i)
+
+"headM/unstream [IVector]" forall v s.
+  headM (new' v (New.unstream s)) = return (Stream.head s)
+
+"lastM/unstream [IVector]" forall v s.
+  lastM (new' v (New.unstream s)) = return (Stream.last s)
+
+ #-}
+
 -- Subarrays
 -- ---------
 
 -- Subarrays
 -- ---------