1 {-# LANGUAGE OverloadedStrings #-}
8 import Control.Monad.State
9 import Data.Bits (shiftR, (.&.))
10 import qualified Data.ByteString.Char8 as BS
11 import qualified Data.ByteString.Lazy as BL
12 import Data.Digest.Adler32 (adler32, adler32Update)
13 import qualified Data.Map as M
14 import Data.Word (Word8, Word32)
15 import qualified Crypto.Hash.MD4 as MD4
17 type Md4digest = BS.ByteString
18 type Adler32checksum = Word32
20 type Signature = (Md4digest, Adler32checksum, Int)
22 fileSignatures :: BL.ByteString -> Integer -> [Signature]
23 fileSignatures bs blockSize = zip3 strongsigs weaksigs [0..]
24 where strongsigs = map blockSig (splitBS bs blockSize)
25 weaksigs = map adler32 (splitBS bs blockSize)
27 splitBS :: BL.ByteString -> Integer -> [BL.ByteString]
28 splitBS bs blockSize | fromIntegral (BL.length bs) < blockSize = [bs]
29 splitBS bs blockSize = BL.take (fromIntegral blockSize) bs :
30 splitBS (BL.drop (fromIntegral blockSize) bs) blockSize
32 -- compute md4 digest (128 bits)
33 blockSig :: BL.ByteString -> BS.ByteString
34 blockSig = MD4.hash . BL.toStrict
36 weakSig :: BL.ByteString -> Adler32checksum
39 data Instruction = RChar Word8
43 genInstructions :: [Signature] -> Integer -> BL.ByteString -> [Instruction]
44 genInstructions f0sigs blockSize fnew =
45 evalState (go 0 fnew) sig0
47 sig0 = weakSig $ BL.take (fromIntegral blockSize) fnew
48 go :: Integer -> BL.ByteString -> State Adler32checksum [Instruction]
49 go startIdx fnew | fnew == BL.empty = return []
51 let (blk, blks) = BL.splitAt (fromIntegral blockSize) fnew
52 endIdx = startIdx + fromIntegral (BL.length blk) - 1
54 let matches = M.lookup adlerSum f0AdlerTable >>
55 M.lookup (blockSig blk) f0MD4Table
58 modify (`adler32Update` blk)
59 is <- go (endIdx + 1) blks
60 return $ RBlk (head idxs) : is
63 modify (`adler32Update` BL.singleton c)
64 is <- go (startIdx + 1) (BL.tail (blk `mappend` blks))
66 f0AdlerTable = toAdlerMap f0sigs
67 f0MD4Table = toMD4Map f0sigs
69 toAdlerMap :: [Signature] -> M.Map Adler32checksum [Int]
70 toAdlerMap = foldr f M.empty
71 where f sig m = let (_, aSig, idx) = sig in M.insertWith (++) aSig [idx] m
73 toMD4Map :: [Signature] -> M.Map Md4digest [Int]
74 toMD4Map = foldr f M.empty
75 where f sig m = let (mSig, _, idx) = sig in M.insertWith (++) mSig [idx] m
77 recreate :: BL.ByteString -> Integer -> [Instruction] -> BL.ByteString
78 recreate f0 blockSize ins =
79 let f0blocks = splitBS f0 blockSize
82 where go f0blocks [] = mempty
83 go f0blocks (inst:insts) =
85 RBlk i -> (f0blocks !! i) `mappend` go f0blocks insts
86 RChar w -> BL.singleton w `mappend` go f0blocks insts
88 rollingChecksum :: Int -> Int -> BL.ByteString -> Word32
89 rollingChecksum strtIdx lenBS bs = a `mod` m + ((fromIntegral b) `mod` m) * m
90 where a = BL.foldl (\acc x -> acc + (fromIntegral x)) 0 bs'
91 b = BL.foldl (\acc x -> acc + x) 0 (BL.pack wbs')
92 bs' = BL.take (fromIntegral lenBS) $ BL.drop (fromIntegral strtIdx) bs
94 wbs' = BL.zipWith (*) (BL.pack (reverse (map fromIntegral [1..(lenBS - strtIdx + 1)]))) bs'
96 -- given the checksum of bytes from index: startIdx to endIdx, find
97 -- the checksum for the block from (startIdx + 1 .. endIdx + 1)
98 rollingChecksumUpdate :: Word32 -> Word8 -> Word8 -> Integer -> Integer -> Word32
99 rollingChecksumUpdate oldChecksum old new strtIdx endIdx =
100 let b_Old = (oldChecksum `shiftR` 16) .&. 0xff
101 a_Old = (oldChecksum .&. 0xff)
102 a_New = (a_Old - (fromIntegral old) + (fromIntegral new)) `mod` m
103 b_New = (b_Old - ((fromIntegral endIdx) - (fromIntegral strtIdx) + 1) * (fromIntegral old) + a_New) `mod` m
106 a_New `mod` m + (b_New `mod` m) * m