diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,13 @@ # Revision history for byteslice +## 0.2.8.0 -- 2022-11-21 + +* Add `Data.Bytes.replace` and `Data.Bytes.findIndices` +* Add `Data.Bytes.fromShortText` +* Add `Data.Bytes.fromPrimArray` +* Add `decodeDecWord` +* Add `concatArray` and `concatArrayU`. + ## 0.2.7.0 -- 2022-02-16 * Add support for GHC 9.2. diff --git a/bench/Main.hs b/bench/Main.hs --- a/bench/Main.hs +++ b/bench/Main.hs @@ -1,7 +1,9 @@ -import Gauge.Main (defaultMain, bench, whnf) +import Gauge.Main (defaultMain, bench, whnf, bgroup) import Data.Bytes.Types import Data.List (permutations) + +import qualified Data.Bytes as Bytes import qualified Data.Bytes.Text.Ascii as Ascii naiveMconcat :: [Bytes] -> Bytes @@ -11,7 +13,33 @@ main = defaultMain [ bench "mconcat" $ whnf mconcat mconcatBytes , bench "naiveMconcat" $ whnf naiveMconcat mconcatBytes + , bgroup "replace" + [ bench "the-dog-and-the-shadow" (whnf replaceMeat theDogAndTheShadow) + ] ] mconcatBytes :: [Bytes] mconcatBytes = fmap Ascii.fromString $ permutations ['a'..'g'] + +replaceMeat :: Bytes -> Bytes +{-# noinline replaceMeat #-} +replaceMeat x = Bytes.replace meat gruel x + +meat :: Bytes +meat = Ascii.fromString "meat" + +gruel :: Bytes +gruel = Ascii.fromString "gruel" + +theDogAndTheShadow :: Bytes +theDogAndTheShadow = Ascii.fromString $ concat + [ "It happened that a Dog had got a piece of meat and was " + , "carrying it home in his mouth to eat it in peace. " + , "Now on his way home he had to cross a plank lying across " + , "a running brook. As he crossed, he looked down and saw his " + , "own shadow reflected in the water beneath. Thinking it was " + , "another dog with another piece of meat, he made up his mind " + , "to have that also. So he made a snap at the shadow in the water, " + , "but as he opened his mouth the piece of meat fell out, dropped " + , "into the water and was never seen more." + ] diff --git a/byteslice.cabal b/byteslice.cabal --- a/byteslice.cabal +++ b/byteslice.cabal @@ -1,6 +1,6 @@ cabal-version: 2.4 name: byteslice -version: 0.2.7.0 +version: 0.2.8.0 synopsis: Slicing managed and unmanaged memory description: This library provides types that allow the user to talk about a slice of @@ -34,6 +34,7 @@ Data.Bytes.Text.Windows1252 Data.Bytes.Types other-modules: + Data.Bytes.Search Data.Bytes.Byte Data.Bytes.Pure Data.Bytes.IO @@ -46,8 +47,9 @@ , primitive-unlifted >=0.1.2 && <0.2 , primitive-addr >=0.1 && <0.2 , run-st >=0.1.1 && <0.2 + , text-short >=0.1.3 && <0.2 , tuples >=0.1 && <0.2 - , vector >=0.12 && <0.13 + , vector >=0.12 && <0.14 hs-source-dirs: src ghc-options: -Wall -O2 if impl(ghc>=9.2) @@ -82,6 +84,7 @@ , tasty-hunit , tasty-quickcheck , transformers + , text >=1.2 benchmark bench type: exitcode-stdio-1.0 diff --git a/src/Data/Bytes.hs b/src/Data/Bytes.hs --- a/src/Data/Bytes.hs +++ b/src/Data/Bytes.hs @@ -23,7 +23,7 @@ , Pure.emptyPinned , Pure.emptyPinnedU -- * Properties - , null + , Pure.null , Pure.length -- * Decompose , uncons @@ -47,6 +47,9 @@ , dropWhile , takeWhileEnd , dropWhileEnd + -- * Traversals + , Pure.map + , Pure.mapU -- * Folds , Pure.foldl , Pure.foldl' @@ -77,6 +80,11 @@ -- * Combining , intercalate , intercalateByte2 + , concatArray + , concatArrayU + -- * Searching + , replace + , findIndices -- * Counting , Byte.count -- * Prefix and Suffix @@ -114,9 +122,10 @@ , Pure.fnv1a32 , Pure.fnv1a64 -- * Unsafe Slicing - , unsafeTake - , unsafeDrop - , unsafeIndex + , Pure.unsafeTake + , Pure.unsafeDrop + , Pure.unsafeIndex + , Pure.unsafeHead -- * Copying , Pure.unsafeCopy -- * Pointers @@ -131,12 +140,14 @@ , fromAsciiString , fromLatinString , Pure.fromByteArray + , Pure.fromPrimArray , toLatinString , fromCString# , Pure.toByteString , Pure.pinnedToByteString , Pure.fromByteString , fromShortByteString + , fromShortText , toShortByteString , toShortByteStringClone , toLowerAsciiByteArrayClone @@ -149,23 +160,24 @@ , unlift ) where -import Prelude hiding (length,takeWhile,dropWhile,null,foldl,foldr,elem,replicate,any,all,readFile) +import Prelude hiding (length,takeWhile,dropWhile,null,foldl,foldr,elem,replicate,any,all,readFile,map) import Control.Monad.Primitive (PrimMonad,primitive_,unsafeIOToPrim) import Control.Monad.ST.Run (runByteArrayST) import Cstrlen (cstringLength#) -import Data.Bits((.&.),(.|.),shiftL,finiteBitSize) -import Data.Bytes.Pure (length,fromByteArray,foldr) -import Data.Bytes.Types (Bytes(Bytes,array,offset)) import Data.ByteString.Short.Internal (ShortByteString(SBS)) -import Data.Maybe (fromMaybe) -import Data.Primitive (ByteArray(ByteArray)) +import Data.Bytes.Pure (length,fromByteArray,foldr,unsafeDrop) +import Data.Bytes.Pure (unsafeIndex) +import Data.Bytes.Types (Bytes(Bytes,array,offset)) +import Data.Primitive (Array,ByteArray(ByteArray)) +import Data.Text.Short (ShortText) import Foreign.C.String (CString) import Foreign.Ptr (Ptr,plusPtr,castPtr) import GHC.Exts (Addr#,Word#,Int#) import GHC.Exts (Int(I#),Ptr(Ptr)) -import GHC.Word (Word8(W8#),Word32) +import GHC.Word (Word8(W8#)) import Reps (Bytes#(..),word8ToWord#) +import Data.Bytes.Search (findIndices,replace,isInfixOf) import qualified Data.Bytes.Byte as Byte import qualified Data.Bytes.Chunks as Chunks @@ -179,13 +191,9 @@ import qualified Data.List as List import qualified Data.Primitive as PM import qualified Data.Primitive.Ptr as PM +import qualified Data.Text.Short as TS import qualified GHC.Exts as Exts --- | Is the byte sequence empty? -null :: Bytes -> Bool -{-# inline null #-} -null (Bytes _ _ len) = len == 0 - -- | Extract the head and tail of the 'Bytes', returning 'Nothing' if -- it is empty. uncons :: Bytes -> Maybe (Word8, Bytes) @@ -239,67 +247,6 @@ then compareByteArrays a aOff b (bOff + bLen - aLen) aLen == EQ else False --- | Is the first argument an infix of the second argument? --- --- Uses the Rabin-Karp algorithm: expected time @O(n+m)@, worst-case @O(nm)@. -isInfixOf :: Bytes -- ^ String to search for - -> Bytes -- ^ String to search in - -> Bool -isInfixOf p s = null p || (not . null) (snd $ breakSubstring p s) - -breakSubstring :: Bytes -- ^ String to search for - -> Bytes -- ^ String to search in - -> (Bytes,Bytes) -- ^ Head and tail of string broken at substring -breakSubstring pat = - case lp of - 0 -> (mempty,) - 1 -> breakByte (unsafeHead pat) - _ -> if lp * 8 <= finiteBitSize (0 :: Word) - then shift - else karpRabin - where - unsafeSplitAt i s = (unsafeTake i s, unsafeDrop i s) - lp = length pat - {-# INLINE breakByte #-} - breakByte b bytes = fromMaybe (mempty,bytes) $ Byte.split1 b bytes - {-# INLINE karpRabin #-} - karpRabin :: Bytes -> (Bytes, Bytes) - karpRabin src - | length src < lp = (src,mempty) - | otherwise = search (rollingHash $ unsafeTake lp src) lp - where - k = 2891336453 :: Word32 - rollingHash = Pure.foldl' (\h b -> h * k + fromIntegral b) 0 - hp = rollingHash pat - m = k ^ lp - get = fromIntegral . unsafeIndex src - search !hs !i - | hp == hs && pat == unsafeTake lp b = u - | length src <= i = (src,mempty) -- not found - | otherwise = search hs' (i + 1) - where - u@(_, b) = unsafeSplitAt (i - lp) src - hs' = hs * k + - get i - - m * get (i - lp) - {-# INLINE shift #-} - shift :: Bytes -> (Bytes, Bytes) - shift !src - | length src < lp = (src,mempty) - | otherwise = search (intoWord $ unsafeTake lp src) lp - where - intoWord :: Bytes -> Word - intoWord = Pure.foldl' (\w b -> (w `shiftL` 8) .|. fromIntegral b) 0 - wp = intoWord pat - mask = (1 `shiftL` (8 * lp)) - 1 - search !w !i - | w == wp = unsafeSplitAt (i - lp) src - | length src <= i = (src, mempty) - | otherwise = search w' (i + 1) - where - b = fromIntegral (unsafeIndex src i) - w' = mask .&. ((w `shiftL` 8) .|. b) - -- | Find the longest string which is a prefix of both arguments. longestCommonPrefix :: Bytes -> Bytes -> Bytes longestCommonPrefix a b = loop 0 @@ -309,7 +256,7 @@ | into < maxLen && unsafeIndex a into == unsafeIndex b into = loop (into + 1) - | otherwise = unsafeTake into a + | otherwise = Pure.unsafeTake into a maxLen = min (length a) (length b) -- | Create a byte sequence with one byte. @@ -412,30 +359,19 @@ -- | Take bytes while the predicate is true. takeWhile :: (Word8 -> Bool) -> Bytes -> Bytes {-# inline takeWhile #-} -takeWhile k b = unsafeTake (countWhile k b) b +takeWhile k b = Pure.unsafeTake (countWhile k b) b -- | Drop bytes while the predicate is true. dropWhile :: (Word8 -> Bool) -> Bytes -> Bytes {-# inline dropWhile #-} -dropWhile k b = unsafeDrop (countWhile k b) b - --- | Index into the byte sequence at the given position. This index --- must be less than the length. -unsafeIndex :: Bytes -> Int -> Word8 -{-# inline unsafeIndex #-} -unsafeIndex (Bytes arr off _) ix = PM.indexByteArray arr (off + ix) - --- | Access the first byte. The given 'Bytes' must be non-empty. -{-# inline unsafeHead #-} -unsafeHead :: Bytes -> Word8 -unsafeHead bs = unsafeIndex bs 0 +dropWhile k b = Pure.unsafeDrop (countWhile k b) b -- | /O(n)/ 'dropWhileEnd' @p@ @b@ returns the prefix remaining after -- dropping characters that satisfy the predicate @p@ from the end of -- @t@. dropWhileEnd :: (Word8 -> Bool) -> Bytes -> Bytes {-# inline dropWhileEnd #-} -dropWhileEnd k !b = unsafeTake (length b - countWhileEnd k b) b +dropWhileEnd k !b = Pure.unsafeTake (length b - countWhileEnd k b) b -- | /O(n)/ 'takeWhileEnd' @p@ @b@ returns the longest suffix of -- elements that satisfy predicate @p@. @@ -445,18 +381,6 @@ let n = countWhileEnd k b in Bytes (array b) (offset b + length b - n) n --- | Take the first @n@ bytes from the argument. Precondition: @n ≤ len@ -unsafeTake :: Int -> Bytes -> Bytes -{-# inline unsafeTake #-} -unsafeTake n (Bytes arr off _) = - Bytes arr off n - --- | Drop the first @n@ bytes from the argument. Precondition: @n ≤ len@ -unsafeDrop :: Int -> Bytes -> Bytes -{-# inline unsafeDrop #-} -unsafeDrop n (Bytes arr off len) = - Bytes arr (off + n) (len - n) - -- Internal. The returns the number of bytes that match the -- predicate until the first non-match occurs. If all bytes -- match the predicate, this will return the length originally @@ -713,6 +637,13 @@ {-# inline fromShortByteString #-} fromShortByteString (SBS x) = fromByteArray (ByteArray x) +-- | /O(1)/ Create 'Bytes' from a 'ShortText'. This encodes the text as UTF-8. +-- It is a no-op. +fromShortText :: ShortText -> Bytes +{-# inline fromShortText #-} +fromShortText t = case TS.toShortByteString t of + SBS x -> fromByteArray (ByteArray x) + -- | /O(n)/ Interpreting the bytes an ASCII-encoded characters, convert -- the string to lowercase. This adds @0x20@ to bytes in the range -- @[0x41,0x5A]@ and leaves all other bytes alone. Unconditionally @@ -730,3 +661,21 @@ {-# inline unlift #-} unlift (Bytes (ByteArray arr) (I# off) (I# len)) = Bytes# (# arr, off, len #) + +concatArrayU :: Array Bytes -> ByteArray +{-# noinline concatArrayU #-} +concatArrayU !xs = runByteArrayST $ do + let !arrLen = PM.sizeofArray xs + let !totalByteLen = F.foldl' (\acc b -> length b + acc) 0 xs + dst <- PM.newByteArray totalByteLen + let go !ix !dstOff = if ix < arrLen + then do + x <- PM.indexArrayM xs ix + Pure.unsafeCopy dst dstOff x + go (ix + 1) (dstOff + length x) + else PM.unsafeFreezeByteArray dst + go 0 0 + +concatArray :: Array Bytes -> Bytes +{-# inline concatArray #-} +concatArray !xs = Pure.fromByteArray (concatArrayU xs) diff --git a/src/Data/Bytes/Byte.hs b/src/Data/Bytes/Byte.hs --- a/src/Data/Bytes/Byte.hs +++ b/src/Data/Bytes/Byte.hs @@ -11,7 +11,8 @@ -- on a particular byte and for counting occurences of that -- byte. module Data.Bytes.Byte - ( count + ( -- Re-exported by Data.Bytes + count , split , splitU , splitNonEmpty @@ -23,6 +24,8 @@ , split3 , split4 , splitEnd1 + -- Used by other internal modules + , elemIndexLoop# ) where import Prelude hiding (length) @@ -273,7 +276,11 @@ -- This returns the offset into the byte array. This is not an index -- that will mean anything to the end user, so it cannot be returned -- to them. +-- +-- Exported for use in other internal modules because it is needed in +-- Data.Bytes.Search. elemIndexLoop# :: Word8 -> Bytes -> Int# +{-# inline elemIndexLoop# #-} elemIndexLoop# !w (Bytes arr off@(I# off# ) len) = case len of 0 -> (-1#) _ -> if PM.indexByteArray arr off == w diff --git a/src/Data/Bytes/Pure.hs b/src/Data/Bytes/Pure.hs --- a/src/Data/Bytes/Pure.hs +++ b/src/Data/Bytes/Pure.hs @@ -17,6 +17,7 @@ , toPinnedByteArray , toPinnedByteArrayClone , fromByteArray + , fromPrimArray , length , foldlM , foldrM @@ -30,16 +31,23 @@ , toByteString , pinnedToByteString , fromByteString + , unsafeDrop + , unsafeTake + , unsafeIndex + , unsafeHead + , map + , mapU + , null ) where -import Prelude hiding (length,foldl,foldr) +import Prelude hiding (length,foldl,foldr,map,null) import Control.Monad.Primitive (PrimState,PrimMonad) import Control.Monad.ST.Run (runByteArrayST) import Data.Bits (xor) import Data.Bytes.Types (Bytes(Bytes)) import Data.ByteString (ByteString) -import Data.Primitive (ByteArray,MutableByteArray) +import Data.Primitive (ByteArray(ByteArray),MutableByteArray,PrimArray(PrimArray)) import Data.Word (Word64,Word32,Word8) import Foreign.Ptr (Ptr,plusPtr) import GHC.IO (unsafeIOToST) @@ -48,7 +56,6 @@ import qualified Data.ByteString.Internal as ByteString import qualified Data.ByteString.Unsafe as ByteString import qualified Data.Primitive as PM -import qualified Data.Primitive.Ptr as PM import qualified GHC.Exts as Exts import qualified GHC.ForeignPtr as ForeignPtr @@ -118,6 +125,11 @@ {-# inline fromByteArray #-} fromByteArray b = Bytes b 0 (PM.sizeofByteArray b) +-- | Create a slice of 'Bytes' that spans the entire 'PrimArray' of 8-bit words. +fromPrimArray :: PrimArray Word8 -> Bytes +{-# inline fromPrimArray #-} +fromPrimArray p@(PrimArray b) = Bytes (ByteArray b) 0 (PM.sizeofPrimArray p) + -- | The length of a slice of bytes. length :: Bytes -> Int {-# inline length #-} @@ -254,3 +266,50 @@ ) 0 len where !len = ByteString.length b + +-- | Drop the first @n@ bytes from the argument. Precondition: @n ≤ len@ +unsafeDrop :: Int -> Bytes -> Bytes +{-# inline unsafeDrop #-} +unsafeDrop n (Bytes arr off len) = + Bytes arr (off + n) (len - n) + +-- | Variant of 'map' that returns unsliced byte sequence. +mapU :: (Word8 -> Word8) -> Bytes -> ByteArray +{-# inline mapU #-} +mapU f (Bytes array ix0 len) = runByteArrayST do + dst <- PM.newByteArray len + let go !srcIx !dstIx = if dstIx < len + then do + let w = PM.indexByteArray array srcIx :: Word8 + PM.writeByteArray dst dstIx (f w) + go (srcIx + 1) (dstIx + 1) + else PM.unsafeFreezeByteArray dst + go ix0 0 + +-- | Map over bytes in a sequence. The result has the same length as +-- the argument. +map :: (Word8 -> Word8) -> Bytes -> Bytes +{-# inline map #-} +map f !b = Bytes (mapU f b) 0 (length b) + +-- | Is the byte sequence empty? +null :: Bytes -> Bool +{-# inline null #-} +null (Bytes _ _ len) = len == 0 + +-- | Take the first @n@ bytes from the argument. Precondition: @n ≤ len@ +unsafeTake :: Int -> Bytes -> Bytes +{-# inline unsafeTake #-} +unsafeTake n (Bytes arr off _) = + Bytes arr off n + +-- | Index into the byte sequence at the given position. This index +-- must be less than the length. +unsafeIndex :: Bytes -> Int -> Word8 +{-# inline unsafeIndex #-} +unsafeIndex (Bytes arr off _) ix = PM.indexByteArray arr (off + ix) + +-- | Access the first byte. The given 'Bytes' must be non-empty. +{-# inline unsafeHead #-} +unsafeHead :: Bytes -> Word8 +unsafeHead bs = unsafeIndex bs 0 diff --git a/src/Data/Bytes/Search.hs b/src/Data/Bytes/Search.hs new file mode 100644 --- /dev/null +++ b/src/Data/Bytes/Search.hs @@ -0,0 +1,234 @@ +{-# language BangPatterns #-} +{-# language BlockArguments #-} +{-# language DuplicateRecordFields #-} +{-# language MagicHash #-} +{-# language NamedFieldPuns #-} +{-# language RankNTypes #-} +{-# language TupleSections #-} +{-# language TypeApplications #-} +{-# language UnboxedTuples #-} + +-- This is broken out into a separate module to make it easier +-- to dump core and investigate performance issues. +module Data.Bytes.Search + ( findIndices + , replace + , isInfixOf + ) where + +import Prelude hiding (length,takeWhile,dropWhile,null,foldl,foldr,elem,replicate,any,all,readFile,map) + +import Control.Monad.ST.Run (runByteArrayST,runPrimArrayST) +import Data.Bits((.&.),(.|.),shiftL,finiteBitSize) +import Data.Bytes.Pure (length,unsafeIndex,unsafeHead) +import Data.Bytes.Types (Bytes(Bytes,array,offset)) +import Data.Primitive (ByteArray,PrimArray) +import GHC.Exts (Int(I#)) +import GHC.Word (Word32) + +import qualified Data.Bytes.Byte as Byte +import qualified Data.Bytes.Pure as Pure +import qualified Data.Bytes.Types as Types +import qualified Data.Primitive as PM + +-- Implementation Notes +-- ===================== +-- For karp rabin, there are some easy performance improvements +-- left on the table. The main optimization that has been done is making +-- sure that there is no unnecessary boxing of Int, Word32, or Bytes +-- going on. Here are some other things that have not been done: +-- +-- * The hash is currently a Word32. It would be better to use either +-- Word or Word64 for this. We would need for hashKey to be different. +-- * In several places, we track an index into a Bytes. This index gets +-- repeatedly added to the base offset as we loop over the bytes. We +-- could instead track the true offset instead of repeatedly +-- recalculating it. + +-- | Replace every non-overlapping occurrence of @needle@ in +-- @haystack@ with @replacement@. +replace :: + Bytes -- ^ needle, must not be empty + -> Bytes -- ^ replacement + -> Bytes -- ^ haystack + -> Bytes +{-# noinline replace #-} +-- Implementation note: there is a lot of room to improve the performance +-- of this function. +replace !needle !replacement !haystack@Bytes{array=haystackArray,offset=haystackIndex,length=haystackLength} + | Pure.length needle == 0 = errorWithoutStackTrace "Data.Bytes.replace: needle of length zero" + | Pure.length haystack == 0 = Pure.empty + | Pure.length needle == 1, Pure.length replacement == 1 = + let !needle0 = unsafeIndex needle 0 + !replacement0 = unsafeIndex replacement 0 + in Pure.map (\w -> if w == needle0 then replacement0 else w) haystack + | otherwise = + let !hp = rollingHash needle + !ixs = findIndicesKarpRabin 0 hp needle haystackArray haystackIndex haystackLength + in Pure.fromByteArray (replaceIndices ixs replacement (Pure.length needle) haystackArray haystackIndex haystackLength) + + +-- This is an internal function because it deals explicitly with +-- an offset into a byte array. +-- +-- Example: +-- * haystack len: 39 +-- * ixs: 7, 19, 33 +-- * patLen: 5 +-- * replacment: foo (len 3) +-- We want to perform these copies: +-- * src[0,7] -> dst[0,7] +-- * foo -> dst[7,3] +-- * src[12,7] -> dst[10,7] +-- * foo -> dst[17,3] +-- * src[24,9] -> dst[20,9] +-- * foo -> dst[29,3] +-- * src[38,1] -> dst[32,1] +replaceIndices :: PrimArray Int -> Bytes -> Int -> ByteArray -> Int -> Int -> ByteArray +replaceIndices !ixs !replacement !patLen !haystack !ix0 !len0 = runByteArrayST $ do + let !ixsLen = PM.sizeofPrimArray ixs + let !delta = Pure.length replacement - patLen + dst <- PM.newByteArray (len0 + ixsLen * delta) + let applyReplacement !ixIx !prevSrcIx = if ixIx < ixsLen + then do + let !srcMatchIx = PM.indexPrimArray ixs ixIx + let !offset = ixIx * delta + let !dstIx = srcMatchIx + offset - ix0 + Pure.unsafeCopy dst (prevSrcIx + offset - ix0) + Bytes{array=haystack,offset=prevSrcIx,length=srcMatchIx - prevSrcIx} + Pure.unsafeCopy dst dstIx replacement + applyReplacement (ixIx + 1) (srcMatchIx + patLen) + else do + let !offset = ixIx * delta + Pure.unsafeCopy dst (prevSrcIx + offset - ix0) + Bytes{array=haystack,offset=prevSrcIx,length=(len0 + ix0) - prevSrcIx} + PM.unsafeFreezeByteArray dst + applyReplacement 0 ix0 + +-- | Find locations of non-overlapping instances of @needle@ within @haystack@. +findIndices :: + Bytes -- ^ needle + -> Bytes -- ^ haystack + -> PrimArray Int +findIndices needle Bytes{array,offset=off,length=len} + | needleLen == 0 = errorWithoutStackTrace "Data.Bytes.findIndices: needle with length zero" + | len == 0 = mempty + | otherwise = + let !hp = rollingHash needle + in findIndicesKarpRabin (negate off) hp needle array off len + where + needleLen = Pure.length needle + +-- Precondition: Haystack has non-zero length +-- Precondition: Pattern has non-zero length +-- Uses karp rabin to search. +-- Easy opportunity to improve implementation. Instead of having karpRabin +-- return two slices, we could have it just return a single index. +findIndicesKarpRabin :: + Int -- Output index modifier. Set to negated initial index to make slicing invisible in results. + -> Word32 -- Hash to search for (must agree with pattern) + -> Bytes -- Pattern to search for + -> ByteArray + -> Int -- initial index + -> Int -- length + -> PrimArray Int +findIndicesKarpRabin !ixModifier !hp !pat !haystack !ix0 !len0 = runPrimArrayST $ do + let dstLen = 1 + quot len0 (Pure.length pat) + dst <- PM.newPrimArray dstLen + let go !ix !len !ixIx = case karpRabin hp pat Bytes{array=haystack,offset=ix,length=len} of + (-1) -> do + PM.shrinkMutablePrimArray dst ixIx + PM.unsafeFreezePrimArray dst + skipCount -> do + let !advancement = skipCount - Pure.length pat + let !advancement' = advancement + Pure.length pat + PM.writePrimArray dst ixIx (ix + advancement + ixModifier) + let !ix' = ix + advancement' + go ix' (len - advancement') (ixIx + 1) + go ix0 len0 0 + +-- Output: Negative one means match not found. Other negative +-- numbers should not occur. Zero may occur. Positive number +-- means the number of bytes to skip to make it past the match. +breakSubstring :: Bytes -- ^ String to search for + -> Bytes -- ^ String to search in + -> Int +breakSubstring !pat !haystack@(Bytes _ off0 _) = + case lp of + 0 -> 0 + 1 -> case Byte.elemIndexLoop# (unsafeHead pat) haystack of + (-1#) -> (-1) + off -> 1 + (I# off) - off0 + _ -> if lp * 8 <= finiteBitSize (0 :: Word) + then shift haystack + else karpRabin (rollingHash pat) pat haystack + where + lp = length pat + {-# INLINE shift #-} + shift :: Bytes -> Int + shift !src + | length src < lp = (-1) + | otherwise = search (intoWord $ Pure.unsafeTake lp src) lp + where + intoWord :: Bytes -> Word + intoWord = Pure.foldl' (\w b -> (w `shiftL` 8) .|. fromIntegral b) 0 + wp = intoWord pat + mask = (1 `shiftL` (8 * lp)) - 1 + search :: Word -> Int -> Int + search !w !i + | w == wp = i + | length src <= i = (-1) + | otherwise = search w' (i + 1) + where + b = fromIntegral (Pure.unsafeIndex src i) + w' = mask .&. ((w `shiftL` 8) .|. b) + +-- Only used for karp rabin +rollingHash :: Bytes -> Word32 +{-# inline rollingHash #-} +rollingHash = Pure.foldl' (\h b -> h * hashKey + fromIntegral b) 0 + +hashKey :: Word32 +{-# inline hashKey #-} +hashKey = 2891336453 + +-- Precondition: Length of bytes is greater than or equal to 1. +-- Precondition: Rolling hash agrees with pattern. +-- Output: Negative one means match not found. Other negative +-- numbers should not occur. Zero should not occur. Positive number +-- means the number of bytes to skip to make it past the match. +karpRabin :: Word32 -> Bytes -> Bytes -> Int +karpRabin !hp !pat !src + | length src < lp = (-1) + | otherwise = search (rollingHash $ Pure.unsafeTake lp src) lp + where + lp :: Int + !lp = Pure.length pat + m :: Word32 + !m = hashKey ^ lp + get :: Int -> Word32 + get !ix = fromIntegral (Pure.unsafeIndex src ix) + search !hs !i + | hp == hs && eqBytesNoShortCut pat (Pure.unsafeTake lp (Pure.unsafeDrop (i - lp) src)) = i + | length src <= i = (-1) + | otherwise = search hs' (i + 1) + where + hs' = hs * hashKey + + get i - + m * get (i - lp) + +-- | Is the first argument an infix of the second argument? +-- +-- Uses the Rabin-Karp algorithm: expected time @O(n+m)@, worst-case @O(nm)@. +isInfixOf :: Bytes -- ^ String to search for + -> Bytes -- ^ String to search in + -> Bool +isInfixOf p s = Pure.null p || breakSubstring p s >= 0 + + +-- Precondition: both arguments have the same length +-- Skips the pointer equality check and the length check. +eqBytesNoShortCut :: Bytes -> Bytes -> Bool +{-# inline eqBytesNoShortCut #-} +eqBytesNoShortCut (Bytes arr1 off1 len1) (Bytes arr2 off2 _) = + PM.compareByteArrays arr1 off1 arr2 off2 len1 == EQ diff --git a/src/Data/Bytes/Text/Ascii.hs b/src/Data/Bytes/Text/Ascii.hs --- a/src/Data/Bytes/Text/Ascii.hs +++ b/src/Data/Bytes/Text/Ascii.hs @@ -8,11 +8,13 @@ -- 'Data.Bytes.Text.AsciiExt'. module Data.Bytes.Text.Ascii ( fromString + , decodeDecWord ) where import Data.Bytes.Types (Bytes) import Data.Char (ord) import Data.Word (Word8) +import Data.Bytes.Text.Latin1 (decodeDecWord) import qualified Data.Bytes.Pure as Bytes import qualified GHC.Exts as Exts diff --git a/src/Data/Bytes/Text/Latin1.hs b/src/Data/Bytes/Text/Latin1.hs --- a/src/Data/Bytes/Text/Latin1.hs +++ b/src/Data/Bytes/Text/Latin1.hs @@ -1,6 +1,8 @@ {-# LANGUAGE BangPatterns #-} {-# LANGUAGE MagicHash #-} {-# LANGUAGE TypeApplications #-} +{-# LANGUAGE UnboxedSums #-} +{-# LANGUAGE UnboxedTuples #-} -- | This module treats 'Bytes' data as holding text encoded in ISO-8859-1. This -- encoding can only encode codepoints strictly below @U+0100@, but this allows @@ -19,6 +21,7 @@ module Data.Bytes.Text.Latin1 ( toString , fromString + , decodeDecWord -- * Specialized Comparisons , equals1 , equals2 @@ -34,12 +37,16 @@ , equals12 ) where +import Prelude hiding (length) + import Data.Bytes.Types (Bytes(..)) import Data.Char (ord,chr) import Data.Primitive (ByteArray(ByteArray)) import Data.Word (Word8) -import GHC.Exts (Int(I#),Char(C#)) +import GHC.Exts (Int(I#),Char(C#),or#,ltWord#,int2Word#) +import GHC.Exts (Word(W#),Word#) +import qualified Data.Primitive as PM import qualified Data.Bytes.Pure as Bytes import qualified GHC.Exts as Exts @@ -212,3 +219,55 @@ indexCharArray :: ByteArray -> Int -> Char indexCharArray (ByteArray arr) (I# off) = C# (Exts.indexCharArray# arr off) + +-- | Decode machine-sized word from decimal representation. Returns +-- Nothing on overflow. Allows any number of leading zeros. Trailing +-- non-digit bytes cause Nothing to be returned. +decodeDecWord :: Bytes -> Maybe Word +{-# inline decodeDecWord #-} +decodeDecWord !b = case decWordStart b of + (# (# #) | #) -> Nothing + (# | w #) -> Just (W# w) + +decWordStart :: + Bytes -- Chunk + -> (# (# #) | Word# #) +{-# noinline decWordStart #-} +decWordStart !chunk0 = if length chunk0 > 0 + then + let !w = fromIntegral @Word8 @Word + (PM.indexByteArray (array chunk0) (offset chunk0)) - 48 + in if w < 10 + then decWordMore w (Bytes.unsafeDrop 1 chunk0) + else (# (# #) | #) + else (# (# #) | #) + where + decWordMore :: + Word -- Accumulator + -> Bytes -- Chunk + -> (# (# #) | Word# #) + decWordMore !acc !chunk = let len = length chunk in case len of + 0 -> (# | unW (fromIntegral acc) #) + _ -> + let !w = fromIntegral @Word8 @Word + (PM.indexByteArray (array chunk) (offset chunk)) - 48 + in if w < 10 + then + let (overflow,acc') = unsignedPushBase10 acc w + in if overflow + then (# (# #) | #) + else decWordMore acc' (Bytes.unsafeDrop 1 chunk) + else (# (# #) | #) + +unsignedPushBase10 :: Word -> Word -> (Bool,Word) +{-# inline unsignedPushBase10 #-} +unsignedPushBase10 (W# a) (W# b) = + let !(# ca, r0 #) = Exts.timesWord2# a 10## + !r1 = Exts.plusWord# r0 b + !cb = int2Word# (ltWord# r1 r0) + !c = ca `or#` cb + in (case c of { 0## -> False; _ -> True }, W# r1) + +unW :: Word -> Word# +{-# inline unW #-} +unW (W# w) = w diff --git a/test/Main.hs b/test/Main.hs --- a/test/Main.hs +++ b/test/Main.hs @@ -17,11 +17,16 @@ import Test.Tasty.HUnit ((@=?),testCase) import Test.Tasty.QuickCheck ((===),testProperty,property,Discard(Discard)) import Test.Tasty.QuickCheck ((==>),Arbitrary) +import Test.Tasty.QuickCheck (ASCIIString(ASCIIString)) import Control.Monad.Trans.Writer (Writer,tell) +import Data.Text (Text) +import qualified Data.Text as Text +import qualified Data.Text.Encoding as Text.Encoding import qualified Data.ByteString as ByteString import qualified Data.Bytes as Bytes import qualified Data.Bytes.Text.Ascii as Ascii +import qualified Data.Bytes.Text.Latin1 as Latin1 import qualified Data.Bytes.Text.AsciiExt as AsciiExt import qualified Data.Bytes.Chunks as Chunks import qualified Data.Foldable as Foldable @@ -77,6 +82,42 @@ , testCase "edge: empty string" $ THU.assertBool "" $ not (Bytes.isInfixOf (bytes "hello hello!") (bytes "")) ] + , testGroup "findIndices" + [ testCase "A" $ + Exts.fromList [4] + @=? + Bytes.findIndices (Bytes.fromByteArray (pack "greatest showman")) (Bytes.fromByteArray (pack "the greatest showman of all")) + , testCase "B" $ + Exts.fromList [4] + @=? + Bytes.findIndices (bytes "greatest showman") (bytes "the greatest showman of all") + , testCase "C" $ + Exts.fromList [0,1,2,3,4,5] + @=? + Bytes.findIndices (bytes "a") (bytes "aaaaaa") + , testCase "D" $ + Exts.fromList [0,2,4] + @=? + Bytes.findIndices (bytes "aa") (bytes "aaaaaaa") + ] + , testGroup "replace-spec" + [ testCase "A" $ + Bytes.empty + @=? + Bytes.replace (Bytes.fromByteArray (pack "hello")) (Bytes.fromByteArray (pack "world")) Bytes.empty + , testCase "B" $ + bytes "xzybbcbbc" + @=? + Bytes.replace (bytes "a") (bytes "b") (bytes "xzyabcabc") + , testCase "C" $ + bytes "my favorite month is March!" + @=? + Bytes.replace (bytes "November") (bytes "March") (bytes "my favorite month is November!") + , testCase "D" $ + bytes "Saturn, Saturn, Mars, Saturn" + @=? + Bytes.replace (bytes "Jupiter") (bytes "Saturn") (bytes "Jupiter, Jupiter, Mars, Jupiter") + ] , testGroup "stripOptionalSuffix" [ testCase "A" $ Ascii.fromString "hey m" @@ -117,6 +158,10 @@ @=? Bytes.takeWhileEnd (/= 0x0) (slicedPack [0x1,0x0,0x1,0x2,0x3]) ] + , testProperty "decodeDecWord" $ \(w :: Word) -> + Just w + === + Latin1.decodeDecWord (Bytes.unsafeDrop 1 (Exts.fromList (0xFF : (Exts.toList (Latin1.fromString (show w)))))) , testProperty "elem" $ \(x :: Word8) (xs :: [Word8]) -> List.elem x xs === @@ -161,6 +206,10 @@ mconcat (List.intersperse x xs) === Bytes.intercalate x xs + , testProperty "concatArray" $ \(xs :: [Bytes]) -> + mconcat xs + === + Bytes.concatArray (Exts.fromList xs) , testProperty "splitNonEmpty" $ \(x :: Word8) (xs :: [Word8]) -> Bytes.split x (slicedPack xs) === @@ -185,6 +234,21 @@ [Ascii.fromString "hello", Ascii.fromString "world"] @=? (Bytes.splitInit 0x0A (Ascii.fromString "hello\nworld\nthere")) + , testProperty "replace" $ \(ws :: ASCIIString) (xs :: ASCIIString) (ys :: ASCIIString) (zs :: ASCIIString) -> + let ws' = asciiStringToBytes ws + xs' = asciiStringToBytes xs + ys' = asciiStringToBytes ys + zs' = asciiStringToBytes zs + ws'' = asciiStringToText ws + xs'' = asciiStringToText xs + ys'' = asciiStringToText ys + zs'' = asciiStringToText zs + in case ws of + ASCIIString [] -> property Discard + _ -> + Bytes.replace ws' xs' (ys' <> ws' <> zs') + === + Bytes.fromByteString (Text.Encoding.encodeUtf8 (Text.replace ws'' xs'' (ys'' <> ws'' <> zs''))) , testProperty "splitEnd1" $ \(x :: Word8) (xs :: [Word8]) -> case ByteString.split x (ByteString.pack xs) of [] -> Bytes.splitEnd1 x (slicedPack xs) === Nothing @@ -288,3 +352,9 @@ lawsToTest :: QCC.Laws -> TestTree lawsToTest (QCC.Laws name pairs) = testGroup name (map (uncurry TQC.testProperty) pairs) + +asciiStringToBytes :: ASCIIString -> Bytes +asciiStringToBytes (ASCIIString x) = bytes x + +asciiStringToText :: ASCIIString -> Text +asciiStringToText (ASCIIString x) = Text.pack x