How can I use the State monad effectively in Haskell?

Question

I have written a simple fibonacci function in Haskell that uses the State monad to save solutions to overlapping sub-problems. Now despite the function working I am not really happy with the solution, especially with the last map lookup. How can I improve my solution? Thank you.

module Lib (
    fib
) where

import Control.Monad.State
import Data.Map (Map)
import qualified Data.Map as Map
import Data.Maybe

type DP = Map Int Int

fib :: Int -> State DP Int
fib 0 = return 1
fib 1 = return 1
fib x = do
    solved <- get
    let m = Map.lookup x solved
    case m of
            Nothing -> do
                prev1 <- fib (x - 1)
                prev2 <- fib (x - 2)  
                put $ Map.insert x (prev1 + prev2) solved
            Just value -> put $ Map.insert x value solved
    gets (fromMaybe 0 . Map.lookup x)

Answer 1

I won't be reviewing this code for efficiency, since I'm not quite so comfortable with analyzing some combination of State monad, immutable data structures, and laziness.

However, there are a few obvious changes you can make to your code.

Passing State Properly

This is what Franky was getting at with his comment, and it's a bit of a tricky bug. When you do

put $ Map.insert x (prev1 + prev2) solved

you use the solved that you get at the beginning of the function invocation. But your recursion goes from top-to-bottom, so if you're computing fib 10, solved is the empty map when you are getting the final answer. Essentially, your memo map keeps getting overwritten with maps that have less information. The fix is pretty easy.

prev1 <- fib (x - 1)
prev2 <- fib (x - 2)
solved' <- get
put $ Map.insert x (prev1 + prev2) solved'

Just get the updated state after the recursive invocations.

Naming

I think solved is ambiguous as to whether it refers to a single solution or the memo map. I would call it something like solutions, maybe sols or fibs if you feel like abbreviating. In light of the previous bug, you may wish to call the first one initialSolutions and later ones updatedSolutions or something like that (I just used solutions and solutions').

The name m is OK, but it doesn't really need to exist at all. You can just case on Map.lookup x solved directly instead of binding m and then casing on it.

The case Statement

It seems like you're trying to use the case statement to set up your memo map so that it always has a value for x. There are two things to address about this.

Indexing

If you, the programmer, are sure that x exists in the map (and in the case above you guarantee it), you could instead use the unsafe lookup Map.!. Given the option between returning an erroneous value silently (your fromMaybe 0) and crashing and burning in case of a bug, I would generally prefer the latter. So your last line would look something like gets (Map.! x).

The Last Lookup

However, doing the lookup itself is inelegant. It might make sense if there was a lot of convoluted stuff happening between, but proper indexing doesn't get checked by the type system and doing a lookup takes (not much, but some) extra time. Fortunately, you don't need to do it. Since I'm going to assume you're learning Haskell, consider how you'd approach a similar problem in an imperative language. What would you do to change this code:

if (x in solutions):
  solutions[x] = solutions[x]
else:
  prev1 = fib(x-1)
  prev2 = fib(x-2)
  solutions[x] = prev1 + prev2
return solutions[x]

There are many right answers, but one thing you can do is as follows (this particular code is nice because it avoids extra lookups):

if (x in solutions):
  return solutions[x]
else:
  prev1 = fib(x-1)
  prev2 = fib(x-2)
  solution = prev1 + prev2
  solutions[x] = solution
  return solution

Your case statement functions like the imperative if statement, except more powerful since you have guarantees on the types! So mirroring the imperative's revision, you can revise your code like so

case Map.lookup x solved of
        Just solution -> return solution
        Nothing -> do
            prev1 <- fib (x - 1)
            prev2 <- fib (x - 2)
            solutions' <- get
            let solution = prev1 + prev2
            put $ Map.insert x solutions' solved
            return solution

Now you don't need the last lookup. Notice how we also avoid the issue entirely of whether x is in solutions, because we explicitly handle the case where it is and isn't. This code doesn't have any unsafe lookups!

Addendum on Lookups

Now, even if you wanted to make your case statement only fill out the memo map instead of also returning the answers, I agree with you that you are doing unnecessary work.

Just value -> put $ Map.insert x value solved

The line above needlessly reinserts the value of x. The memo map already has x, and x is already set to value. If you wanted to otherwise keep your code the same, at least change this to

Just value -> return ()

put has type a -> State a (). It's a convention for monads to pass () as their return value if they perform an action that doesn't return anything (like how putStrLn has type String -> IO ()). You can simply return () to do nothing instead of actually modifying the memo map, which I assume you did to fix a type error.

Revised Function

Included are comments noting the revision

fib :: Int -> State DP Int
fib 0 = return 1
fib 1 = return 1
fib x = do
    -- Change to a more descriptive name
    solutions <- get
    -- Case directly on the value without intermediate variable
    case Map.lookup x solutions of
            Nothing -> do
                prev1 <- fib (x - 1)
                prev2 <- fib (x - 2)
                -- Get updated solutions
                solutions' <- get
                let solution = prev1 + prev2
                put $ Map.insert x solution solutions'
                -- Return solution directly
                return solution
            Just solution ->
                -- Return solution directly
                return solution
    -- Elide previous lookup