src/Lib.hs

   1 {-# LANGUAGE OverloadedStrings #-}
   2 module Lib
   3     ( fileSignatures
   4     , genInstructions
   5     , recreate
   6     ) where
   7
   8 import Control.Monad.State
   9 import Data.Bits (shiftL, shiftR, (.&.))
  10 import qualified Data.ByteString.Char8 as BS
  11 import qualified Data.ByteString.Lazy as BL
  12 import Data.Digest.Adler32 (adler32)
  13 import qualified Data.Map as M
  14 import Data.Word (Word8, Word32)
  15 import qualified Crypto.Hash.MD4 as MD4
  16
  17 type Md4digest       = BS.ByteString
  18 type Adler32checksum = Word32
  19
  20 type Checksum        = (Word32, Int, Int)
  21
  22 type Signature     = (Md4digest, Adler32checksum, Int)
  23
  24 fileSignatures :: BL.ByteString -> Integer -> [Signature]
  25 fileSignatures bs blockSize = zip3 strongsigs weaksigs [0..]
  26   where strongsigs = map blockSig (splitBS bs blockSize)
  27         weaksigs   = map adler32 (splitBS bs blockSize)
  28
  29 splitBS :: BL.ByteString -> Integer -> [BL.ByteString]
  30 splitBS bs blockSize | fromIntegral (BL.length bs) < blockSize = [bs]
  31 splitBS bs blockSize = BL.take (fromIntegral blockSize) bs :
  32   splitBS (BL.drop (fromIntegral blockSize) bs) blockSize
  33
  34 -- compute md4 digest (128 bits)
  35 blockSig :: BL.ByteString -> BS.ByteString
  36 blockSig = MD4.hash . BL.toStrict
  37
  38 weakSig :: BL.ByteString -> Adler32checksum
  39 weakSig = adler32
  40
  41 data Instruction = RChar Word8
  42                  | RBlk  Int
  43                  deriving Show
  44
  45 genInstructions :: [Signature] -> Integer -> BL.ByteString -> [Instruction]
  46 genInstructions f0sigs blockSize fnew =
  47   evalState (go 0 fnew) sig0
  48   where
  49     sig0 = weakSig $ BL.take (fromIntegral blockSize) fnew
  50     go :: Integer -> BL.ByteString -> State Adler32checksum [Instruction]
  51     go startIdx fnew | fnew == BL.empty = return []
  52                      | otherwise = do
  53                          let (blk, blks) = BL.splitAt (fromIntegral blockSize) fnew
  54                              endIdx      = startIdx + fromIntegral (BL.length blk) - 1
  55                          adlerSum <- get
  56                          let matches = M.lookup adlerSum f0AdlerTable >>
  57                                M.lookup (blockSig blk) f0MD4Table
  58                          case matches of
  59                            Just idxs -> do
  60                              -- modify (`adler32Update` blk)
  61                              put $ rollingChecksum (fromIntegral startIdx) (fromIntegral endIdx) fnew
  62                              is <- go (endIdx + 1) blks
  63                              return $ RBlk (head idxs) : is
  64                            Nothing -> do
  65                              let c  = BL.head blk
  66                                  c' = BL.head blk -- FIX (should have been blks)
  67                              -- modify (`adler32Update`  BL.singleton c)
  68                              put $ rollingChecksumUpdate adlerSum c c' startIdx endIdx
  69                              is <- go (startIdx + 1) (BL.tail (blk `mappend` blks))
  70                              return $ RChar c : is
  71     f0AdlerTable = toAdlerMap f0sigs
  72     f0MD4Table   = toMD4Map f0sigs
  73
  74 toAdlerMap :: [Signature] -> M.Map Adler32checksum [Int]
  75 toAdlerMap = foldr f M.empty
  76   where f sig m = let (_, aSig, idx) = sig in M.insertWith (++) aSig [idx] m
  77
  78 toMD4Map :: [Signature] -> M.Map Md4digest [Int]
  79 toMD4Map = foldr f M.empty
  80   where f sig m = let (mSig, _, idx) = sig in M.insertWith (++) mSig [idx] m
  81
  82 recreate :: BL.ByteString -> Integer -> [Instruction] -> BL.ByteString
  83 recreate f0 blockSize ins =
  84   let f0blocks = splitBS f0 blockSize
  85   in
  86     go f0blocks ins
  87   where go f0blocks [] = mempty
  88         go f0blocks (inst:insts) =
  89           case inst of
  90             RBlk i  -> (f0blocks !! i) `mappend` go f0blocks insts
  91             RChar w -> BL.singleton w `mappend` go f0blocks insts
  92
  93 rollingChecksum :: BL.ByteString -> Int -> Int -> Checksum
  94 rollingChecksum bs strtIdx endIdx = (csval, strtIdx, endIdx)
  95   where csval   = a `mod` m + ((fromIntegral b) `mod` m) `shiftL` size
  96         buffer  = map fromIntegral $ take (endIdx - strtIdx) $ drop strtIdx $ BL.unpack bs
  97         indices = map fromIntegral [1..(endIdx - strtIdx + 1)]
  98         a       = sum buffer
  99         b       = sum $ zipWith (*) (reverse indices) buffer
 100         m       = 2^size
 101         size    = 16
 102
 103 -- given the checksum a(k, l) and b(k, l), find checksum a(k+1, l+1), b(k+1, l+1)
 104 rollingChecksumUpdate :: Checksum -> BL.ByteString -> Checksum
 105 rollingChecksumUpdate curCheckSum bs = (csval, oldStrtIdx + 1, oldEndIdx + 1)
 106   where (oldChecksum, oldStrtIdx, oldEndIdx) = curCheckSum
 107         csval   = a `mod` m + ((fromIntegral b) `mod` m) `shiftL` size
 108         m       = 2^size
 109         size    = 16
 110         bold    = oldChecksum `shiftR` size
 111         aold    = oldChecksum .&. (m - 1)
 112         xk      = head $ drop oldStrtIdx $ BL.unpack bs
 113         xlPlus1 = head $ drop (oldEndIdx + 1) $BL.unpack bs
 114         a       = aold - fromIntegral xk + fromIntegral xlPlus1
 115         b       = a + bold - (fromIntegral (oldEndIdx - oldStrtIdx + 1))