Search code examples
haskellmonad-transformersstate-monad

Why do I need so much memory for WriterT State?


Trying to get to grips with the concepts I am trying to solve an exercise in Haskell using WriterT and State (it's advent of code day 15). For some reason I do not understand I end up using loads of memory and my notebook (just 4G Ram) comes to a halt.

My first idea was to use strictness and sprinkle bangs around - but the issue persists.

Could someone explain me where I did go wrong?

Here's cleaned up code:

{-# LANGUAGE BangPatterns #-}
module Main where
import Control.Monad.State.Strict
import Control.Monad.Writer.Strict

main = do
  let generators = (Generator 65 16807, Generator 8921 48271)
      res1 = compute generators (4*10^7) 
  putStrLn "Answer 1"
  print res1

data Generator = Generator { _value :: Int
                           , _factor :: Int
                           }
    deriving Show

newtype Value = Value Int
  deriving (Show, Eq)

newtype Counter = Counter Int
  deriving (Show, Eq)

instance Monoid Counter where
  mempty = Counter 0
  mappend (Counter !a) (Counter !b) = Counter (a+b)

generate :: Generator -> (Value, Generator)
generate (Generator v f) = (Value newval, Generator newval f)
  where newval = (v * f) `mod` 2147483647

agree (Value a) (Value b) = (a `mod` mf) == (b `mod` mf)
  where mf = 2^16

oneComp :: State (Generator, Generator) Bool
oneComp = do
  (!ga, !gb) <- get
  let (va, gan) = generate ga
      (vb, gbn) = generate gb
      !ag = agree va vb
  put (gan, gbn)
  pure ag

counterStep :: WriterT Counter (State (Generator, Generator)) ()
counterStep = do
  !ag <- lift oneComp
  when ag $ tell (Counter 1)

afterN :: Int -> WriterT Counter (State (Generator, Generator)) ()
afterN n = replicateM_ n counterStep

compute s0 n = evalState (execWriterT (afterN n)) s0

I compile it with stack. The entry in the cabal file is:

executable day15
  hs-source-dirs:      app
  main-is:             day15.hs
  ghc-options:         -threaded -rtsopts -with-rtsopts=-N 
  build-depends:       base
                     , advent
                     , hspec
                     , mtl
  default-language:    Haskell2010

update

I had a little more time and followed the suggestion to make Generator strict. However still something is using too much memory.

Here's the part of the prof file that I think may be relevant.

            Fri Dec 15 16:28 2017 Time and Allocation Profiling Report  (Final)

       day15 +RTS -N -p -RTS

    total time  =       71.66 secs   (71662 ticks @ 1000 us, 1 processor)
    total alloc = 17,600,423,088 bytes  (excludes profiling overheads)

COST CENTRE    MODULE    SRC                          %time %alloc

afterN         Main      app/day15.hs:79:1-36          41.1   20.0
mappend        Main      app/day15.hs:51:3-51          31.0    3.6
oneComp        Main      app/day15.hs:(64,1)-(71,9)     9.2   49.1
generate.(...) Main      app/day15.hs:55:9-42           8.5   14.5

Solution

  • The cause is likely to be the WriterT layer.

    Even the "strict" WriterT is completely lazy in the accumulator —it is strict in another sense unrelated to the accumulator.

    For example, this program runs without errors:

    import Data.Monoid
    import Control.Monad.Trans.Writer
    import Control.Exception
    
    main :: IO ()
    main = do
      let (x,_) = runWriter $ do
            tell $ Sum (1::Float)
            tell (throw $ AssertionFailed "oops")
            tell (error "oops")
            tell undefined
            tell (let z = z in z)
            return True
      print x
    

    Furthermore, it is impossible to "strictify" the accumulator from within WriterT, because there's no way to get to it.

    For long computations, thunks will accumulate and consume a lot of memory.

    One solution is to store the counter in a StateT layer instead. The strict modify' function is helpful here.


    Using StateT for an append-only accumulator is a bit unsatisfactory though. Another option is to use Accum with judiciously positioned BangPatterns. This program throws an error:

    import Control.Monad.Trans.Accum
    
    main :: IO ()
    main = do
      let (x,_) = flip runAccum mempty $ do
            add $ Sum (1::Float)
            add $ error "oops"
            !_ <- look
            return True
      print x
    

    Accum is like a Writer that lets you access the accumulator. It doesn't let you change it at will, only add to it. But getting hold of it is enough to introduce strictness.