diff --git a/benchmark/Bench.hs b/benchmark/Bench.hs new file mode 100644 --- /dev/null +++ b/benchmark/Bench.hs @@ -0,0 +1,61 @@ +{-# LANGUAGE BangPatterns #-} +module Main where + +import Control.Exception (evaluate) +import Data.IORef +import qualified Data.ByteString as BS +import Data.Time.Clock.POSIX (getPOSIXTime) +import Data.Word (Word8) +import Snappy + +main :: IO () +main = do + src <- BS.readFile "./data/Isaac.Newton-Opticks.txt" + let !payloadMB = fromIntegral (BS.length src) / (1024 * 1024) :: Double + iters = 200 :: Int + + putStrLn $ "Input: " ++ show (BS.length src) ++ " bytes, " ++ show iters ++ " iterations" + + -- Pre-generate varied inputs (16 KB + 1 byte each, tiny overhead) + let inputs = [ BS.snoc src (fromIntegral i :: Word8) | i <- [1..iters] ] + + -- Compress benchmark + cRef <- newIORef undefined + t0c <- getPOSIXTime + mapM_ (\tweaked -> do + let !c = compress tweaked + _ <- evaluate (BS.length c) + writeIORef cRef c) inputs + t1c <- getPOSIXTime + compressed <- readIORef cRef + let totalCompMs = realToFrac (t1c - t0c) * 1000 :: Double + compSecs = realToFrac (t1c - t0c) / fromIntegral iters :: Double + putStrLn $ "compress: " ++ show (BS.length compressed) ++ " bytes -> " + ++ showMBps payloadMB compSecs ++ " MB/s" + ++ " (total " ++ showMs totalCompMs ++ "ms)" + + -- Pre-compress all inputs so decompress benchmark is decompress-only + let compInputs = map (compress . BS.snoc src . (fromIntegral :: Int -> Word8)) [1..iters] + + -- Decompress benchmark (decompress only, no compress overhead) + dRef <- newIORef undefined + t0d <- getPOSIXTime + mapM_ (\c -> + case decompress c of + Left e -> putStrLn ("error: " ++ show e) + Right !d -> do + _ <- evaluate (BS.length d) + writeIORef dRef d) compInputs + t1d <- getPOSIXTime + decompResult <- readIORef dRef + let totalDecompMs = realToFrac (t1d - t0d) * 1000 :: Double + decompSecs = realToFrac (t1d - t0d) / fromIntegral iters :: Double + putStrLn $ "decomp: " ++ show (BS.length decompResult) ++ " bytes -> " + ++ showMBps payloadMB decompSecs ++ " MB/s" + ++ " (total " ++ showMs totalDecompMs ++ "ms)" + +showMBps :: Double -> Double -> String +showMBps mb secs = show (round (mb / secs) :: Int) + +showMs :: Double -> String +showMs x = show (round x :: Int) diff --git a/snappy-hs.cabal b/snappy-hs.cabal --- a/snappy-hs.cabal +++ b/snappy-hs.cabal @@ -1,6 +1,6 @@ cabal-version: 3.0 name: snappy-hs -version: 0.1.0.5 +version: 0.1.1.0 synopsis: Snappy compression library. description: A pure Haskell implementation of the Snappy compression spec. license: MIT @@ -50,7 +50,20 @@ main-is: Main.hs build-depends: base >=4.11 && <5, + bytestring >= 0.11 && < 0.14, snappy-hs + +executable snappy-bench + import: warnings + main-is: Bench.hs + hs-source-dirs: benchmark + build-depends: + base >=4.11 && <5, + bytestring >= 0.11 && < 0.14, + time >= 1.14 && < 1.15, + snappy-hs >= 0.1 && < 1 + default-language: Haskell2010 + ghc-options: -rtsopts -- benchmark snappy-benchmark -- type: exitcode-stdio-1.0 diff --git a/src/Snappy.hs b/src/Snappy.hs --- a/src/Snappy.hs +++ b/src/Snappy.hs @@ -1,7 +1,8 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE LambdaCase #-} - -{-# LANGUAGE Strict #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE Strict #-} +{-# LANGUAGE UnboxedTuples #-} +{-# LANGUAGE UnliftedNewtypes #-} module Snappy ( compress, @@ -10,182 +11,202 @@ ) where import Control.Exception (Exception) -import Data.Bits +import Data.Bits (unsafeShiftL, unsafeShiftR, (.&.), (.|.)) import Data.ByteString.Internal (ByteString (..)) -import Data.Word -import Foreign.ForeignPtr -import Foreign.Marshal.Alloc (callocBytes, free) +import Data.Word (Word64) +import Foreign.ForeignPtr (mallocForeignPtrBytes, withForeignPtr) import Foreign.Marshal.Utils (copyBytes, fillBytes) -import Foreign.Ptr -import Foreign.Storable +import Foreign.Ptr (Ptr) +import Foreign.Storable (sizeOf) +import GHC.Exts + ( Addr# + , Int (..) + , MutableByteArray# + , Ptr (..) + , RealWorld + , (+#) + , eqWord8# + , indexWord8OffAddr# + , isTrue# + , newByteArray# + , or# + , plusAddr# + , readIntArray# + , readWord8OffAddr# + , setByteArray# + , uncheckedShiftL# + , word8ToWord# + , wordToWord32# + , writeIntArray# + , writeWord8OffAddr# + ) +import GHC.IO (IO (..)) +import GHC.Word (Word32 (W32#), Word8 (W8#)) import System.IO.Unsafe (unsafeDupablePerformIO) -data DecodeError - = EmptyInput - | Overflow - | VarIntTooLong - | UnexpectedEOF - | UnsupportedLiteralLength - | InvalidOffset - | InvalidTag - deriving (Show, Eq) -instance Exception DecodeError - decompress :: ByteString -> Either DecodeError ByteString decompress (BS srcFp srcLen) | srcLen == 0 = Left EmptyInput | otherwise = unsafeDupablePerformIO $ - withForeignPtr srcFp $ \srcPtr -> do - let srcEnd = srcPtr `plusPtr` srcLen - vr <- parseVarint srcPtr srcEnd + withForeignPtr srcFp $ \(Ptr srcAddr#) -> do + vr <- parseVarint srcAddr# srcLen case vr of Left e -> pure (Left e) Right (outLen, hdrLen) -> do - let payPtr = srcPtr `plusPtr` hdrLen outFp <- mallocForeignPtrBytes (max 1 outLen) - withForeignPtr outFp $ \outPtr -> do - r <- decodeLoop payPtr srcEnd outPtr 0 outLen + withForeignPtr outFp $ \(Ptr outAddr#) -> do + r <- decodeLoop srcAddr# hdrLen srcLen outAddr# 0 outLen case r of - Left e -> pure (Left e) + Left e -> pure (Left e) Right _ -> pure (Right (BS outFp outLen)) +compress :: ByteString -> ByteString +compress (BS srcFp srcLen) = unsafeDupablePerformIO $ + withForeignPtr srcFp $ \(Ptr srcAddr#) -> do + let maxOut = 10 + srcLen + (srcLen `div` 6) + 64 + outFp <- mallocForeignPtrBytes maxOut + withForeignPtr outFp $ \(Ptr outAddr#) -> do + hdrLen <- writeVarint outAddr# (fromIntegral srcLen :: Word64) + written <- + if srcLen < 4 + then writeLiteral srcAddr# 0 srcLen outAddr# hdrLen + else do + let !tableSize = 1 `unsafeShiftL` 15 + !tableBytes = tableSize * sizeOfInt + tbl <- newZeroedByteArray tableBytes + compressLoop srcAddr# srcLen tbl outAddr# hdrLen 0 0 + pure (BS outFp written) + + {-# INLINE parseVarint #-} -parseVarint :: Ptr Word8 -> Ptr Word8 -> IO (Either DecodeError (Int, Int)) -parseVarint !ptr !end = go 0 0 0 +parseVarint :: Addr# -> Int -> IO (Either DecodeError (Int, Int)) +parseVarint !addr# !srcLen = go 0 0 0 where go :: Word64 -> Int -> Int -> IO (Either DecodeError (Int, Int)) go !acc !s !n - | n > 10 = pure (Left VarIntTooLong) - | s > 63 = pure (Left Overflow) - | ptr `plusPtr` n >= end = pure (Left UnexpectedEOF) + | n > 10 = pure (Left VarIntTooLong) + | s > 63 = pure (Left Overflow) + | n >= srcLen = pure (Left UnexpectedEOF) + | otherwise = + let b = fromIntegral (ixByte addr# n) :: Word64 + in if b < 0x80 + then pure (Right (fromIntegral (acc .|. (b `unsafeShiftL` s)), n + 1)) + else go + (acc .|. ((b .&. 0x7f) `unsafeShiftL` s)) + (s + 7) + (n + 1) + +{-# INLINE writeVarint #-} +writeVarint :: Addr# -> Word64 -> IO Int +writeVarint !addr# = go 0 + where + go !off !x + | x < 0x80 = do + writeW8 addr# off (fromIntegral x :: Word8) + pure (off + 1) | otherwise = do - b <- peekByteOff ptr n :: IO Word8 - if b < 0x80 - then pure (Right (fromIntegral (acc .|. (fromIntegral b `unsafeShiftL` s)), n + 1)) - else - go - (acc .|. ((fromIntegral b .&. 0x7f) `unsafeShiftL` s)) - (s + 7) - (n + 1) + writeW8 addr# off (fromIntegral (0x80 .|. (x .&. 0x7f)) :: Word8) + go (off + 1) (x `unsafeShiftR` 7) -decodeLoop :: Ptr Word8 -> Ptr Word8 -> Ptr Word8 -> Int -> Int -> IO (Either DecodeError Int) -decodeLoop !src0 !srcEnd !dst = go src0 +decodeLoop + :: Addr# -> Int -> Int -- src base, current offset, end offset + -> Addr# -> Int -> Int -- dst base, current offset, capacity + -> IO (Either DecodeError Int) +decodeLoop !srcAddr# !srcOff0 !srcEnd !outAddr# = go srcOff0 where - go :: Ptr Word8 -> Int -> Int -> IO (Either DecodeError Int) - go !src !di !outLen - | src >= srcEnd = pure (Right di) + go :: Int -> Int -> Int -> IO (Either DecodeError Int) + go !srcOff !outOff !outLen + | srcOff >= srcEnd = pure (Right outOff) | otherwise = do - tag <- peek src :: IO Word8 - let src' = src `plusPtr` 1 + let !tag = ixByte srcAddr# srcOff + !srcOff1 = srcOff + 1 case tag .&. 0x03 of 0 -> do - let hdr = fromIntegral (tag `unsafeShiftR` 2) :: Int + let !hdr = fromIntegral (tag `unsafeShiftR` 2) :: Int if hdr < 60 then do - let litLen = hdr + 1 - if src' `plusPtr` litLen > srcEnd || di + litLen > outLen + let !litLen = hdr + 1 + if srcOff1 + litLen > srcEnd || outOff + litLen > outLen then pure (Left UnexpectedEOF) else do - copyBytes (dst `plusPtr` di) src' litLen - go (src' `plusPtr` litLen) (di + litLen) outLen + copyBytes (atOff outAddr# outOff) (atOff srcAddr# srcOff1) litLen + go (srcOff1 + litLen) (outOff + litLen) outLen else do - let nb = hdr - 59 + let !nb = hdr - 59 if nb < 1 || nb > 4 then pure (Left UnsupportedLiteralLength) - else - if src' `plusPtr` nb > srcEnd - then pure (Left UnexpectedEOF) - else do - lm1 <- readLEn src' nb - let litLen = lm1 + 1 - src'' = src' `plusPtr` nb - if src'' `plusPtr` litLen > srcEnd || di + litLen > outLen - then pure (Left UnexpectedEOF) - else do - copyBytes (dst `plusPtr` di) src'' litLen - go (src'' `plusPtr` litLen) (di + litLen) outLen + else if srcOff1 + nb > srcEnd + then pure (Left UnexpectedEOF) + else do + lm1 <- readLEn srcAddr# srcOff1 nb + let !litLen = lm1 + 1 + !srcOff2 = srcOff1 + nb + if srcOff2 + litLen > srcEnd || outOff + litLen > outLen + then pure (Left UnexpectedEOF) + else do + copyBytes (atOff outAddr# outOff) (atOff srcAddr# srcOff2) litLen + go (srcOff2 + litLen) (outOff + litLen) outLen 1 -> do - if src' >= srcEnd + if srcOff1 >= srcEnd then pure (Left UnexpectedEOF) else do - lo <- peek src' :: IO Word8 - let cpLen = (fromIntegral ((tag `unsafeShiftR` 2) .&. 0x07) :: Int) + 4 - offHi = fromIntegral (tag `unsafeShiftR` 5) :: Int - offset = (offHi `unsafeShiftL` 8) .|. fromIntegral lo - if offset <= 0 || offset > di + let !lo = fromIntegral (ixByte srcAddr# srcOff1) :: Int + !cpLen = (fromIntegral ((tag `unsafeShiftR` 2) .&. 0x07) :: Int) + 4 + !offHi = fromIntegral (tag `unsafeShiftR` 5) :: Int + !offset = (offHi `unsafeShiftL` 8) .|. lo + if offset <= 0 || offset > outOff then pure (Left InvalidOffset) else do - copyOverlap dst di offset cpLen - go (src' `plusPtr` 1) (di + cpLen) outLen + copyOverlap outAddr# outOff offset cpLen + go (srcOff1 + 1) (outOff + cpLen) outLen 2 -> do - if src' `plusPtr` 2 > srcEnd + if srcOff1 + 2 > srcEnd then pure (Left UnexpectedEOF) else do - off <- readLEn src' 2 - let cpLen = (fromIntegral ((tag `unsafeShiftR` 2) .&. 0x3f) :: Int) + 1 - if off <= 0 || off > di + off <- readLEn srcAddr# srcOff1 2 + let !cpLen = (fromIntegral ((tag `unsafeShiftR` 2) .&. 0x3f) :: Int) + 1 + if off <= 0 || off > outOff then pure (Left InvalidOffset) else do - copyOverlap dst di off cpLen - go (src' `plusPtr` 2) (di + cpLen) outLen + copyOverlap outAddr# outOff off cpLen + go (srcOff1 + 2) (outOff + cpLen) outLen 3 -> do - if src' `plusPtr` 4 > srcEnd + if srcOff1 + 4 > srcEnd then pure (Left UnexpectedEOF) else do - off <- readLEn src' 4 - let cpLen = (fromIntegral ((tag `unsafeShiftR` 2) .&. 0x3f) :: Int) + 1 - if off <= 0 || off > di + off <- readLEn srcAddr# srcOff1 4 + let !cpLen = (fromIntegral ((tag `unsafeShiftR` 2) .&. 0x3f) :: Int) + 1 + if off <= 0 || off > outOff then pure (Left InvalidOffset) else do - copyOverlap dst di off cpLen - go (src' `plusPtr` 4) (di + cpLen) outLen + copyOverlap outAddr# outOff off cpLen + go (srcOff1 + 4) (outOff + cpLen) outLen _ -> pure (Left InvalidTag) -compress :: ByteString -> ByteString -compress (BS srcFp srcLen) = unsafeDupablePerformIO $ - withForeignPtr srcFp $ \srcPtr -> do - -- Worst-case expansion: varint header (≤10 B) + per-literal overhead - let maxOut = 10 + srcLen + (srcLen `div` 6) + 64 - outFp <- mallocForeignPtrBytes maxOut - withForeignPtr outFp $ \outPtr -> do - hdrLen <- writeVarint outPtr (fromIntegral srcLen :: Word64) - written <- - if srcLen < 4 - then writeLiteral srcPtr 0 srcLen outPtr hdrLen - else do - let !tableSize = 1 `unsafeShiftL` 15 - !tableBytes = tableSize * sizeOfInt - tbl <- callocBytes tableBytes :: IO (Ptr Int) - w <- compressLoop srcPtr srcLen tbl outPtr hdrLen 0 0 - free tbl - pure w - pure (BS outFp written) - -sizeOfInt :: Int -sizeOfInt = sizeOf (0 :: Int) -{-# INLINE sizeOfInt #-} - -compressLoop :: Ptr Word8 -> Int -> Ptr Int -> Ptr Word8 -> Int -> Int -> Int -> IO Int +compressLoop + :: Addr# -> Int -> HashTable + -> Addr# -> Int -> Int -> Int + -> IO Int compressLoop !src !srcLen !tbl !dst !dstOff0 = go dstOff0 where go :: Int -> Int -> Int -> IO Int go !dstOff !i !litStart | i + 3 >= srcLen = do - let remLen = srcLen - litStart + let !remLen = srcLen - litStart if remLen > 0 then writeLiteral src litStart remLen dst dstOff else pure dstOff | otherwise = do - w <- readU32LE src i - let !h = hash32 w - prev0 <- peekByteOff tbl (h * sizeOfInt) :: IO Int - pokeByteOff tbl (h * sizeOfInt) (i + 1 :: Int) - let prev = prev0 - 1 + let !w = readU32LE src i + !h = hash32 w + prev0 <- readTbl tbl h + writeTbl tbl h (i + 1) + let !prev = prev0 - 1 if prev < 0 then go dstOff (i + 1) litStart else do - pw <- readU32LE src prev + let !pw = readU32LE src prev if pw /= w then go dstOff (i + 1) litStart else do @@ -193,77 +214,67 @@ if i > litStart then writeLiteral src litStart (i - litStart) dst dstOff else pure dstOff - extra <- extendMatchIO src (i + 4) (prev + 4) srcLen - let matchLen = 4 + extra - offset = i - prev + let !extra = extendMatch src (i + 4) (prev + 4) srcLen + !matchLen = 4 + extra + !offset = i - prev dstOff'' <- writeCopies offset matchLen dst dstOff' - let i' = i + matchLen + let !i' = i + matchLen go dstOff'' i' i' -extendMatchIO :: Ptr Word8 -> Int -> Int -> Int -> IO Int -extendMatchIO !ptr !a0 !b0 !limit = go 0 +-- | Extend a match as far as possible; pure since source is immutable. +extendMatch :: Addr# -> Int -> Int -> Int -> Int +extendMatch addr# a0 b0 limit = go 0 where + go :: Int -> Int go !n - | a0 + n >= limit || b0 + n >= limit = pure n - | otherwise = do - ba <- peekByteOff ptr (a0 + n) :: IO Word8 - bb <- peekByteOff ptr (b0 + n) :: IO Word8 - if ba == bb then go (n + 1) else pure n -{-# INLINE extendMatchIO #-} + | a0 + n >= limit || b0 + n >= limit = n + | otherwise = + let !(I# an#) = a0 + n + !(I# bn#) = b0 + n + in if isTrue# (eqWord8# + (indexWord8OffAddr# addr# an#) + (indexWord8OffAddr# addr# bn#)) + then go (n + 1) + else n +{-# INLINE extendMatch #-} --- -------------------------------------------------------------------------- --- Helpers – reading --- -------------------------------------------------------------------------- {-# INLINE readLEn #-} -readLEn :: Ptr Word8 -> Int -> IO Int -readLEn !ptr = \case - 1 -> do - b0 <- peek ptr :: IO Word8 - pure (fromIntegral b0) - 2 -> do - b0 <- peek ptr :: IO Word8 - b1 <- peekByteOff ptr 1 :: IO Word8 - pure (fromIntegral b0 .|. (fromIntegral b1 `unsafeShiftL` 8)) - 3 -> do - b0 <- peek ptr :: IO Word8 - b1 <- peekByteOff ptr 1 :: IO Word8 - b2 <- peekByteOff ptr 2 :: IO Word8 - pure - ( fromIntegral b0 - .|. (fromIntegral b1 `unsafeShiftL` 8) - .|. (fromIntegral b2 `unsafeShiftL` 16) - ) - 4 -> do - b0 <- peek ptr :: IO Word8 - b1 <- peekByteOff ptr 1 :: IO Word8 - b2 <- peekByteOff ptr 2 :: IO Word8 - b3 <- peekByteOff ptr 3 :: IO Word8 - pure - ( fromIntegral b0 - .|. (fromIntegral b1 `unsafeShiftL` 8) - .|. (fromIntegral b2 `unsafeShiftL` 16) - .|. (fromIntegral b3 `unsafeShiftL` 24) - ) - n -> go 0 0 n -- fallback +readLEn :: Addr# -> Int -> Int -> IO Int +readLEn !addr# !off !n + | n == 1 = pure $! + fromIntegral (ixByte addr# off) + | n == 2 = pure $! + fromIntegral (ixByte addr# off) + .|. (fromIntegral (ixByte addr# (off + 1)) `unsafeShiftL` 8) + | n == 3 = pure $! + fromIntegral (ixByte addr# off) + .|. (fromIntegral (ixByte addr# (off + 1)) `unsafeShiftL` 8) + .|. (fromIntegral (ixByte addr# (off + 2)) `unsafeShiftL` 16) + | n == 4 = pure $! + fromIntegral (ixByte addr# off) + .|. (fromIntegral (ixByte addr# (off + 1)) `unsafeShiftL` 8) + .|. (fromIntegral (ixByte addr# (off + 2)) `unsafeShiftL` 16) + .|. (fromIntegral (ixByte addr# (off + 3)) `unsafeShiftL` 24) + | otherwise = pure $! go 0 0 where - go !acc !_ 0 = pure acc - go !acc !k m = do - b <- peekByteOff ptr k :: IO Word8 - go (acc .|. (fromIntegral b `unsafeShiftL` (k * 8))) (k + 1) (m - 1) + go !acc !k + | k >= n = acc + | otherwise = + go (acc .|. (fromIntegral (ixByte addr# (off + k)) `unsafeShiftL` (k * 8))) (k + 1) +-- | Read a 32-bit little-endian word; pure since source is immutable. {-# INLINE readU32LE #-} -readU32LE :: Ptr Word8 -> Int -> IO Word32 -readU32LE !ptr !i = do - b0 <- peekByteOff ptr i :: IO Word8 - b1 <- peekByteOff ptr (i + 1) :: IO Word8 - b2 <- peekByteOff ptr (i + 2) :: IO Word8 - b3 <- peekByteOff ptr (i + 3) :: IO Word8 - pure $! - fromIntegral b0 - .|. (fromIntegral b1 `unsafeShiftL` 8) - .|. (fromIntegral b2 `unsafeShiftL` 16) - .|. (fromIntegral b3 `unsafeShiftL` 24) +readU32LE :: Addr# -> Int -> Word32 +readU32LE addr# (I# i#) = + let b0 = word8ToWord# (indexWord8OffAddr# addr# i#) + b1 = word8ToWord# (indexWord8OffAddr# addr# (i# +# 1#)) + b2 = word8ToWord# (indexWord8OffAddr# addr# (i# +# 2#)) + b3 = word8ToWord# (indexWord8OffAddr# addr# (i# +# 3#)) + w = b0 `or#` (b1 `uncheckedShiftL#` 8#) + `or#` (b2 `uncheckedShiftL#` 16#) + `or#` (b3 `uncheckedShiftL#` 24#) + in W32# (wordToWord32# w) {-# INLINE hash32 #-} hash32 :: Word32 -> Int @@ -271,113 +282,167 @@ let !kMul = 0x1e35a7bd :: Word32 in fromIntegral ((w * kMul) `unsafeShiftR` (32 - 15)) -{-# INLINE writeVarint #-} -writeVarint :: Ptr Word8 -> Word64 -> IO Int -writeVarint !ptr = go 0 - where - go !off !x - | x < 0x80 = do - pokeByteOff ptr off (fromIntegral x :: Word8) - pure (off + 1) - | otherwise = do - pokeByteOff ptr off (fromIntegral (0x80 .|. (x .&. 0x7f)) :: Word8) - go (off + 1) (x `unsafeShiftR` 7) - -writeLiteral :: Ptr Word8 -> Int -> Int -> Ptr Word8 -> Int -> IO Int +writeLiteral :: Addr# -> Int -> Int -> Addr# -> Int -> IO Int writeLiteral _ _ 0 _ !dstOff = pure dstOff writeLiteral !src !off !len !dst !dstOff | len <= 60 = do let !tag = fromIntegral (((len - 1) `unsafeShiftL` 2) .|. 0x00) :: Word8 - pokeByteOff dst dstOff tag - copyBytes (dst `plusPtr` (dstOff + 1)) (src `plusPtr` off) len + writeW8 dst dstOff tag + copyBytes (atOff dst (dstOff + 1)) (atOff src off) len pure (dstOff + 1 + len) | otherwise = do - let l1 = len - 1 - nb = bytesFor l1 + let !l1 = len - 1 + !nb = bytesFor l1 !tag = fromIntegral (((59 + nb) `unsafeShiftL` 2) .|. 0x00) :: Word8 - pokeByteOff dst dstOff tag + writeW8 dst dstOff tag writeLEn dst (dstOff + 1) nb (fromIntegral l1 :: Word32) - copyBytes (dst `plusPtr` (dstOff + 1 + nb)) (src `plusPtr` off) len + copyBytes (atOff dst (dstOff + 1 + nb)) (atOff src off) len pure (dstOff + 1 + nb + len) -writeCopies :: Int -> Int -> Ptr Word8 -> Int -> IO Int +writeCopies :: Int -> Int -> Addr# -> Int -> IO Int writeCopies !offset = go where - go :: Int -> Ptr Word8 -> Int -> IO Int + go :: Int -> Addr# -> Int -> IO Int go !remLen !dst !dstOff | remLen <= 0 = pure dstOff | offset <= 0x7FF && remLen >= 4 = do - let c = min 11 remLen + let !c = min 11 remLen dstOff' <- writeCopy1 c offset dst dstOff go (remLen - c) dst dstOff' | offset <= 0xFFFF = do - let c = min 64 remLen + let !c = min 64 remLen dstOff' <- writeCopy2 c offset dst dstOff go (remLen - c) dst dstOff' | otherwise = do - let c = min 64 remLen + let !c = min 64 remLen dstOff' <- writeCopy4 c offset dst dstOff go (remLen - c) dst dstOff' {-# INLINE writeCopies #-} {-# INLINE writeCopy1 #-} -writeCopy1 :: Int -> Int -> Ptr Word8 -> Int -> IO Int +writeCopy1 :: Int -> Int -> Addr# -> Int -> IO Int writeCopy1 !len !off !dst !dstOff = do let !lenField = (len - 4) .&. 0x7 - !offHi = (off `unsafeShiftR` 8) .&. 0x7 - !tag = fromIntegral ((lenField `unsafeShiftL` 2) .|. (offHi `unsafeShiftL` 5) .|. 0x01) :: Word8 - !lo = fromIntegral (off .&. 0xFF) :: Word8 - pokeByteOff dst dstOff tag - pokeByteOff dst (dstOff + 1) lo + !offHi = (off `unsafeShiftR` 8) .&. 0x7 + !tag = fromIntegral ((lenField `unsafeShiftL` 2) .|. (offHi `unsafeShiftL` 5) .|. 0x01) :: Word8 + !lo = fromIntegral (off .&. 0xFF) :: Word8 + writeW8 dst dstOff tag + writeW8 dst (dstOff + 1) lo pure (dstOff + 2) {-# INLINE writeCopy2 #-} -writeCopy2 :: Int -> Int -> Ptr Word8 -> Int -> IO Int +writeCopy2 :: Int -> Int -> Addr# -> Int -> IO Int writeCopy2 !len !off !dst !dstOff = do let !tag = fromIntegral (((len - 1) `unsafeShiftL` 2) .|. 0x02) :: Word8 - pokeByteOff dst dstOff tag + writeW8 dst dstOff tag writeLEn dst (dstOff + 1) 2 (fromIntegral off :: Word32) pure (dstOff + 3) {-# INLINE writeCopy4 #-} -writeCopy4 :: Int -> Int -> Ptr Word8 -> Int -> IO Int +writeCopy4 :: Int -> Int -> Addr# -> Int -> IO Int writeCopy4 !len !off !dst !dstOff = do let !tag = fromIntegral (((len - 1) `unsafeShiftL` 2) .|. 0x03) :: Word8 - pokeByteOff dst dstOff tag + writeW8 dst dstOff tag writeLEn dst (dstOff + 1) 4 (fromIntegral off :: Word32) pure (dstOff + 5) {-# INLINE writeLEn #-} -writeLEn :: Ptr Word8 -> Int -> Int -> Word32 -> IO () +writeLEn :: Addr# -> Int -> Int -> Word32 -> IO () writeLEn !dst !off !n !v = go 0 where go !k - | k >= n = pure () + | k >= n = pure () | otherwise = do - pokeByteOff dst (off + k) (fromIntegral ((v `unsafeShiftR` (8 * k)) .&. 0xFF) :: Word8) + writeW8 dst (off + k) (fromIntegral ((v `unsafeShiftR` (8 * k)) .&. 0xFF) :: Word8) go (k + 1) {-# INLINE bytesFor #-} bytesFor :: Int -> Int bytesFor x - | x < 0x100 = 1 - | x < 0x10000 = 2 + | x < 0x100 = 1 + | x < 0x10000 = 2 | x < 0x1000000 = 3 - | otherwise = 4 + | otherwise = 4 {-# INLINE copyOverlap #-} -copyOverlap :: Ptr Word8 -> Int -> Int -> Int -> IO () +copyOverlap :: Addr# -> Int -> Int -> Int -> IO () copyOverlap !base !outOff !offset !len - | offset >= len = copyBytes (base `plusPtr` outOff) (base `plusPtr` srcStart) len + | offset >= len = + copyBytes (atOff base outOff) (atOff base srcStart) len | offset == 1 = do - b <- peekByteOff base srcStart :: IO Word8 - fillBytes (base `plusPtr` outOff) b len + b <- readW8 base srcStart + fillBytes (atOff base outOff) b len | otherwise = go 0 where !srcStart = outOff - offset go !i - | i >= len = pure () + | i >= len = pure () | otherwise = do - b <- peekByteOff base (srcStart + (i `rem` offset)) :: IO Word8 - pokeByteOff base (outOff + i) b + b <- readW8 base (srcStart + (i `rem` offset)) + writeW8 base (outOff + i) b go (i + 1) + +sizeOfInt :: Int +sizeOfInt = sizeOf (0 :: Int) +{-# INLINE sizeOfInt #-} + +-- | Lifted wrapper around an unlifted 'MutableByteArray#'. +data HashTable = HashTable (MutableByteArray# RealWorld) + +data DecodeError + = EmptyInput + | Overflow + | VarIntTooLong + | UnexpectedEOF + | UnsupportedLiteralLength + | InvalidOffset + | InvalidTag + deriving (Show, Eq) + +instance Exception DecodeError + +-- | Allocate a GC-managed byte array, zero-initialised. +newZeroedByteArray :: Int -> IO HashTable +newZeroedByteArray (I# n#) = IO $ \s0 -> + case newByteArray# n# s0 of + (# s1, arr #) -> case setByteArray# arr 0# n# 0# s1 of + s2 -> (# s2, HashTable arr #) +{-# INLINE newZeroedByteArray #-} + +-- | Read the Int at slot index @i@ (each slot is one machine Int). +readTbl :: HashTable -> Int -> IO Int +readTbl (HashTable arr) (I# i#) = IO $ \s -> + case readIntArray# arr i# s of + (# s', v# #) -> (# s', I# v# #) +{-# INLINE readTbl #-} + +-- | Write an Int to slot index @i@. +writeTbl :: HashTable -> Int -> Int -> IO () +writeTbl (HashTable arr) (I# i#) (I# v#) = IO $ \s -> + case writeIntArray# arr i# v# s of + s' -> (# s', () #) +{-# INLINE writeTbl #-} + +-- | Pure read: one byte from an immutable foreign address. +ixByte :: Addr# -> Int -> Word8 +ixByte addr# (I# i#) = W8# (indexWord8OffAddr# addr# i#) +{-# INLINE ixByte #-} + +-- | Mutable read: one byte from a (possibly mutable) address. +readW8 :: Addr# -> Int -> IO Word8 +readW8 addr# (I# i#) = IO $ \s -> + case readWord8OffAddr# addr# i# s of + (# s', b# #) -> (# s', W8# b# #) +{-# INLINE readW8 #-} + +-- | Write one byte to an address. +writeW8 :: Addr# -> Int -> Word8 -> IO () +writeW8 addr# (I# i#) (W8# b#) = IO $ \s -> + case writeWord8OffAddr# addr# i# b# s of + s' -> (# s', () #) +{-# INLINE writeW8 #-} + +-- | Construct a @Ptr@ from a base @Addr#@ plus an @Int@ byte offset. +atOff :: Addr# -> Int -> Ptr a +atOff addr# (I# i#) = Ptr (addr# `plusAddr#` i#) +{-# INLINE atOff #-} diff --git a/test/Main.hs b/test/Main.hs --- a/test/Main.hs +++ b/test/Main.hs @@ -1,4 +1,23 @@ +{-# LANGUAGE OverloadedStrings #-} module Main (main) where +import qualified Data.ByteString as BS +import Snappy + main :: IO () -main = putStrLn "Test suite not yet implemented." +main = do + roundTrip "short repetitive" (BS.replicate 100 65) + roundTrip "single byte" "x" + roundTrip "short string" "Hello, world!" + roundTrip "longer repetitive" (BS.replicate 1000 65) + roundTrip "varied" "Hello, Snappy! Hello, Snappy! ABCDEFGHIJKLMNOP 1234567890" + roundTrip "already compressed-like" (BS.pack [0..255]) + putStrLn "All tests passed." + +roundTrip :: String -> BS.ByteString -> IO () +roundTrip label bs = + case decompress (compress bs) of + Left e -> fail $ label ++ ": decode error: " ++ show e + Right d -> if d == bs + then putStrLn $ "OK: " ++ label + else fail $ label ++ ": round-trip mismatch"