Commits

Bryan O'Sullivan committed 55ac947

Implement count.

  • Participants
  • Parent commits 9e2cdf4

Comments (0)

Files changed (3)

File Data/Text.hs

     , findIndices
     , elemIndex
     , elemIndices
-    -- , count
+    , count
 
     -- * Zipping and unzipping
     , zipWith
     | n >= len  = t
     | otherwise = Text arr off (loop 0 0)
   where
-      loop !i !count
-           | i >= len || count >= n = i
-           | otherwise              = loop (i+d) (count+1)
+      loop !i !cnt
+           | i >= len || cnt >= n = i
+           | otherwise            = loop (i+d) (cnt+1)
            where d = iter_ t i
 {-# INLINE [1] take #-}
 
     | n >= len  = empty
     | otherwise = loop 0 0
   where end = off + len
-        loop !i !count
-            | i >= end || count >= n   = Text arr (off+i) (len-i)
-            | otherwise                = loop (i+d) (count+1)
+        loop !i !cnt
+            | i >= end || cnt >= n   = Text arr (off+i) (len-i)
+            | otherwise              = loop (i+d) (cnt+1)
             where d = iter_ t i
 {-# INLINE [1] drop #-}
 
     | n >= len  = (t, empty)
     | otherwise = (Text arr off k, Text arr (off+k) (len-k))
   where k = loop 0 0
-        loop !i !count
-            | i >= len || count >= n = i
-            | otherwise              = loop (i+d) (count+1)
-            where d                  = iter_ t i
+        loop !i !cnt
+            | i >= len || cnt >= n = i
+            | otherwise            = loop (i+d) (cnt+1)
+            where d                = iter_ t i
 {-# INLINE splitAt #-}
 
 -- | /O(n)/ 'span', applied to a predicate @p@ and text @t@, returns a
 elemIndices c t = S.elemIndices c (stream t)
 {-# INLINE elemIndices #-}
 
+-- | /O(n)/ The 'count' function returns the number of times the query
+-- element appears in the given 'Text'. This function is subject to
+-- fusion.
+count :: Char -> Text -> Int
+count c t = S.count c (stream t)
+{-# INLINE count #-}
+
 -------------------------------------------------------------------------------
 -- * Zipping
 

File Data/Text/Fusion.hs

     , findIndexOrEnd
     , elemIndex
     , elemIndices
+    , count
 
     -- * Zipping and unzipping
     , zipWith
                  | otherwise -> loop (i+1) s'
 {-# INLINE [0] elemIndices #-}
 
+-- | /O(n)/ The 'count' function returns the number of times the query
+-- element appears in the given stream.
+count :: Char -> Stream Char -> Int
+count a (Stream next s0 _len) = loop 0 s0
+  where
+    loop !i !s = case next s of
+      Done                   -> i
+      Skip    s'             -> loop i s'
+      Yield x s' | a == x    -> loop (i+1) s'
+                 | otherwise -> loop i s'
+{-# INLINE [0] count #-}
+
 -------------------------------------------------------------------------------
 -- * Zipping
 

File tests/Properties.hs

 prop_findIndices p   = L.findIndices p`eqP` T.findIndices p
 prop_elemIndex c     = L.elemIndex c `eqP` T.elemIndex c
 prop_elemIndices c   = L.elemIndices c`eqP` T.elemIndices c
+prop_count c         = (L.length . L.elemIndices c) `eqP` T.count c
 prop_zipWith c s     = L.zipWith c s `eqP` (unpack . T.zipWith c (pack s))
 
 main = run tests =<< getArgs
   ("prop_findIndices", mytest prop_findIndices),
   ("prop_elemIndex", mytest prop_elemIndex),
   ("prop_elemIndices", mytest prop_elemIndices),
+  ("prop_count", mytest prop_count),
   ("prop_zipWith", mytest prop_zipWith)
   ]