Haskell State Monad and Binary not outputting everything

Question

I've been playing around with some simple binary encoding and it seemed to be working correctly for the most part, up until I added the state monad. The plan was to use the state to keep a lookup table of what I have written to the bytestring so far, and then write out offsets to previous instances of strings rather than duplicate them.

I got everything type checked and running, but then I noticed that it was only writing out the final instruction in the chain. I changed to using the Control.Monad.State.Strict but that made no difference, so I suspect I'm making a fundamental error somewhere else, but I'm not sure where - I've trimmed the code down to the basic functionality. Also, is there a more idiomatic way of doing this?

{-# LANGUAGE OverloadedStrings #-}


import           Control.Applicative
import qualified Control.Monad.State.Strict as S
import           Data.Binary.Put
import qualified Data.ByteString            as BS
import qualified Data.ByteString.Lazy       as BL

data SState = SState {
   wsPosition :: Int
   -- plus whatever else
}

initialState = SState 0
type StatePut = S.State SState Put

class StateBinary a where
   sput :: a -> StatePut

incPos :: Int -> S.State SState ()
incPos amnt = do
   (SState p) <- S.get
   S.put $ SState (p + amnt)

writeSized :: Int -> (a -> Put) -> a -> StatePut
writeSized n f x = do
                    incPos n
                    return (f x)

writeInt32 :: Int -> StatePut
writeInt32 = writeSized 32 putWord32be . fromIntegral

writeBS :: BS.ByteString -> StatePut
writeBS b = writeSized (BS.length b) putByteString b

data SomeData = SomeData {
    sdName :: BS.ByteString
  , sdAge  :: Int
  , sdN    :: Int
} deriving (Show, Eq)

instance StateBinary SomeData where
    sput (SomeData nm a n) = do
           writeBS nm
           writeInt32 a
           writeInt32 n

testData = SomeData "TestName" 30 100

runSPut :: StateBinary a => a -> BL.ByteString
runSPut a = runPut $ S.evalState (sput a) initialState

-- runSPut testData returns "\NUL\NUL\NULd"

Answer 1

The problem is that writeSized is not actually writing to the bytestring. return only wraps the Put computation into the state monad without actually running it. There might be smarter ways to solve it, but the obvious one would be taking advantage of the fact that Put (actually PutM) is a monad and using monad transformers to compose it with the state monad:

{-# LANGUAGE OverloadedStrings #-}


import           Control.Applicative
import qualified Control.Monad.State.Strict as S
import           Data.Binary.Put
import qualified Data.ByteString            as BS
import qualified Data.ByteString.Lazy       as BL

data SState = SState {
   wsPosition :: Int
   -- plus whatever else
}

initialState = SState 0
-- S.StateT SState PutM is a composed monad, with a state layer above PutM.
type StatePut = S.StateT SState PutM ()

class StateBinary a where
   sput :: a -> StatePut

incPos :: Int -> StatePut
incPos amnt = do
   (SState p) <- S.get
   S.put $ SState (p + amnt)

writeSized :: Int -> (a -> Put) -> a -> StatePut
writeSized n f x = do
                    incPos n
                    -- lift runs a computation in the underlying monad.
                    S.lift (f x)

writeInt32 :: Int -> StatePut
writeInt32 = writeSized 32 putWord32be . fromIntegral

writeBS :: BS.ByteString -> StatePut
writeBS b = writeSized (BS.length b) putByteString b

data SomeData = SomeData {
    sdName :: BS.ByteString
  , sdAge  :: Int
  , sdN    :: Int
} deriving (Show, Eq)

instance StateBinary SomeData where
    sput (SomeData nm a n) = do
           writeBS nm
           writeInt32 a
           writeInt32 n

testData = SomeData "TestName" 30 100

runSPut :: StateBinary a => a -> BL.ByteString
runSPut a = runPut $ S.evalStateT (sput a) initialState

-- *Main> runSPut testData
-- "TestName\NUL\NUL\NUL\RS\NUL\NUL\NULd"

Answer 2

You can use a bytestring Builder (edit: now using the one from binary instead of from bytestring):

{-# LANGUAGE OverloadedStrings #-}

import           Data.Monoid
import qualified Data.Binary                as B
import qualified Data.Binary.Builder        as BU
import qualified Data.ByteString            as BS
import qualified Data.ByteString.Lazy       as BL

data SomeData = SomeData {
    sdName :: BS.ByteString
  , sdAge  :: Int
  , sdN    :: Int
} deriving (Show, Eq)

testData :: SomeData
testData  = SomeData "TestName" 30 100

renderData :: SomeData -> BU.Builder
renderData (SomeData n a i) = mconcat $
  BU.fromByteString n : map (BU.fromLazyByteString . B.encode) [a,i]

test :: BL.ByteString
test = BU.toLazyByteString . renderData $ testData

The idea is that the introduction (BU.fromX) and appending operations are O(1), so you only pay at the end when you convert back to a ByteString.