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