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 qualified Data.ByteString.Char8 as BS
  10 import qualified Data.ByteString.Lazy as BL
  11 import Data.Digest.Adler32 (adler32, adler32Update)
  12 import qualified Data.Map as M
  13 import Data.Word (Word8, Word32)
  14 import qualified Crypto.Hash.MD4 as MD4
  15
  16 type Md4digest       = BS.ByteString
  17 type Adler32checksum = Word32
  18
  19 type Signature     = (Md4digest, Adler32checksum, Int)
  20
  21 fileSignatures :: BL.ByteString -> Integer -> [Signature]
  22 fileSignatures bs blockSize = zip3 strongsigs weaksigs [0..]
  23   where strongsigs = map blockSig (splitBS bs blockSize)
  24         weaksigs   = map adler32 (splitBS bs blockSize)
  25
  26 splitBS :: BL.ByteString -> Integer -> [BL.ByteString]
  27 splitBS bs blockSize | fromIntegral (BL.length bs) < blockSize = [bs]
  28 splitBS bs blockSize = BL.take (fromIntegral blockSize) bs :
  29   splitBS (BL.drop (fromIntegral blockSize) bs) blockSize
  30
  31 -- compute md4 digest (128 bits)
  32 blockSig :: BL.ByteString -> BS.ByteString
  33 blockSig = MD4.hash . BL.toStrict
  34
  35 weakSig :: BL.ByteString -> Adler32checksum
  36 weakSig = adler32
  37
  38 data Instruction = RChar Word8
  39                  | RBlk  Int
  40                  deriving Show
  41
  42 genInstructions :: [Signature] -> Integer -> BL.ByteString -> [Instruction]
  43 genInstructions f0sigs blockSize fnew =
  44   evalState (go fnew) sig0
  45   where
  46     sig0 = weakSig $ BL.take (fromIntegral blockSize) fnew
  47     go :: BL.ByteString -> State Adler32checksum [Instruction]
  48     go fnew | fnew == BL.empty = return []
  49             | otherwise = do
  50                 let (blk, blks) = BL.splitAt (fromIntegral blockSize) fnew
  51                 adlerSum <- get
  52                 let matches = M.lookup adlerSum f0AdlerTable >>
  53                       M.lookup (blockSig blk) f0MD4Table
  54                 case matches of
  55                   Just idxs -> do
  56                     modify (`adler32Update` blk)
  57                     is <- go blks
  58                     return $ RBlk (head idxs) : is
  59                   Nothing -> do
  60                     let c = BL.head blk
  61                     modify (`adler32Update`  BL.singleton c)
  62                     is <- go (BL.tail (blk `mappend` blks))
  63                     return $ RChar c : is
  64     f0AdlerTable = toAdlerMap f0sigs
  65     f0MD4Table   = toMD4Map f0sigs
  66
  67 toAdlerMap :: [Signature] -> M.Map Adler32checksum [Int]
  68 toAdlerMap = foldr f M.empty
  69   where f sig m = let (_, aSig, idx) = sig in M.insertWith (++) aSig [idx] m
  70
  71 toMD4Map :: [Signature] -> M.Map Md4digest [Int]
  72 toMD4Map = foldr f M.empty
  73   where f sig m = let (mSig, _, idx) = sig in M.insertWith (++) mSig [idx] m
  74
  75 recreate :: BL.ByteString -> Integer -> [Instruction] -> BL.ByteString
  76 recreate f0 blockSize ins =
  77   let f0blocks = splitBS f0 blockSize
  78   in
  79     go f0blocks ins
  80   where go f0blocks [] = mempty
  81         go f0blocks (inst:insts) =
  82           case inst of
  83             RBlk i  -> (f0blocks !! i) `mappend` go f0blocks insts
  84             RChar w -> BL.singleton w `mappend` go f0blocks insts
  85