@@ -36,6 +36,7 @@ import qualified Data.Text.Array as A
3636import Data.Int (Int64)
3737import Data.Word (Word16, Word64)
3838import qualified Data.Text.Internal as T
39+import qualified Data.Text as T
3940import Data.Text.Fusion.Internal (PairS(..))
4041import Data.Text.Lazy.Internal (Text(..))
4142import Data.Bits ((.|.), (.&.))
@@ -49,12 +50,12 @@ import Data.Text.UnsafeShift (shiftL)
4950indices :: Text -- ^ Substring to search for (@needle@)
5051 -> Text -- ^ Text to search in (@haystack@)
5152 -> [Int64]
52−indices needle Empty = []
53−indices needle@(Chunk n ns) haystack@(Chunk c cs)
53+indices needle@(Chunk n ns) haystack@(Chunk k ks)
5454 | nlen <= 0 || ldiff < 0 = []
55− | otherwise = scan 0 0 c cs
55+ | nlen == 1 = scanOne (nindex 0) 0 k ks
56+ | otherwise = scan 0 0 k ks
5657 where
57− scan !g !i x@(T.Text harr hoff l) xs
58+ scan !g !i x@(T.Text _ _ l) xs
5859 | g > ldiff = []
5960 | i >= m = case xs of
6061 Empty -> []
@@ -79,17 +80,28 @@ indices needle@(Chunk n ns) haystack@(Chunk c cs)
7980 nlast = nlen - 1
8081 nindex = index n ns
8182 z = foldChunks fin 0 needle
82− where fin _ (T.Text narr noff nlen) = A.unsafeIndex narr (noff+nlen-1)
83+ where fin _ (T.Text farr foff flen) = A.unsafeIndex farr (foff+flen-1)
8384 mask :*: skip = buildTable needle
85+ scanOne c i (T.Text oarr ooff olen) os = go 0
86+ where
87+ go h | h >= olen = case os of
88+ Empty -> []
89+ Chunk y ys -> scanOne c (i+fromIntegral olen) y ys
90+ | on == c = i + fromIntegral h : go (h+1)
91+ | otherwise = go (h+1)
92+ where on = A.unsafeIndex oarr (ooff+h)
93+indices _ _ = []
8494
8595index :: T.Text -> Text -> Int64 -> Word16
8696index (T.Text arr off len) xs i
87− | j <= len = A.unsafeIndex arr (off+j)
97+ | j < len = A.unsafeIndex arr (off+j)
8898 | otherwise = case xs of
89− Empty -> error "empty"
99+ Empty | j == len -> 0
100+ | otherwise -> error "empty"
90101 Chunk c cs -> index c cs (i-fromIntegral len)
91102 where j = fromIntegral i
92103
104+swizzle :: Word16 -> Word64
93105swizzle k = 1 `shiftL` (fromIntegral k .&. 0x3f)
94106
95107foldChunks :: (a -> T.Text -> a) -> a -> Text -> a
@@ -97,23 +109,25 @@ foldChunks _ z Empty = z
97109foldChunks f z (Chunk c cs) = let z' = f z c
98110 in z' `seq` foldChunks f z' cs
99111
112+wordLength :: Text -> Int64
100113wordLength = foldChunks sumLength 0
101114 where sumLength i (T.Text _ _ l) = i + fromIntegral l
102115
103116buildTable :: Text -> PairS Word64 Int64
104117buildTable Empty = 0 :*: 0
105−buildTable needle@(Chunk c cs) = outer c cs 0 0 (nlen-2)
118+buildTable needle@(Chunk k ks) = outer k ks 0 0 0 (nlen-2)
106119 where
107− outer x@(T.Text xarr xoff xlen) xs = go
120+ outer (T.Text xarr xoff xlen) xs = go
108121 where
109− go !i !mask !skip
110− | i >= xlen = case xs of
111− Empty -> (mask .|. swizzle z) :*: skip
112− Chunk y ys -> outer y ys 0 mask skip
113− | otherwise = go (i+1) (mask .|. swizzle c) skip'
122+ go !(g::Int64) !i !mask !skip
123+ | i >= xlast = case xs of
124+ Empty -> (mask .|. swizzle z) :*: skip
125+ Chunk y ys -> outer y ys g 0 mask skip
126+ | otherwise = go (g+1) (i+1) (mask .|. swizzle c) skip'
114127 where c = A.unsafeIndex xarr (xoff+i)
115− skip' | c == z = nlen - fromIntegral i - 2
128+ skip' | c == z = nlen - fromIntegral g - 2
116129 | otherwise = skip
130+ xlast = xlen - 1
117131 nlen = wordLength needle
118132 z = foldChunks fin 0 needle
119− where fin _ (T.Text narr noff nlen) = A.unsafeIndex narr (noff+nlen-1)
133+ where fin _ (T.Text farr foff flen) = A.unsafeIndex farr (foff+flen-1)