Add (!?) for streams and make vector (!?) fusible
authorRoman Leshchinskiy <rl@cse.unsw.edu.au>
Thu, 18 Aug 2011 22:21:15 +0000 (22:21 +0000)
committerRoman Leshchinskiy <rl@cse.unsw.edu.au>
Thu, 18 Aug 2011 22:21:15 +0000 (22:21 +0000)
Data/Vector/Fusion/Stream.hs
Data/Vector/Fusion/Stream/Monadic.hs
Data/Vector/Generic.hs

index 52998d1..18bbabf 100644 (file)
@@ -29,7 +29,7 @@ module Data.Vector.Fusion.Stream (
   empty, singleton, cons, snoc, replicate, generate, (++),
 
   -- * Accessing individual elements
-  head, last, (!!),
+  head, last, (!!), (!?),
 
   -- * Substreams
   slice, init, tail, take, drop,
@@ -204,6 +204,12 @@ last = unId . M.last
 {-# INLINE (!!) #-}
 s !! i = unId (s M.!! i)
 
+infixl 9 !?
+-- | Element at the given position or 'Nothing' if out of bounds
+(!?) :: Stream a -> Int -> Maybe a
+{-# INLINE (!?) #-}
+s !? i = unId (s M.!? i)
+
 -- Substreams
 -- ----------
 
index 55cd2d7..b781e58 100644 (file)
@@ -25,7 +25,7 @@ module Data.Vector.Fusion.Stream.Monadic (
   empty, singleton, cons, snoc, replicate, replicateM, generate, generateM, (++),
 
   -- * Accessing elements
-  head, last, (!!),
+  head, last, (!!), (!?),
 
   -- * Substreams
   slice, init, tail, take, drop,
@@ -271,6 +271,22 @@ Stream step s _ !! i | i < 0     = BOUNDS_ERROR(error) "!!" "negative index"
             Skip    s'             -> index_loop SPEC s' i
             Done                   -> BOUNDS_ERROR(emptyStream) "!!"
 
+infixl 9 !?
+-- | Element at the given position or 'Nothing' if out of bounds
+(!?) :: Monad m => Stream m a -> Int -> m (Maybe a)
+{-# INLINE (!?) #-}
+Stream step s _ !? i = index_loop SPEC s i
+  where
+    index_loop !sPEC s i
+      = i `seq`
+        do
+          r <- step s
+          case r of
+            Yield x s' | i == 0    -> return (Just x)
+                       | otherwise -> index_loop SPEC s' (i-1)
+            Skip    s'             -> index_loop SPEC s' i
+            Done                   -> return Nothing
+
 -- Substreams
 -- ----------
 
index fd97d42..e9916c1 100644 (file)
@@ -237,6 +237,7 @@ null v = basicLength v == 0
 v ! i = BOUNDS_CHECK(checkIndex) "(!)" i (length v)
       $ unId (basicUnsafeIndexM v i)
 
+infixl 9 !?
 -- | O(1) Safe indexing
 (!?) :: Vector v a => v a -> Int -> Maybe a
 {-# INLINE_STREAM (!?) #-}
@@ -274,6 +275,9 @@ unsafeLast v = unsafeIndex v (length v - 1)
 "(!)/unstream [Vector]" forall i s.
   new (New.unstream s) ! i = s Stream.!! i
 
+"(!?)/unstream [Vector]" forall i s.
+  new (New.unstream s) !? i = s Stream.!? i
+
 "head/unstream [Vector]" forall s.
   head (new (New.unstream s)) = Stream.head s