Search code examples
haskellmonadsstate-monad

Haskell State Monad and Binary not outputting everything


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"

Solution

  • 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"