Avoid calling length on chunks in lazy splitAt

dolio · Bodigrim · commit fc86f1c833f9 · 2025-12-16T20:08:09.000Z
This instead uses the `measureOff` function used in the strict `splitAt`
to count only as many characters as are needed.
diff --git a/src/Data/Text/Lazy.hs b/src/Data/Text/Lazy.hs
@@ -1306,12 +1306,18 @@ splitAt = loop
     loop :: Int64 -> Text -> (Text, Text)
     loop !_ Empty     = (empty, empty)
     loop n t | n <= 0 = (empty, t)
-    loop n (Chunk t ts)
-         | n < len   = let (t',t'') = T.splitAt (int64ToInt n) t
-                       in (Chunk t' Empty, Chunk t'' ts)
-         | otherwise = let (ts',ts'') = loop (n - len) ts
+    loop n (Chunk t@(T.Text arr off len) ts)
+         | m > 0, m >= len = (Chunk t Empty, ts)
+         | m > 0 = let t' = T.Text arr off m
+                       t'' = T.Text arr (off+m) (len-m)
+                   in (Chunk t' Empty, Chunk t'' ts)
+         | otherwise = let (ts', ts'') = loop (n + intToInt64 m) ts
                        in (Chunk t ts', ts'')
-         where len = intToInt64 (T.length t)
+         where
+         k | n > intToInt64 len = len+1
+           | otherwise = int64ToInt n
+         m = T.measureOff k t
+
 
 -- | /O(n)/ 'splitAtWord' @n t@ returns a strict pair whose first
 -- element is a prefix of @t@ whose chunks contain @n@ 'Word8'
