diff --git a/Benchmark.hs b/Benchmark.hs new file mode 100644 --- /dev/null +++ b/Benchmark.hs @@ -0,0 +1,35 @@ +import Codec.Compression.Zlib(ZlibDecoder(..), decompressIncremental) +import Control.Monad(unless) +import qualified Data.ByteString as S +import qualified Data.ByteString.Lazy as L +import Data.Time.Clock(getCurrentTime, diffUTCTime) +import Prelude hiding (readFile, writeFile) + +main :: IO () +main = + do zbstr <- L.readFile "test/test-cases/tor-list.z" + goldbstr <- L.readFile "test/test-cases/tor-list.gold" + before <- getCurrentTime + runDecompression (L.toChunks zbstr) goldbstr decompressIncremental + after <- getCurrentTime + putStrLn ("Decompression took " ++ show (diffUTCTime after before)) + +runDecompression :: [S.ByteString] -> L.ByteString -> ZlibDecoder -> IO () +runDecompression ls real decoder = + case decoder of + Done -> + do unless (null ls) $ + fail "ERROR: Finished decompression with data left." + unless (L.null real) $ + fail "ERROR: Did not completely decompress file." + return () + DecompError e -> + fail ("ERROR: " ++ show e) + NeedMore f | (x:rest) <- ls -> runDecompression rest real (f x) + | otherwise -> + fail "ERROR: Ran out of data mid-decompression." + Chunk c m -> + let (realfirst, realrest) = L.splitAt (L.length c) real + in if realfirst == c + then runDecompression ls realrest m + else fail "Mismatch in decompression" diff --git a/Deflate.hs b/Deflate.hs --- a/Deflate.hs +++ b/Deflate.hs @@ -1,8 +1,11 @@ -import Codec.Compression.Zlib(decompress) -import Data.ByteString.Lazy(readFile, writeFile) +import Codec.Compression.Zlib(ZlibDecoder(..), decompressIncremental) +import Control.Monad(unless) +import qualified Data.ByteString as S +import qualified Data.ByteString.Lazy as L import Data.List(isSuffixOf) import Prelude hiding (readFile, writeFile) -import System.Environment +import System.Environment(getArgs) +import System.IO(IOMode(..), Handle, openFile, hClose) main :: IO () main = @@ -10,10 +13,28 @@ case args of [ifile] -> if ".z" `isSuffixOf` ifile - then do bstr <- readFile ifile - case decompress bstr of - Left err -> putStrLn (show err) - Right bs -> writeFile (take (length ifile - 2) ifile) bs + then do bstr <- L.readFile ifile + let outname = take (length ifile - 2) ifile + hndl <- openFile outname WriteMode + runDecompression hndl (L.toChunks bstr) decompressIncremental else putStrLn "Unexpected file name." _ -> putStrLn "USAGE: deflate [filename]" + +runDecompression :: Handle -> [S.ByteString] -> ZlibDecoder -> IO () +runDecompression hndl ls decoder = + case decoder of + Done -> + do unless (null ls) $ + putStrLn "WARNING: Finished decompression with data left." + hClose hndl + DecompError e -> + do putStrLn ("ERROR: " ++ show e) + hClose hndl + NeedMore f | (x:rest) <- ls -> runDecompression hndl rest (f x) + | otherwise -> + do putStrLn "ERROR: Ran out of data mid-decompression." + hClose hndl + Chunk c m -> + do L.hPut hndl c + runDecompression hndl ls m diff --git a/pure-zlib.cabal b/pure-zlib.cabal --- a/pure-zlib.cabal +++ b/pure-zlib.cabal @@ -1,5 +1,5 @@ name: pure-zlib -version: 0.5 +version: 0.6 synopsis: A Haskell-only implementation of zlib / DEFLATE homepage: http://github.com/GaloisInc/pure-zlib license: BSD3 @@ -18,13 +18,13 @@ ghc-options: -Wall hs-source-dirs: src build-depends: + array >= 0.4 && < 0.9, base >= 4.6 && < 5.0, base-compat >= 0.9.1 && < 0.11, bytestring >= 0.10 && < 0.11, bytestring-builder >= 0.10 && < 0.11, containers >= 0.5 && < 0.7, - fingertree >= 0.1 && < 0.3, - monadLib >= 3.7 && < 3.9 + fingertree >= 0.1 && < 0.3 exposed-modules: Codec.Compression.Zlib, Codec.Compression.Zlib.Adler32, @@ -67,6 +67,18 @@ tasty >= 0.11.0.4 && < 0.13, tasty-hunit >= 0.9.2 && < 0.11, tasty-quickcheck >= 0.8.4 && < 0.11 + +benchmark bench-zlib + type: exitcode-stdio-1.0 + main-is: Benchmark.hs + default-language: Haskell2010 + ghc-options: -Wall + build-depends: + base >= 4.6 && < 5.0, + base-compat >= 0.9.1 && < 0.11, + bytestring >= 0.10 && < 0.11, + pure-zlib >= 0.5 && < 0.8, + time >= 1.4.2 && < 1.8 source-repository head type: git diff --git a/src/Codec/Compression/Zlib.hs b/src/Codec/Compression/Zlib.hs --- a/src/Codec/Compression/Zlib.hs +++ b/src/Codec/Compression/Zlib.hs @@ -1,42 +1,58 @@ {-# LANGUAGE MultiWayIf #-} module Codec.Compression.Zlib( DecompressionError(..) + , ZlibDecoder(NeedMore, Chunk, Done, DecompError) , decompress + , decompressIncremental ) where -import Codec.Compression.Zlib.Deflate -import Codec.Compression.Zlib.Monad -import Data.Bits -import Data.ByteString.Lazy(ByteString) -import qualified Data.ByteString.Lazy as BS -import Data.Word +import Codec.Compression.Zlib.Deflate(inflate) +import Codec.Compression.Zlib.Monad(ZlibDecoder(..), DeflateM, + DecompressionError(..), + runDeflateM, raise, nextByte) +import Control.Monad(unless, when, replicateM_) +import Data.Bits((.|.), (.&.), shiftL, shiftR, testBit) +import Data.ByteString.Builder(lazyByteString,toLazyByteString) +import qualified Data.ByteString.Lazy as L +import Data.Monoid((<>)) +import Data.Word(Word16) +import Prelude() +import Prelude.Compat -decompress :: ByteString -> Either DecompressionError ByteString -decompress ifile = - case BS.uncons ifile of - Nothing -> Left (HeaderError "Could not read CMF.") - Just (cmf, rest) -> - case BS.uncons rest of - Nothing -> Left (HeaderError "Could not read FLG.") - Just (flg, body) -> - runDecompression cmf flg body +decompressIncremental :: ZlibDecoder +decompressIncremental = runDeflateM inflateWithHeaders -runDecompression :: Word8 -> Word8 -> ByteString -> - Either DecompressionError ByteString -runDecompression cmf flg body - | both `mod` 31 /= 0 = Left (HeaderError ("Header checksum failed")) - | cm /= 8 = Left (HeaderError ("Bad method ("++show cm++")")) - | cinfo > 7 = Left (HeaderError "Window size too big.") - | otherwise = runDeflateM inflate body' +decompress :: L.ByteString -> Either DecompressionError L.ByteString +decompress ifile = run decompressIncremental (L.toChunks ifile) mempty where - cm = cmf .&. 0x0f - cinfo = cmf `shiftR` 4 - fdict = testBit flg 5 --- flevel = flg `shiftR` 6 - -- - body' | fdict = BS.drop 4 body - | otherwise = body - -- - both :: Word16 - both = (fromIntegral cmf `shiftL` 8) .|. fromIntegral flg + run (NeedMore _) [] _ = + Left (DecompressionError "Ran out of data mid-decompression 2.") + run (NeedMore f) (first:rest) acc = + run (f first) rest acc + run (Chunk c m) ls acc = + run m ls (acc <> lazyByteString c) + run Done [] acc = + Right (toLazyByteString acc) + run Done (_:_) _ = + Left (DecompressionError "Finished with data remaining.") + run (DecompError e) _ _ = + Left e + +inflateWithHeaders :: DeflateM () +inflateWithHeaders = + do cmf <- nextByte + flg <- nextByte + let both = fromIntegral cmf `shiftL` 8 .|. fromIntegral flg + cm = cmf .&. 0x0f + cinfo = cmf `shiftR` 4 + fdict = testBit flg 5 +-- flevel = flg `shiftR` 6 + unless ((both :: Word16) `mod` 31 == 0) $ + raise (HeaderError "Header checksum failed") + unless (cm == 8) $ + raise (HeaderError ("Bad compression method: " ++ show cm)) + unless (cinfo <= 7) $ + raise (HeaderError ("Window size too big: " ++ show cinfo)) + when fdict $ replicateM_ 4 nextByte -- just skip them for now (FIXME) + inflate diff --git a/src/Codec/Compression/Zlib/Adler32.hs b/src/Codec/Compression/Zlib/Adler32.hs --- a/src/Codec/Compression/Zlib/Adler32.hs +++ b/src/Codec/Compression/Zlib/Adler32.hs @@ -6,8 +6,8 @@ ) where -import Data.Bits -import Data.Word +import Data.Bits(shiftL, (.|.)) +import Data.Word(Word8, Word16, Word32) data AdlerState = AdlerState { adlerA :: !Word16, adlerB :: !Word16 } diff --git a/src/Codec/Compression/Zlib/Deflate.hs b/src/Codec/Compression/Zlib/Deflate.hs --- a/src/Codec/Compression/Zlib/Deflate.hs +++ b/src/Codec/Compression/Zlib/Deflate.hs @@ -5,42 +5,48 @@ ) where -import Codec.Compression.Zlib.HuffmanTree -import Codec.Compression.Zlib.Monad -import Control.Monad -import Data.Bits -import Data.ByteString.Lazy(ByteString) -import qualified Data.ByteString.Lazy as BS -import Data.Int -import Data.List -import Data.Map.Strict(Map) -import qualified Data.Map.Strict as Map -import Data.Word -import MonadLib(raise) +import Codec.Compression.Zlib.HuffmanTree(HuffmanTree, + createHuffmanTree) +import Codec.Compression.Zlib.Monad(DeflateM, DecompressionError(..), + raise,nextBits,nextCode, + nextBlock,nextWord16,nextWord32, + emitByte,emitBlock,emitPastChunk, + advanceToByte, moveWindow, + finalAdler, finalize) +import Control.Monad(unless, replicateM) +import Data.Array(Array, array, (!)) +import Data.Bits(shiftL, complement) +import Data.Int(Int64) +import Data.List(sortBy) +import Data.IntMap.Strict(IntMap) +import qualified Data.IntMap.Strict as Map +import Data.Word(Word8) +import Numeric(showHex) -inflate :: DeflateM ByteString +inflate :: DeflateM () inflate = - do isFinal <- inflateBlock - if isFinal - then checkChecksum >> finalOutput - else inflate + do fixedLit <- buildFixedLitTree + fixedDist <- buildFixedDistanceTree + go fixedLit fixedDist where - shiftAdd x y = (x `shiftL` 8) .|. fromIntegral y + go fixedLit fixedDist = + do isFinal <- inflateBlock fixedLit fixedDist + moveWindow + if isFinal + then checkChecksum >> finalize + else go fixedLit fixedDist -- checkChecksum = do advanceToByte - rest <- readRest - ourAdler <- finalAdler - let theirAdler = BS.foldl shiftAdd 0 rest - if | BS.length rest < 4 -> raise (ChecksumError "checksum missing") - | BS.length rest > 4 -> raise (FormatError "Ends in middle of file") - | theirAdler /= ourAdler -> raise (ChecksumError "checksum mismatch") - | otherwise -> return () - + ourAdler <- finalAdler + theirAdler <- nextWord32 + unless (theirAdler == ourAdler) $ + raise (ChecksumError ("checksum mismatch: " ++ showHex theirAdler "" ++ + " != " ++ showHex ourAdler "")) -inflateBlock :: DeflateM Bool -inflateBlock = - do bfinal <- nextBit +inflateBlock :: HuffmanTree Int -> HuffmanTree Int -> DeflateM Bool +inflateBlock fixedLitTree fixedDistanceTree = + do bfinal <- (== (1::Word8)) `fmap` nextBits 1 btype <- nextBits 2 case btype :: Word8 of 0 -> -- no compression @@ -52,9 +58,7 @@ emitBlock =<< nextBlock len return bfinal 1 -> -- compressed with fixed Huffman codes - do flt <- fixedLitTree - fdt <- fixedDistanceTree - runInflate flt fdt + do runInflate fixedLitTree fixedDistanceTree return bfinal 2 -> -- compressed with dynamic Huffman codes do hlit <- (257+) `fmap` nextBits 5 @@ -80,21 +84,22 @@ runInflate :: HuffmanTree Int -> HuffmanTree Int -> DeflateM () runInflate litTree distTree = do code <- nextCode litTree - if | code < 256 -> do emitByte (fromIntegral code) - runInflate litTree distTree - | code == 256 -> return () - | code > 256 -> do len <- getLength code - distCode <- nextCode distTree - dist <- getDistance distCode - emitPastChunk dist len - runInflate litTree distTree + case compare code 256 of + LT -> do emitByte (fromIntegral code) + runInflate litTree distTree + EQ -> return () + GT -> do len <- getLength code + distCode <- nextCode distTree + dist <- getDistance distCode + emitPastChunk dist len + runInflate litTree distTree -- ----------------------------------------------------------------------------- getCodeLengths :: HuffmanTree Int -> Int -> Int -> Int -> - Map Int Int -> - DeflateM (Map Int Int) + IntMap Int -> + DeflateM (IntMap Int) getCodeLengths tree n maxl prev acc | n >= maxl = return acc | otherwise = @@ -119,13 +124,11 @@ -- ----------------------------------------------------------------------------- getLength :: Int -> DeflateM Int64 -getLength c = - case Map.lookup c getLengthMap of - Nothing -> raise (DecompressionError ("getLength for bad code: "++show c)) - Just m -> m +getLength c = lengthArray ! c +{-# INLINE getLength #-} -getLengthMap :: Map Int (DeflateM Int64) -getLengthMap = Map.fromList [ +lengthArray :: Array Int (DeflateM Int64) +lengthArray = array (257,285) [ (257, return 3) , (258, return 4) , (259, return 5) @@ -158,13 +161,11 @@ ] getDistance :: Int -> DeflateM Int -getDistance c = - case Map.lookup c getDistanceMap of - Nothing -> raise (DecompressionError ("getDistance for bad code: "++show c)) - Just m -> m +getDistance c = distanceArray ! c +{-# INLINE getDistance #-} -getDistanceMap :: Map Int (DeflateM Int) -getDistanceMap = Map.fromList [ +distanceArray :: Array Int (DeflateM Int) +distanceArray = array (0,29) [ (0, return 1) , (1, return 2) , (2, return 3) @@ -199,15 +200,15 @@ -- ----------------------------------------------------------------------------- -fixedLitTree :: DeflateM (HuffmanTree Int) -fixedLitTree = computeHuffmanTree +buildFixedLitTree :: DeflateM (HuffmanTree Int) +buildFixedLitTree = computeHuffmanTree ([(x, 8) | x <- [0 .. 143]] ++ [(x, 9) | x <- [144 .. 255]] ++ [(x, 7) | x <- [256 .. 279]] ++ [(x, 8) | x <- [280 .. 287]]) -fixedDistanceTree :: DeflateM (HuffmanTree Int) -fixedDistanceTree = computeHuffmanTree [(x,5) | x <- [0..31]] +buildFixedDistanceTree :: DeflateM (HuffmanTree Int) +buildFixedDistanceTree = computeHuffmanTree [(x,5) | x <- [0..31]] -- ----------------------------------------------------------------------------- @@ -217,7 +218,7 @@ Left err -> raise (HuffmanTreeError err) Right x -> return x -computeCodeValues :: Ord a => [(a, Int)] -> [(a, Int, Int)] +computeCodeValues :: [(Int, Int)] -> [(Int, Int, Int)] computeCodeValues vals = Map.foldrWithKey (\ v (l, c) a -> (v,l,c):a) [] codes where valsNo0s = filter (\ (_, b) -> (b /= 0)) vals diff --git a/src/Codec/Compression/Zlib/HuffmanTree.hs b/src/Codec/Compression/Zlib/HuffmanTree.hs --- a/src/Codec/Compression/Zlib/HuffmanTree.hs +++ b/src/Codec/Compression/Zlib/HuffmanTree.hs @@ -6,7 +6,8 @@ ) where -import Data.Bits +import Data.Bits(testBit) +import Data.Word(Word8) data HuffmanTree a = HuffmanNode (HuffmanTree a) (HuffmanTree a) | HuffmanValue a @@ -59,13 +60,14 @@ Left err -> Left err Right l' -> Right (HuffmanNode l' r) -advanceTree :: Bool -> HuffmanTree a -> AdvanceResult a +advanceTree :: Word8 -> HuffmanTree a -> AdvanceResult a advanceTree x node = case node of HuffmanEmpty -> AdvanceError "Tried to advance empty tree!" HuffmanValue _ -> AdvanceError "Tried to advance value!" HuffmanNode l r -> - case if x then r else l of + case if (x == 1) then r else l of HuffmanEmpty -> AdvanceError "Advanced to empty tree!" HuffmanValue y -> Result y t -> NewTree t +{-# INLINE advanceTree #-} diff --git a/src/Codec/Compression/Zlib/Monad.hs b/src/Codec/Compression/Zlib/Monad.hs --- a/src/Codec/Compression/Zlib/Monad.hs +++ b/src/Codec/Compression/Zlib/Monad.hs @@ -1,52 +1,67 @@ {-# LANGUAGE DeriveDataTypeable #-} {-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE MultiWayIf #-} +{-# LANGUAGE Rank2Types #-} module Codec.Compression.Zlib.Monad( DeflateM , runDeflateM + , ZlibDecoder(..) + , raise , DecompressionError(..) -- * Getting data from the input stream. - , nextBit , nextBits , nextByte , nextWord16 + , nextWord32 , nextBlock , nextCode - , readRest -- * Aligning , advanceToByte - -- * Emitting data + -- * Emitting data into the output window , emitByte , emitBlock , emitPastChunk - -- * Getting output + -- * Getting and publishing output , finalAdler - , finalOutput + , moveWindow + , finalize ) where -import Codec.Compression.Zlib.Adler32 -import Codec.Compression.Zlib.HuffmanTree -import Codec.Compression.Zlib.OutputWindow -import Control.Exception(Exception) -import Control.Monad -import Data.Bits -import Data.ByteString.Lazy(ByteString) -import qualified Data.ByteString.Lazy as BS -import Data.Int -import Data.Typeable -import Data.Word -import MonadLib -import Prelude() -import Prelude.Compat +import Codec.Compression.Zlib.Adler32(AdlerState, initialAdlerState, + advanceAdler, finalizeAdler) +import Codec.Compression.Zlib.HuffmanTree(HuffmanTree, advanceTree, + AdvanceResult(..)) +import Codec.Compression.Zlib.OutputWindow(OutputWindow, emptyWindow, + emitExcess, addByte, + addChunk, addOldChunk, + finalizeWindow) +import Control.Exception(Exception) +import Control.Monad(Monad) +import Data.Bits(Bits(..)) +import qualified Data.ByteString as S +import qualified Data.ByteString.Lazy as L +import Data.Int(Int64) +import Data.Typeable(Typeable) +import Data.Word(Word32, Word16, Word8) +import Prelude() +import Prelude.Compat -data DecompressState = DecompressState { +data DecompressionState = DecompressionState { dcsNextBitNo :: !Int , dcsCurByte :: !Word8 , dcsAdler32 :: !AdlerState - , dcsInput :: !ByteString + , dcsInput :: !S.ByteString , dcsOutput :: !OutputWindow } +instance Show DecompressionState where + show dcs = "DecompressionState" + +-- ----------------------------------------------------------------------------- + data DecompressionError = HuffmanTreeError String | FormatError String | DecompressionError String @@ -65,71 +80,133 @@ instance Exception DecompressionError -newtype DeflateM a = DeflateM (StateT DecompressState - (ExceptionT DecompressionError Id) - a) - deriving (Applicative, Functor, Monad) +-- ----------------------------------------------------------------------------- -instance StateM DeflateM DecompressState where - get = DeflateM get - set x = DeflateM (set x) +newtype DeflateM a = DeflateM { + unDeflateM :: DecompressionState -> + (DecompressionState -> a -> ZlibDecoder) -> + ZlibDecoder + } -instance ExceptionM DeflateM DecompressionError where - raise e = DeflateM (lift (raise e)) +instance Applicative DeflateM where + pure x = DeflateM (\ s k -> k s x) -initialState :: ByteString -> DecompressState -initialState bstr = - case BS.uncons bstr of - Nothing -> error "No compressed data to inflate." - Just (f,rest) -> DecompressState 0 f initialAdlerState rest emptyWindow + f <*> x = DeflateM $ \ s1 k -> + unDeflateM f s1 $ \ s2 g -> + unDeflateM x s2 $ \ s3 y -> k s3 (g y) -runDeflateM :: DeflateM a -> ByteString -> Either DecompressionError a -runDeflateM (DeflateM m) i = - case runId (runExceptionT (runStateT (initialState i) m)) of - Left err -> Left err - Right (res, _) -> Right res + m *> n = DeflateM $ \ s1 k -> + unDeflateM m s1 $ \ s2 _ -> unDeflateM n s2 k + {-# INLINE pure #-} + {-# INLINE (<*>) #-} + {-# INLINE (*>) #-} + + +instance Functor DeflateM where + fmap f m = DeflateM (\s k -> unDeflateM m s (\s' a -> k s' (f a))) + {-# INLINE fmap #-} + +instance Monad DeflateM where + {-# INLINE return #-} + return = pure + + {-# INLINE (>>=) #-} + m >>= f = DeflateM $ \ s1 k -> + unDeflateM m s1 $ \ s2 a -> unDeflateM (f a) s2 k + + (>>) = (*>) + {-# INLINE (>>) #-} + +get :: DeflateM DecompressionState +get = DeflateM (\ s k -> k s s) +{-# INLINE get #-} + +set :: DecompressionState -> DeflateM () +set !s = DeflateM (\ _ k -> k s ()) +{-# INLINE set #-} + +raise :: DecompressionError -> DeflateM a +raise e = DeflateM (\ _ _ -> DecompError e) +{-# INLINE raise #-} + +initialState :: DecompressionState +initialState = DecompressionState { + dcsNextBitNo = 8 + , dcsCurByte = 0 + , dcsAdler32 = initialAdlerState + , dcsInput = S.empty + , dcsOutput = emptyWindow + } + -- ----------------------------------------------------------------------------- -nextBit :: DeflateM Bool -nextBit = - do dcs <- get - let v = dcsCurByte dcs `testBit` dcsNextBitNo dcs - set $ advanceBit dcs - return v +data ZlibDecoder = NeedMore (S.ByteString -> ZlibDecoder) + | Chunk L.ByteString ZlibDecoder + | Done + | DecompError DecompressionError + +runDeflateM :: DeflateM () -> ZlibDecoder +runDeflateM m = unDeflateM m initialState (\ _ _ -> Done) +{-# INLINE runDeflateM #-} + +-- ----------------------------------------------------------------------------- + +getNextChunk :: DeflateM () +getNextChunk = DeflateM $ \ st k -> NeedMore (loadChunk st k) where - advanceBit dcs - | dcsNextBitNo dcs == 7 = - case BS.uncons (dcsInput dcs) of - Nothing -> - error "Bit required, but no bits available!" - Just (nextb, rest) -> - dcs{ dcsNextBitNo = 0, dcsCurByte = nextb, dcsInput = rest } - | otherwise = - dcs{ dcsNextBitNo = dcsNextBitNo dcs + 1 } + loadChunk st k bstr = + case S.uncons bstr of + Nothing -> NeedMore (loadChunk st k) + Just (nextb, rest) -> + k st { dcsNextBitNo = 0, dcsCurByte = nextb, dcsInput = rest } () +{-# SPECIALIZE nextBits :: Int -> DeflateM Word8 #-} +{-# SPECIALIZE nextBits :: Int -> DeflateM Int #-} +{-# SPECIALIZE nextBits :: Int -> DeflateM Int64 #-} +{-# INLINE nextBits #-} nextBits :: (Num a, Bits a) => Int -> DeflateM a -nextBits x - | x < 1 = error "nextBits called with x < 1" - | x == 1 = toNum `fmap` nextBit - | otherwise = do cur <- toNum `fmap` nextBit - rest <- nextBits (x - 1) - return ((rest `shiftL` 1) .|. cur) - where - toNum False = 0 - toNum True = 1 +nextBits x = nextBits' x 0 0 +{-# SPECIALIZE nextBits' :: Int -> Int -> Word8 -> DeflateM Word8 #-} +{-# SPECIALIZE nextBits' :: Int -> Int -> Int -> DeflateM Int #-} +{-# SPECIALIZE nextBits' :: Int -> Int -> Int64 -> DeflateM Int64 #-} +{-# INLINE nextBits' #-} +nextBits' :: (Num a, Bits a) => Int -> Int -> a -> DeflateM a +nextBits' !x' !shiftNum !acc + | x' == 0 = return acc + | otherwise = + do dcs <- get + case dcsNextBitNo dcs of + 8 -> case S.uncons (dcsInput dcs) of + Nothing -> + do getNextChunk + nextBits' x' shiftNum acc + Just (nextb, rest) -> + do set dcs{dcsNextBitNo=0,dcsCurByte=nextb,dcsInput=rest} + nextBits' x' shiftNum acc + nextBitNo -> + do let !myBits = min x' (8 - nextBitNo) + !base = dcsCurByte dcs `shiftR` nextBitNo + !mask = complement (0xFF `shiftL` myBits) + !res = fromIntegral (base .&. mask) + !acc' = acc .|. (res `shiftL` shiftNum) + set dcs { dcsNextBitNo=nextBitNo + myBits } + nextBits' (x' - myBits) (shiftNum + myBits) acc' + nextByte :: DeflateM Word8 nextByte = do dcs <- get - case BS.uncons (dcsInput dcs) of - _ | dcsNextBitNo dcs /= 0 -> - nextBits 8 - Nothing -> - error "nextByte called with no more data." - Just (nextb, rest) -> - do set dcs{ dcsNextBitNo = 0, dcsCurByte = nextb, dcsInput = rest } - return (dcsCurByte dcs) + if | dcsNextBitNo dcs == 0 -> do set dcs{ dcsNextBitNo = 8 } + return (dcsCurByte dcs) + | dcsNextBitNo dcs /= 8 -> nextBits 8 -- we're not aligned. sigh. + | otherwise -> case S.uncons (dcsInput dcs) of + Nothing -> getNextChunk >> nextByte + Just (nextb, rest) -> + do set dcs{ dcsNextBitNo = 8, + dcsCurByte = nextb, + dcsInput = rest } + return nextb nextWord16 :: DeflateM Word16 nextWord16 = @@ -137,64 +214,93 @@ high <- fromIntegral `fmap` nextByte return ((high `shiftL` 8) .|. low) -nextBlock :: Integral a => a -> DeflateM ByteString +nextWord32 :: DeflateM Word32 +nextWord32 = + do a <- fromIntegral `fmap` nextByte + b <- fromIntegral `fmap` nextByte + c <- fromIntegral `fmap` nextByte + d <- fromIntegral `fmap` nextByte + return ((a `shiftL` 24) .|. (b `shiftL` 16) .|. (c `shiftL` 8) .|. d) + +nextBlock :: Integral a => a -> DeflateM L.ByteString nextBlock amt = do dcs <- get - unless (dcsNextBitNo dcs == 0) $ - fail "Can't get a block on a non-byte boundary." - let curBlock = BS.cons (dcsCurByte dcs) (dcsInput dcs) - (block, rest) = BS.splitAt (fromIntegral amt) curBlock - case BS.uncons rest of - Nothing -> - fail "Not enough data left after nextBlock." - Just (first, rest') -> - do set dcs{ dcsNextBitNo = 0, dcsCurByte = first, dcsInput = rest' } - return block + if | dcsNextBitNo dcs == 0 -> + do let startByte = dcsCurByte dcs + set dcs{ dcsNextBitNo = 8 } + rest <- nextBlock (amt - 1) + return (L.cons startByte rest) + | dcsNextBitNo dcs == 8 -> + getBlock (fromIntegral amt) (dcsInput dcs) + | otherwise -> + fail "Can't get a block on a non-byte boundary." + where + getBlock len bstr + | len < S.length bstr = do let (mine, rest) = S.splitAt len bstr + dcs <- get + set dcs{ dcsNextBitNo = 8, dcsInput = rest } + return (L.fromStrict mine) + | S.null bstr = do getNextChunk + dcs <- get + let byte1 = dcsCurByte dcs + rest <- getBlock (len - 1) (dcsInput dcs) + return (L.cons byte1 rest) + | otherwise = do rest <- getBlock (len - S.length bstr) S.empty + return (L.fromStrict bstr `L.append` rest) nextCode :: Show a => HuffmanTree a -> DeflateM a nextCode tree = - do b <- nextBit + do b <- nextBits 1 case advanceTree b tree of AdvanceError str -> raise (HuffmanTreeError str) NewTree tree' -> nextCode tree' Result x -> return x - -readRest :: DeflateM ByteString -readRest = - do dcs <- get - return (BS.cons (dcsCurByte dcs) (dcsInput dcs)) +{-# INLINE nextCode #-} advanceToByte :: DeflateM () advanceToByte = do dcs <- get - when (dcsNextBitNo dcs /= 0) $ - case BS.uncons (dcsInput dcs) of - Nothing -> error "Advanced with no bytes left!" - Just (nextb, rest) -> - set dcs{ dcsNextBitNo = 0, dcsCurByte = nextb, dcsInput = rest } + set dcs{ dcsNextBitNo = 8 } emitByte :: Word8 -> DeflateM () emitByte b = do dcs <- get set dcs{ dcsOutput = dcsOutput dcs `addByte` b , dcsAdler32 = advanceAdler (dcsAdler32 dcs) b } +{-# INLINE emitByte #-} -emitBlock :: ByteString -> DeflateM () +emitBlock :: L.ByteString -> DeflateM () emitBlock b = do dcs <- get set dcs { dcsOutput = dcsOutput dcs `addChunk` b - , dcsAdler32 = BS.foldl advanceAdler (dcsAdler32 dcs) b } + , dcsAdler32 = L.foldl advanceAdler (dcsAdler32 dcs) b } emitPastChunk :: Int -> Int64 -> DeflateM () emitPastChunk dist len = do dcs <- get let (output', newChunk) = addOldChunk (dcsOutput dcs) dist len set dcs { dcsOutput = output' - , dcsAdler32 = BS.foldl advanceAdler (dcsAdler32 dcs) newChunk } + , dcsAdler32 = L.foldl advanceAdler (dcsAdler32 dcs) newChunk } +{-# INLINE emitPastChunk #-} finalAdler :: DeflateM Word32 finalAdler = (finalizeAdler . dcsAdler32) `fmap` get -finalOutput :: DeflateM ByteString -finalOutput = (outByteString . dcsOutput) `fmap` get +moveWindow :: DeflateM () +moveWindow = + do dcs <- get + case emitExcess (dcsOutput dcs) of + Nothing -> + return () + Just (builtChunks, output') -> + do set dcs{ dcsOutput = output' } + publishLazy builtChunks +finalize :: DeflateM () +finalize = + do dcs <- get + publishLazy (finalizeWindow (dcsOutput dcs)) + +{-# INLINE publishLazy #-} +publishLazy :: L.ByteString -> DeflateM () +publishLazy lbstr = DeflateM (\ st k -> Chunk lbstr (k st ())) diff --git a/src/Codec/Compression/Zlib/OutputWindow.hs b/src/Codec/Compression/Zlib/OutputWindow.hs --- a/src/Codec/Compression/Zlib/OutputWindow.hs +++ b/src/Codec/Compression/Zlib/OutputWindow.hs @@ -4,66 +4,93 @@ module Codec.Compression.Zlib.OutputWindow( OutputWindow , emptyWindow + , emitExcess + , finalizeWindow , addByte , addChunk , addOldChunk - , outByteString ) where -import Data.ByteString.Builder -import Data.ByteString.Lazy(ByteString) -import qualified Data.ByteString as SBS -import qualified Data.ByteString.Lazy as BS -import Data.Int -import Data.FingerTree -import Data.Foldable(foldMap) -import Data.Monoid -import Data.Word +import Data.ByteString.Builder(Builder, toLazyByteString, word8, + lazyByteString, byteString) +import qualified Data.ByteString as S +import qualified Data.ByteString.Lazy as L +import Data.FingerTree(FingerTree, Measured, ViewL(..), + empty, (|>), split, measure, viewl) +import Data.Foldable.Compat(foldMap) +import Data.Int(Int64) +import Data.Monoid.Compat((<>)) +import Data.Word(Word8) +import Prelude() +import Prelude.Compat -data OutputWindow = OutputWindow { - owCommitted :: !(FingerTree Int SBS.ByteString) - , owRecent :: !Builder - } +type WindowType = FingerTree Int S.ByteString instance Monoid Int where mempty = 0 + {-# INLINE mempty #-} mappend = (+) + {-# INLINE mappend #-} -instance Measured Int SBS.ByteString where - measure = SBS.length +instance Measured Int S.ByteString where + measure = S.length + {-# INLINE measure #-} +data OutputWindow = OutputWindow { + owWindow :: WindowType + , owRecent :: Builder + } + emptyWindow :: OutputWindow emptyWindow = OutputWindow empty mempty +emitExcess :: OutputWindow -> Maybe (L.ByteString, OutputWindow) +emitExcess ow + | totalMeasure < 65536 = Nothing + | otherwise = Just (excess, ow{ owWindow = window' }) + where + window = owWindow ow + totalMeasure = measure window + excessAmount = totalMeasure - 32768 + (excessFT, window') = split (>= excessAmount) window + excess = toLazyByteString (foldMap byteString excessFT) + +finalizeWindow :: OutputWindow -> L.ByteString +finalizeWindow ow = + toLazyByteString (foldMap byteString (owWindow ow) <> owRecent ow) + +-- ----------------------------------------------------------------------------- + addByte :: OutputWindow -> Word8 -> OutputWindow -addByte !ow !b = ow{ owRecent = owRecent ow <> word8 b } +addByte ow b = ow{ owRecent = owRecent ow <> word8 b } -addChunk :: OutputWindow -> ByteString -> OutputWindow -addChunk !ow !bs = ow{ owRecent = owRecent ow <> lazyByteString bs } +addChunk :: OutputWindow -> L.ByteString -> OutputWindow +addChunk ow bs = ow{ owRecent = owRecent ow <> lazyByteString bs } -addOldChunk :: OutputWindow -> Int -> Int64 -> (OutputWindow, ByteString) -addOldChunk !ow !dist !len = (OutputWindow output (lazyByteString chunk), chunk) +addOldChunk :: OutputWindow -> Int -> Int64 -> (OutputWindow, L.ByteString) +addOldChunk ow dist len = (OutputWindow output (lazyByteString chunk), chunk) where - output = owCommitted ow |> BS.toStrict (toLazyByteString (owRecent ow)) + output = L.foldlChunks (|>) (owWindow ow) (toLazyByteString (owRecent ow)) dropAmt = measure output - dist (prev, sme) = split (> dropAmt) output s :< rest = viewl sme - start = SBS.take (fromIntegral len) (SBS.drop (dropAmt-measure prev) s) - len' = fromIntegral len - SBS.length start - (m, rest') = split (> len') rest - middle = BS.toStrict (toLazyByteString (outFinger m)) - end = case viewl rest' of - EmptyL -> SBS.empty - bs2 :< _ -> SBS.take (len' - measure m) bs2 - chunkInf = BS.fromChunks [start, middle, end] `BS.append` chunk - chunk = BS.take len chunkInf - -outFinger :: FingerTree Int SBS.ByteString -> Builder -outFinger = foldMap byteString - -outByteString :: OutputWindow -> ByteString -outByteString ow = - toLazyByteString (outFinger (owCommitted ow) <> owRecent ow) - + start = S.take (fromIntegral len) (S.drop (dropAmt-measure prev) s) + len' = fromIntegral len - S.length start + chunkBase = getChunk rest len' (byteString start) + chunkInf = chunkBase `L.append` chunkInf + chunk = L.take len chunkInf +getChunk :: WindowType -> Int -> Builder -> L.ByteString +getChunk win len acc + | len <= 0 = toLazyByteString acc + | otherwise = + case viewl win of + EmptyL -> toLazyByteString acc + cur :< rest -> + let curlen = S.length cur + in case compare (S.length cur) len of + LT -> getChunk rest (len - curlen) (acc <> byteString cur) + EQ -> toLazyByteString (acc <> byteString cur) + GT -> let (mine, _notMine) = S.splitAt len cur + in toLazyByteString (acc <> byteString mine) diff --git a/test/Test.hs b/test/Test.hs --- a/test/Test.hs +++ b/test/Test.hs @@ -1,6 +1,7 @@ import Codec.Compression.Zlib import Codec.Compression.Zlib.Deflate import Data.ByteString.Lazy(readFile) +import Data.Char (ord) import Data.List(last, isPrefixOf) import Prelude hiding (readFile) import System.FilePath @@ -9,28 +10,28 @@ -- ----------------------------------------------------------------------------- -rfcSimpleTestLengths :: [(Char, Int)] +rfcSimpleTestLengths :: [(Int, Int)] rfcSimpleTestLengths = [ - ('A', 3) - , ('B', 3) - , ('C', 3) - , ('D', 3) - , ('E', 3) - , ('F', 2) - , ('G', 4) - , ('H', 4) + (ord 'A', 3) + , (ord 'B', 3) + , (ord 'C', 3) + , (ord 'D', 3) + , (ord 'E', 3) + , (ord 'F', 2) + , (ord 'G', 4) + , (ord 'H', 4) ] -rfcSimpleTestResults :: [(Char, Int, Int)] +rfcSimpleTestResults :: [(Int, Int, Int)] rfcSimpleTestResults = [ - ('A', 3, 2) -- 010 - , ('B', 3, 3) -- 011 - , ('C', 3, 4) -- 100 - , ('D', 3, 5) -- 101 - , ('E', 3, 6) -- 110 - , ('F', 2, 0) -- 00 - , ('G', 4, 14) -- 1110 - , ('H', 4, 15) -- 1111 + (ord 'A', 3, 2) -- 010 + , (ord 'B', 3, 3) -- 011 + , (ord 'C', 3, 4) -- 100 + , (ord 'D', 3, 5) -- 101 + , (ord 'E', 3, 6) -- 110 + , (ord 'F', 2, 0) -- 00 + , (ord 'G', 4, 14) -- 1110 + , (ord 'H', 4, 15) -- 1111 ] fixedHuffmanLengths :: [(Int, Int)]