Use new Stream in length and null
authorRoman Leshchinskiy <rl@cse.unsw.edu.au>
Sun, 29 Jan 2012 10:53:46 +0000 (10:53 +0000)
committerRoman Leshchinskiy <rl@cse.unsw.edu.au>
Sun, 29 Jan 2012 10:53:46 +0000 (10:53 +0000)
Data/Vector/Fusion/Stream.hs
Data/Vector/Fusion/Stream/Monadic.hs
Data/Vector/Generic.hs

index 2bb0d34..9f4c80c 100644 (file)
@@ -142,12 +142,12 @@ sized = M.sized
 -- ------
 
 -- | Length of a 'Stream'
-length :: Stream v a -> Int
+length :: Vector v a => Stream v a -> Int
 {-# INLINE length #-}
 length = unId . M.length
 
 -- | Check if a 'Stream' is empty
-null :: Stream v a -> Bool
+null :: Vector v a => Stream v a -> Bool
 {-# INLINE null #-}
 null = unId . M.null
 
@@ -490,13 +490,18 @@ eq M.Stream{M.sElems = M.Unf step1 s1}
     eq_loop0 !sPEC s1 s2 = case unId (step1 s1) of
                              Yield x s1' -> eq_loop1 SPEC x s1' s2
                              Skip    s1' -> eq_loop0 SPEC   s1' s2
-                             Done        -> null (M.simple step2 s2 Unknown)
+                             Done        -> eq_null s2
 
     eq_loop1 !sPEC x s1 s2 = case unId (step2 s2) of
                                Yield y s2' -> x == y && eq_loop0 SPEC   s1 s2'
                                Skip    s2' ->           eq_loop1 SPEC x s1 s2'
                                Done        -> False
 
+    eq_null s2 = case unId (step2 s2) of
+                   Yield _ _ -> False
+                   Skip s2'  -> eq_null s2'
+                   Done      -> True
+
 -- | Lexicographically compare two 'Stream's
 cmp :: Ord a => Stream v a -> Stream v a -> Ordering
 {-# INLINE_STREAM cmp #-}
@@ -506,8 +511,7 @@ cmp M.Stream{M.sElems = M.Unf step1 s1}
     cmp_loop0 !sPEC s1 s2 = case unId (step1 s1) of
                               Yield x s1' -> cmp_loop1 SPEC x s1' s2
                               Skip    s1' -> cmp_loop0 SPEC   s1' s2
-                              Done        -> if null (M.simple step2 s2 Unknown)
-                                               then EQ else LT
+                              Done        -> cmp_null s2
 
     cmp_loop1 !sPEC x s1 s2 = case unId (step2 s2) of
                                 Yield y s2' -> case x `compare` y of
@@ -516,6 +520,11 @@ cmp M.Stream{M.sElems = M.Unf step1 s1}
                                 Skip    s2' -> cmp_loop1 SPEC x s1 s2'
                                 Done        -> GT
 
+    cmp_null s2 = case unId (step2 s2) of
+                    Yield _ _ -> LT
+                    Skip s2'  -> cmp_null s2'
+                    Done      -> EQ
+
 instance Eq a => Eq (M.Stream Id v a) where
   {-# INLINE (==) #-}
   (==) = eq
index f3a089f..857fd03 100644 (file)
@@ -169,15 +169,27 @@ sized s sz = s { sSize = sz }
 -- ------
 
 -- | Length of a 'Stream'
-length :: Monad m => Stream m v a -> m Int
+length :: (Monad m, Vector v a) => Stream m v a -> m Int
 {-# INLINE_STREAM length #-}
-length s = foldl' (\n _ -> n+1) 0 s
+length Stream{sVector = Just v} = return (basicLength v)
+length Stream{sSize = Exact n}  = return n
+length s = vfoldl' (\n (Chunk k _) -> n+k) 0 s
 
 -- | Check if a 'Stream' is empty
-null :: Monad m => Stream m v a -> m Bool
+null :: (Monad m, Vector v a) => Stream m v a -> m Bool
 {-# INLINE_STREAM null #-}
-null s = foldr (\_ _ -> False) True s
-
+null Stream{sVector = Just v} = return (basicLength v == 0)
+null Stream{sSize = Exact n} = return (n == 0)
+null Stream{sChunks = Unf step s} = null_loop s
+  where
+    null_loop s = do
+      r <- step s
+      case r of
+        Yield (Chunk n _) s'
+          | n /= 0    -> return False
+          | otherwise -> null_loop s'
+        Skip s'       -> null_loop s'
+        Done          -> return True
 
 -- Construction
 -- ------------
@@ -897,6 +909,10 @@ foldl' :: Monad m => (a -> b -> a) -> a -> Stream m v b -> m a
 {-# INLINE foldl' #-}
 foldl' f = foldlM' (\a b -> return (f a b))
 
+vfoldl' :: Monad m => (a -> Chunk v b -> a) -> a -> Stream m v b -> m a
+{-# INLINE vfoldl' #-}
+vfoldl' f = vfoldlM' (\a b -> return (f a b))
+
 -- | Left fold with a strict accumulator and a monadic operator
 foldlM' :: Monad m => (a -> b -> m a) -> a -> Stream m v b -> m a
 {-# INLINE_STREAM foldlM' #-}
index 06c5492..f3c2992 100644 (file)
@@ -212,27 +212,13 @@ mkNoRepType = mkNorepType
 
 -- | /O(1)/ Yield the length of the vector.
 length :: Vector v a => v a -> Int
-{-# INLINE_STREAM length #-}
-length v = basicLength v
-
-{-# RULES
-
-"length/unstream [Vector]" forall s.
-  length (new (New.unstream s)) = Stream.length s
-
-  #-}
+{-# INLINE length #-}
+length = Stream.length . stream
 
 -- | /O(1)/ Test whether a vector if empty
 null :: Vector v a => v a -> Bool
-{-# INLINE_STREAM null #-}
-null v = basicLength v == 0
-
-{-# RULES
-
-"null/unstream [Vector]" forall s.
-  null (new (New.unstream s)) = Stream.null s
-
-  #-}
+{-# INLINE null #-}
+null = Stream.null . stream
 
 -- Indexing
 -- --------