parsing binary data in haskell

slim.hs

module Main where 

import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as BL
import Data.Binary.Strict.Get
import qualified Data.Binary.Strict.BitGet as BG
import Data.Word
import Control.Monad
import Control.Applicative
import Data.Binary.Put

data Header = Header
  { _signature :: Word16
  , _fileSize :: Int
  , _res1 :: Word16
  , _res2 :: Word16
  , _PAOffset :: Int
  } deriving (Eq, Show)

data DIBHeader = DIBHeader
  { _headerSize      :: Int
  , _width           :: Int
  , _height          :: Int
  , _colorPlanes     :: Word16
  , _bpp             :: Word16
  , _compression     :: Int
  , _rawSize         :: Int
  , _ppmH            :: Int
  , _ppmV            :: Int
  , _paletteSize     :: Int
  , _importantColors :: Int
  } deriving (Eq, Show)

data ColorTable = ColorTable
  { _colorTable :: B.ByteString
  } deriving (Eq, Show)

data PixelArray = PixelArray 
  { _pixelArray :: [Word8]
  } deriving (Eq, Show)

getInt32 = liftM fromIntegral getWord32le

parseHeader :: Get Header
parseHeader = do
  signature <- getWord16le
  size <- getInt32
  res1 <- getWord16le
  res2 <- getWord16le
  offset <- getInt32
  return $ Header signature size res1 res2 offset

parseDIBHeader = 
  DIBHeader <$> 
  getInt32 <*> -- size
  getInt32 <*> -- width
  getInt32 <*> -- height
  getWord16le <*> -- planes
  getWord16le <*> -- bpp
  getInt32 <*> -- compression
  getInt32 <*> -- raw size
  getInt32 <*> -- ppm horizontal
  getInt32 <*> -- ppm vertical
  getInt32 <*> -- palette size 
  getInt32     -- important colors

parseColorTable :: Int -> Get ColorTable
parseColorTable size = ColorTable <$> getByteString size

parseFile :: Get (Header, DIBHeader)
parseFile = liftM2 (,) parseHeader parseDIBHeader

parseRow :: Int -> Int -> BG.BitGet [Word8]
parseRow bpp width = replicateM width $ BG.getAsWord8 bpp

parsePA :: Int -> Int -> Int -> Get PixelArray
parsePA bpp width height = do
  let bytesPerRow = 4 * ((bpp * width + 31) `div` 32)
  rows <- replicateM height $ getByteString bytesPerRow
  let Right bytes = mapM (\x -> BG.runBitGet x $ parseRow bpp width) rows
  return . PixelArray $ concat bytes

data BMP = BMP
  { _header :: Header
  , _dibHeader :: DIBHeader
  , _colors :: ColorTable
  , _pa :: PixelArray
  } deriving (Eq, Show)

putInt :: Int -> Put
putInt = putWord32le . fromIntegral

putHeader :: Header -> Put
putHeader (Header sig size res1 res2 off) = do
  putWord16le sig
  putInt size
  putWord16le res1
  putWord16le res2
  putInt off

putDIB :: DIBHeader -> Put
putDIB header = do
  putInt $ _headerSize header
  putInt $ _width header
  putInt $ _height header
  putWord16le $ _colorPlanes header
  putWord16le $ _bpp header
  putInt $ _compression header
  putInt $ _rawSize header
  putInt $ _ppmH header
  putInt $ _ppmV header
  putInt $ _paletteSize header
  putInt $ _importantColors header

putColors :: ColorTable -> Put
putColors (ColorTable bytes) = putByteString bytes

-- XXX: assumes multiple of 32 width
putPA :: PixelArray -> Put
putPA (PixelArray pa) = mapM_ putWord8 pa

putBMP :: BMP -> Put
putBMP (BMP header dib colors pa) = do
  putHeader header
  putDIB dib
  putColors colors
  putPA pa

main :: IO ()
main = do
  contents <- B.readFile "picnic.bmp"
  let (Right headers, rest) = runGet parseFile contents
  let (header, dibHeader) = headers
  print headers
  let paOffset = _PAOffset header 
  let width = _width dibHeader
  let height = _height dibHeader
  let bpp = fromIntegral $ _bpp dibHeader
  let (Right colorTable, rest') = runGet (parseColorTable $ _paletteSize dibHeader) rest
  let (Right pixelArray, rest) = runGet (skip paOffset >>= \_ -> parsePA bpp width height) contents

  let dibHeader' = dibHeader { _bpp = 8 }
  let bmp = BMP header dibHeader' colorTable pixelArray
  BL.writeFile "picnic-good.bmp" $ runPut $ putBMP bmp
--  print pixelArray
  print rest