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 Signature     = (Md4digest, Adler32checksum, Int)
  21
  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)
  26
  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
  31
  32 -- compute md4 digest (128 bits)
  33 blockSig :: BL.ByteString -> BS.ByteString
  34 blockSig = MD4.hash . BL.toStrict
  35
  36 weakSig :: BL.ByteString -> Adler32checksum
  37 weakSig = adler32
  38
  39 data Instruction = RChar Word8
  40                  | RBlk  Int
  41                  deriving Show
  42
  43 genInstructions :: [Signature] -> Integer -> BL.ByteString -> [Instruction]
  44 genInstructions f0sigs blockSize fnew =
  45   evalState (go 0 fnew) sig0
  46   where
  47     sig0 = weakSig $ BL.take (fromIntegral blockSize) fnew
  48     go :: Integer -> BL.ByteString -> State Adler32checksum [Instruction]
  49     go startIdx fnew | fnew == BL.empty = return []
  50                      | otherwise = do
  51                          let (blk, blks) = BL.splitAt (fromIntegral blockSize) fnew
  52                              endIdx      = startIdx + fromIntegral (BL.length blk) - 1
  53                          adlerSum <- get
  54                          let matches = M.lookup adlerSum f0AdlerTable >>
  55                                M.lookup (blockSig blk) f0MD4Table
  56                          case matches of
  57                            Just idxs -> do
  58                              -- modify (`adler32Update` blk)
  59                              put $ rollingChecksum (fromIntegral startIdx) (fromIntegral endIdx) fnew
  60                              is <- go (endIdx + 1) blks
  61                              return $ RBlk (head idxs) : is
  62                            Nothing -> do
  63                              let c  = BL.head blk
  64                                  c' = BL.head blk -- FIX (should have been blks)
  65                              -- modify (`adler32Update`  BL.singleton c)
  66                              put $ rollingChecksumUpdate adlerSum c c' startIdx endIdx
  67                              is <- go (startIdx + 1) (BL.tail (blk `mappend` blks))
  68                              return $ RChar c : is
  69     f0AdlerTable = toAdlerMap f0sigs
  70     f0MD4Table   = toMD4Map f0sigs
  71
  72 toAdlerMap :: [Signature] -> M.Map Adler32checksum [Int]
  73 toAdlerMap = foldr f M.empty
  74   where f sig m = let (_, aSig, idx) = sig in M.insertWith (++) aSig [idx] m
  75
  76 toMD4Map :: [Signature] -> M.Map Md4digest [Int]
  77 toMD4Map = foldr f M.empty
  78   where f sig m = let (mSig, _, idx) = sig in M.insertWith (++) mSig [idx] m
  79
  80 recreate :: BL.ByteString -> Integer -> [Instruction] -> BL.ByteString
  81 recreate f0 blockSize ins =
  82   let f0blocks = splitBS f0 blockSize
  83   in
  84     go f0blocks ins
  85   where go f0blocks [] = mempty
  86         go f0blocks (inst:insts) =
  87           case inst of
  88             RBlk i  -> (f0blocks !! i) `mappend` go f0blocks insts
  89             RChar w -> BL.singleton w `mappend` go f0blocks insts
  90
  91 rollingChecksum :: Int -> Int -> BL.ByteString -> Word32
  92 rollingChecksum strtIdx endIdx bs = a `mod` m + ((fromIntegral b) `mod` m) `shiftL` mb
  93   where a    = BL.foldl (\acc x -> acc + (fromIntegral x)) 0 bs'
  94         b    = BL.foldl (\acc x -> acc + x) 0 (BL.pack wbs')
  95         bs'  = BL.take (fromIntegral (endIdx - strtIdx + 1)) bs
  96         m    = 2^16
  97         mb   = 16
  98         wbs' = BL.zipWith (*) (BL.pack (reverse (map fromIntegral [1..(endIdx - strtIdx + 1)]))) bs'
  99
 100 checksum :: Int -> Int -> BL.ByteString -> Word32
 101 checksum strtIdx endIdx bs = a `mod` m + ((fromIntegral b) `mod` m) `shiftL` mb
 102   where buffer  = map fromIntegral $ take (endIdx - strtIdx) $ drop strtIdx $ BL.unpack bs
 103         indices = map fromIntegral [1..(endIdx - strtIdx + 1)]
 104         a       = sum buffer
 105         b       = sum $ zipWith (*) (reverse indices) buffer
 106         m       = 2^size
 107         size    = 16
 108
 109 -- given the checksum of bytes from index: startIdx to endIdx, find
 110 -- the checksum for the block from (startIdx + 1 .. endIdx + 1)
 111 rollingChecksumUpdate :: Word32 -> Word8 -> Word8 -> Integer -> Integer -> Word32
 112 rollingChecksumUpdate oldChecksum old new strtIdx endIdx =
 113   let b_Old = (oldChecksum `shiftR` 16) .&. 0xff
 114       a_Old = (oldChecksum .&. 0xff)
 115       a_New = (a_Old - (fromIntegral old) + (fromIntegral new)) `mod` m
 116       b_New = (b_Old - ((fromIntegral endIdx) - (fromIntegral strtIdx) + 1) * (fromIntegral old) + a_New) `mod` m
 117       m     = 2^16
 118   in
 119     a_New `mod` m + (b_New `mod` m) * m