Commits

Bryan O'Sullivan committed 27e0e98

Implement splitAt

Comments (0)

Files changed (2)

     , drop
     , takeWhile
     , dropWhile
-    -- , splitAt
+    , splitAt
     -- , span
     -- , break
     -- , group
     unstream (S.dropWhile p (stream t)) = dropWhile p t
   #-}
 
+-- | /O(n)/ 'splitAt' @n t@ returns a pair whose first element is a
+-- prefix of @t@ of length @n@ and second element is the remainder of
+-- the string. It is equivalent to @('take' n t, 'drop' n t)@.
+splitAt :: Int -> Text -> (Text, Text)
+splitAt n t@(Text arr off len)
+    | n <= 0    = (empty, t)
+    | n >= len  = (t, empty)
+    | otherwise = (Text arr off k, Text arr (off+k) (len-k))
+  where k = loop off 0
+        end = off + len
+        loop !i !count
+            | i >= end || count >= n   = i - off
+            | otherwise                = loop (i+d) (count+1)
+            where d = iter_ arr i
+{-# INLINE splitAt #-}
+
 -- | /O(n)/ Return all initial segments of the given 'Text', shortest
 -- first.
 inits :: Text -> [Text]

tests/Properties.hs

 prop_takeWhileS p    = L.takeWhile p `eqP` (unpack . unstream . S.takeWhile p . stream)
 prop_dropWhile p     = L.dropWhile p `eqP` (unpack . T.dropWhile p)
 prop_dropWhileS p    = L.dropWhile p `eqP` (unpack . unstream . S.dropWhile p . stream)
+prop_splitAt n       = L.splitAt n   `eqP` ((unpack *** unpack) . T.splitAt n)
 prop_inits           = L.inits       `eqP` (map unpack . T.inits)
 prop_tails           = L.tails       `eqP` (map unpack . T.tails)
 
   ("prop_takeWhileS", mytest prop_takeWhileS),
   ("prop_dropWhile", mytest prop_dropWhile),
   ("prop_dropWhileS", mytest prop_dropWhileS),
+  ("prop_splitAt", mytest prop_splitAt),
   ("prop_inits", mytest prop_inits),
   ("prop_tails", mytest prop_tails),