# LANGUAGE BangPatterns      #
# LANGUAGE TypeApplications #

module Digest (readDict, extract)
  where

import MD5
import Config

import Control.Applicative
import Control.Monad
import Data.ByteString.Internal                                     ( w2c )
import Data.Label
import Data.Word
import Prelude
import qualified Data.ByteString                                    as S
import qualified Data.ByteString.Lazy.Char8                         as L
import qualified Data.Serialize                                     as S

import Data.Array.Accelerate.Array.Data                             as A
import Data.Array.Accelerate.Sugar.Array
import Data.Array.Accelerate.Sugar.Elt
import Data.Array.Accelerate.Sugar.Shape
import qualified Data.Array.Accelerate                              as A
import qualified Data.Array.Accelerate.Representation.Array         as R


-- MD5 block sizes
--
blockSizeBits, blockSizeBytes, blockSizeWords :: Int
blockSizeBits  = 512
blockSizeBytes = blockSizeBits `div` 8
blockSizeWords = blockSizeBytes `div` 4


-- Create an MD5 block from the given message. This appends The '1' bit to the
  message , pads the block with zeros until the length in bits is 448 , and
  finally appends the length in bits ( mod 2 ^ 64 ) .
--
md5block :: L.ByteString -> S.ByteString
md5block msg = do
  let
      len               = fromIntegral (L.length msg)
      lenBits           = 8 * fromIntegral len
      lenZeroPad
        | len + 1 <= blockSizeBytes - 8 = (blockSizeBytes - 8) - (len + 1)
        | otherwise                     = (2 * blockSizeBytes - 8) - (len + 1)
  --
  S.runPut $! do
    S.putLazyByteString msg
    S.putWord8 0x80
    mapM_ S.putWord8 (replicate lenZeroPad 0)
    S.putWord64le lenBits


-- Create a dictionary of blocks ready to digest from the given bytestring. This
  reads one entry per line . Because we only do a single MD5 chunk , we discard
  any entries with ( length > blockSizeBytes - 8 = 55 ) characters . Because of
-- this restriction, we also need to a precise count of how many entries will be
-- in the final array, since it is stored column-major and hence there is no
-- easy way to truncate it.
--
readDict :: Config -> FilePath -> IO Dictionary
readDict c fp = do
  entries <- length       . chunk <$> L.readFile fp
  blocks  <- map md5block . chunk <$> L.readFile fp

  let sh        = Z :. blockSizeWords :. entries
      (adata,_) = runArrayData $ do
        arr <- newArrayData (eltR @Word32) (size sh)

        let go !_ []     = return ()
            go !n (b:bs) = do
              foldM_ (\i w -> do writeArrayData (eltR @Word32) arr (toIndex sh (Z:.i:.n)) (fromElt w)
                                 return (i+1)) 0 (bytes b)
              go (n+1) bs

        go 0 blocks
        return (arr, 0 :: Word32)

  adata `seq` return $ Array (R.Array (fromElt sh) adata)

  where
    chunk = maybe id take (get configMaxWords c)
          . drop (get configSkipWords c)
          . filter (\w -> fromIntegral (L.length w) < blockSizeBytes - 8)
          . L.lines

    bytes = either error id
          . S.runGet (replicateM blockSizeWords S.getWord32le)

-- Extract a word from the dictionary at a given index
--
extract :: Dictionary -> Int -> L.ByteString
extract dict i
  | i > n       = error "extract: index too large"
  | otherwise   = L.takeWhile (/= w2c 0x80) bytes
  where
    Z :. _ :. n = A.arrayShape dict
    bytes       = S.runPutLazy $
      forM_ [0 .. blockSizeWords-1] $ \c -> S.putWord32le (dict `A.indexArray` (Z:.c:.i))


 MD5 block sizes

 Create an MD5 block from the given message. This appends The '1' bit to the


 Create a dictionary of blocks ready to digest from the given bytestring. This
 this restriction, we also need to a precise count of how many entries will be
 in the final array, since it is stored column-major and hence there is no
 easy way to truncate it.

 Extract a word from the dictionary at a given index

# LANGUAGE BangPatterns      #
# LANGUAGE TypeApplications #

module Digest (readDict, extract)
  where

import MD5
import Config

import Control.Applicative
import Control.Monad
import Data.ByteString.Internal                                     ( w2c )
import Data.Label
import Data.Word
import Prelude
import qualified Data.ByteString                                    as S
import qualified Data.ByteString.Lazy.Char8                         as L
import qualified Data.Serialize                                     as S

import Data.Array.Accelerate.Array.Data                             as A
import Data.Array.Accelerate.Sugar.Array
import Data.Array.Accelerate.Sugar.Elt
import Data.Array.Accelerate.Sugar.Shape
import qualified Data.Array.Accelerate                              as A
import qualified Data.Array.Accelerate.Representation.Array         as R


blockSizeBits, blockSizeBytes, blockSizeWords :: Int
blockSizeBits  = 512
blockSizeBytes = blockSizeBits `div` 8
blockSizeWords = blockSizeBytes `div` 4


  message , pads the block with zeros until the length in bits is 448 , and
  finally appends the length in bits ( mod 2 ^ 64 ) .
md5block :: L.ByteString -> S.ByteString
md5block msg = do
  let
      len               = fromIntegral (L.length msg)
      lenBits           = 8 * fromIntegral len
      lenZeroPad
        | len + 1 <= blockSizeBytes - 8 = (blockSizeBytes - 8) - (len + 1)
        | otherwise                     = (2 * blockSizeBytes - 8) - (len + 1)
  S.runPut $! do
    S.putLazyByteString msg
    S.putWord8 0x80
    mapM_ S.putWord8 (replicate lenZeroPad 0)
    S.putWord64le lenBits


  reads one entry per line . Because we only do a single MD5 chunk , we discard
  any entries with ( length > blockSizeBytes - 8 = 55 ) characters . Because of
readDict :: Config -> FilePath -> IO Dictionary
readDict c fp = do
  entries <- length       . chunk <$> L.readFile fp
  blocks  <- map md5block . chunk <$> L.readFile fp

  let sh        = Z :. blockSizeWords :. entries
      (adata,_) = runArrayData $ do
        arr <- newArrayData (eltR @Word32) (size sh)

        let go !_ []     = return ()
            go !n (b:bs) = do
              foldM_ (\i w -> do writeArrayData (eltR @Word32) arr (toIndex sh (Z:.i:.n)) (fromElt w)
                                 return (i+1)) 0 (bytes b)
              go (n+1) bs

        go 0 blocks
        return (arr, 0 :: Word32)

  adata `seq` return $ Array (R.Array (fromElt sh) adata)

  where
    chunk = maybe id take (get configMaxWords c)
          . drop (get configSkipWords c)
          . filter (\w -> fromIntegral (L.length w) < blockSizeBytes - 8)
          . L.lines

    bytes = either error id
          . S.runGet (replicateM blockSizeWords S.getWord32le)

extract :: Dictionary -> Int -> L.ByteString
extract dict i
  | i > n       = error "extract: index too large"
  | otherwise   = L.takeWhile (/= w2c 0x80) bytes
  where
    Z :. _ :. n = A.arrayShape dict
    bytes       = S.runPutLazy $
      forM_ [0 .. blockSizeWords-1] $ \c -> S.putWord32le (dict `A.indexArray` (Z:.c:.i))

