Commits

Bryan O'Sullivan committed eaee653

Add strict and lazy count functions.

Comments (0)

Files changed (3)

 {-# INLINE elemIndices #-}
 
 -- | /O(n*m)/ The 'count' function returns the number of times the
--- query string appears in the given 'Text'.
+-- query string appears in the given 'Text'. An empty query string is
+-- invalid, and will cause an error to be raised.
 count :: Text -> Text -> Int
 count pat src0
-    | l == 0    = length src0 + 1
+    | l == 0    = emptyError "count"
     | l == 1    = countChar (unsafeHead pat) src0
     | otherwise = go 0 src0
   where

Data/Text/Lazy.hs

     , findIndices
     , elemIndex
     , elemIndices
+    , count
     , countChar
 
     -- * Zipping and unzipping
 elemIndices c t = S.elemIndices c (stream t)
 {-# INLINE elemIndices #-}
 
+-- | /O(n*m)/ The 'count' function returns the number of times the
+-- query string appears in the given 'Text'. An empty query string is
+-- invalid, and will cause an error to be raised.
+count :: Text -> Text -> Int64
+count pat src0
+    | l == 0    = emptyError "count"
+    | l == 1    = countChar (head pat) src0
+    | otherwise = go 0 src0
+  where
+    l = length pat
+    go !n src = search src
+      where
+        search s | null s             = n
+                 | pat `isPrefixOf` s = go (n+1) (drop l s)
+                 | otherwise          = search (tail s)
+{-# INLINE [1] count #-}
+
 -- | /O(n)/ The 'countChar' function returns the number of times the
 -- query element appears in the given 'Text'. This function is subject
 -- to fusion.

tests/Properties.hs

 sf_elemIndices p c= (L.elemIndices c . L.filter p) `eqP` (S.elemIndices c . S.filter p)
 t_elemIndices c   = L.elemIndices c`eqP` T.elemIndices c
 tl_elemIndices c  = (fmap fromIntegral . L.elemIndices c) `eqP` TL.elemIndices c
-sf_countChar p c      = (L.length . L.elemIndices c . L.filter p) `eqP` (S.countChar c . S.filter p)
-t_countChar c         = (L.length . L.elemIndices c) `eqP` T.countChar c
-tl_countChar c        = (fromIntegral . L.length . L.elemIndices c) `eqP` TL.countChar c
+t_count t         = (subtract 1 . L.length . T.split t) `eq` T.count t
+tl_count t        = (subtract 1 . L.genericLength . TL.split t) `eq` TL.count t
+sf_countChar p c  = (L.length . L.elemIndices c . L.filter p) `eqP` (S.countChar c . S.filter p)
+t_countChar c     = (L.length . L.elemIndices c) `eqP` T.countChar c
+tl_countChar c    = (fromIntegral . L.length . L.elemIndices c) `eqP` TL.countChar c
 t_zip s           = L.zip s `eqP` T.zip (packS s)
 tl_zip s          = L.zip s `eqP` TL.zip (packS s)
 sf_zipWith p c s  = (L.zipWith c (L.filter p s) . L.filter p) `eqP` (unpackS . S.zipWith c (S.filter p $ packS s) . S.filter p)
     testProperty "sf_elemIndices" sf_elemIndices,
     testProperty "t_elemIndices" t_elemIndices,
     testProperty "tl_elemIndices" tl_elemIndices,
+    testProperty "t_count" t_count,
+    testProperty "tl_count" tl_count,
     testProperty "sf_countChar" sf_countChar,
     testProperty "t_countChar" t_countChar,
     testProperty "tl_countChar" tl_countChar