CODEWITHSUNDEEP
loopshaskell
Janthelme
Janthelme

Reputation: 999

Better way to write loops in IO

Given a png file, I am trying to get the list of its chunks' offset and sizes.

In a nutshell, png files are made of chunks, and each chunk is made of three 4-byte fields plus one variable-length field (the chunk's data field). The size of the data field is stored in the first 4-byte field (called the "length" field).

Therefore given the current chunk's offset and size, (ofs, sz), one derives the next chunk's offset and size, (ofs', sz'), that way :

ofs' = ofs + sz

read sz' at offset = ofs'

Given the initial chunk's offset and size, always (0, 8) in png files, one can loop through the file until one reaches its end. That's how I did it :

import Data.Word
import qualified Data.ByteString.Lazy as BS
import Data.Binary.Get

size :: BS.ByteString -> Int -> IO (Int)
size bytes offset = do
    let ln = runGet (do skip offset
                        l <- getWord32be
                        return l)
                    bytes
    return $ 3*4 + fromIntegral ln

offsetSizes :: Int -> BS.ByteString -> [(Int, Int)] -> IO [(Int, Int)]
offsetSizes fLen bytes oss = do
        let (offset, sz) = last oss
            offset' = offset + sz
        sz' <- size bytes offset'
        let nextOffset = offset' + sz'
        if nextOffset < fLen then offsetSizes fLen bytes $ oss ++ [(offset', sz')]
                              else return oss
main = do
    contents <- BS.readFile "myfile.png"
    let fLen = fromIntegral $ BS.length contents :: Int

    ofszs <- offsetSizes fLen contents [(0,8)]
    putStrLn $ "# of chunks: " ++ (show $ length ofszs)
    putStrLn $ "chunks [(offset,size)]: " ++ show ofszs

My question : I am not really satisfied with the loop. I was wondering whether there is a more idiomatic way to achieve the same in Haskell?

Upvotes: 2

Views: 158

Answers (1)

Cactus
Cactus

Reputation: 27626

offsetSizes is repeatedly fed some state (the (offset, sz) pair) to produce a new pair or finish. All the pairs created are collected into a list.

This recursion scheme is captured by unfoldrM from the monad-loops package, allowing you to write offsetSizes as

offsetSizes :: Int -> BS.ByteString -> IO [(Int, Int)]
offsetSizes fLen bytes = unfoldrM step (0, 8)
  where
    step (offset, sz) = do
        let offset' = offset + sz
        sz' <- size bytes offset'
        let state' = (offset', sz')
        return $ if offset' + sz' < fLen then Just (state', state') else Nothing

Upvotes: 1

Related Questions