Implement index via measureOff

Author: Bodigrim

src/Data/Text.hs

@@ -1845,7 +1845,15 @@ breakOnAll pat src@(Text arr off slen)
 
 -- | /O(n)/ 'Text' index (subscript) operator, starting from 0.
 index :: HasCallStack => Text -> Int -> Char
-index t n = S.index (stream t) n
+index t@(Text _ _ lenInBytes) ix
+  | ix < 0
+  = P.error $ "Data.Text.index: negative index " ++ P.show ix
+  | off < 0 || off == lenInBytes
+  = P.error $ "Data.Text.index: index " ++ P.show ix ++ " is too large"
+  | otherwise = ch
+  where
+    off = measureOff ix t
+    Iter ch _ = iter t off
 {-# INLINE index #-}
 
 -- | /O(n)/ The 'findIndex' function takes a predicate and a 'Text'

src/Data/Text/Lazy.hs

@@ -223,7 +223,7 @@ import Prelude (Char, Bool(..), Maybe(..), String,
 import qualified Prelude as P
 import Control.Arrow (first)
 import Control.DeepSeq (NFData(..))
-import Data.Bits (finiteBitSize)
+import Data.Bits (finiteBitSize, toIntegralSized)
 import Data.Int (Int64)
 import qualified Data.List as L hiding (head, tail)
 import Data.Char (isSpace)
@@ -1808,7 +1808,22 @@ partition p t = (filter p t, filter (not . p) t)
 
 -- | /O(n)/ 'Text' index (subscript) operator, starting from 0.
 index :: HasCallStack => Text -> Int64 -> Char
-index t n = S.index (stream t) n
+index lazyText ix
+  | ix < 0 = P.error $ "Data.Text.Lazy.index: negative index " ++ P.show ix
+  | otherwise = go lazyText ix
+  where
+    go :: Text -> Int64 -> Char
+    go Empty _ = P.error $ "Data.Text.index: index " ++ P.show ix ++ " is too large"
+    go (Chunk t@(T.Text _ _ lenInBytes) ts) n = case toIntegralSized n of
+      Nothing ->
+        go ts (n - fromIntegral (T.length t))
+      Just n'
+        | off < 0 -> go ts (n + fromIntegral off)
+        | off == lenInBytes -> go ts 0
+        | otherwise -> ch
+        where
+          off = T.measureOff n' t
+          T.Iter ch _ = T.iter t off
 {-# INLINE index #-}
 
 -- | /O(n+m)/ The 'count' function returns the number of times the

tests/Tests/Properties/Text.hs

@@ -1,15 +1,19 @@
 -- | Tests for operations that don't fit in the other @Test.Properties.*@ modules.
 
-{-# LANGUAGE BangPatterns #-}
-{-# LANGUAGE CPP          #-}
-{-# LANGUAGE ViewPatterns #-}
+{-# LANGUAGE BangPatterns        #-}
+{-# LANGUAGE CPP                 #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE ViewPatterns        #-}
 
 {-# OPTIONS_GHC  -fno-warn-missing-signatures #-}
+{-# OPTIONS_GHC -Wno-unrecognised-pragmas #-}
+{-# HLINT ignore "Avoid restricted function" #-}
 
 module Tests.Properties.Text
     ( testText
     ) where
 
+import Control.Exception (SomeException, evaluate, try)
 import Data.Char (isLower, isLetter, isUpper)
 import Data.Maybe (mapMaybe)
 import Data.Text.Internal.Fusion.Size
@@ -263,13 +267,21 @@ tl_partition (applyFun -> p) = L.partition p `eqP` (unpack2 . TL.partition p)
 sf_index (applyFun -> p) s i = ((L.filter p s L.!!) `eq` S.index (S.filter p $ packS s)) j
     where l = L.length s
           j = if l == 0 then 0 else i `mod` (3 * l) - l
-t_index s i       = ((s L.!!) `eq` T.index (packS s)) j
-    where l = L.length s
-          j = if l == 0 then 0 else i `mod` (3 * l) - l
 
-tl_index s i      = ((s L.!!) `eq` (TL.index (packS s) . fromIntegral)) j
-    where l = L.length s
-          j = if l == 0 then 0 else i `mod` (3 * l) - l
+t_index :: T.Text -> Int -> Property
+t_index xs i = ioProperty $ do
+    ch <- try (evaluate (T.index xs i))
+    pure $ case ch of
+        Left (_ :: SomeException) -> i < 0 .||. i >= T.length xs
+        Right c -> i >= 0 .&&. i < T.length xs .&&. c === T.unpack xs L.!! i
+
+tl_index :: TL.Text -> Int -> Property
+tl_index xs i = ioProperty $ do
+    let i' = fromIntegral i
+    ch <- try (evaluate (TL.index xs i'))
+    pure $ case ch of
+        Left (_ :: SomeException) -> i' < 0 .||. i' >= TL.length xs
+        Right c -> i >= 0 .&&. i' < TL.length xs .&&. c === TL.unpack xs L.!! i
 
 t_findIndex (applyFun -> p) = L.findIndex p `eqP` T.findIndex p
 t_count (NotEmpty t)  = (subtract 1 . L.length . T.splitOn t) `eq` T.count t