awkward monad transformer stack

Question

Solving a problem from Google Code Jam (2009.1A.A: "Multi-base happiness") I came up with an awkward (code-wise) solution, and I'm interested in how it could be improved.

The problem description, shortly, is: Find the smallest number bigger than 1 for which iteratively calculating the sum of squares of digits reaches 1, for all bases from a given list.

Or description in pseudo-Haskell (code that would solve it if elem could always work for infinite lists):

solution =
  head . (`filter` [2..]) .
  all ((1 `elem`) . (`iterate` i) . sumSquareOfDigitsInBase)

And my awkward solution:

• By awkward I mean it has this kind of code: happy <- lift . lift . lift $ isHappy Set.empty base cur
• I memoize results of the isHappy function. Using the State monad for the memoized results Map.
• Trying to find the first solution, I did not use head and filter (like the pseudo-haskell above does), because the computation isn't pure (changes state). So I iterated by using StateT with a counter, and a MaybeT to terminate the computation when condition holds.
• Already inside a MaybeT (StateT a (State b)), if the condition doesn't hold for one base, there is no need to check the other ones, so I have another MaybeT in the stack for that.

Code:

import Control.Monad.Maybe
import Control.Monad.State
import Data.Maybe
import qualified Data.Map as Map
import qualified Data.Set as Set

type IsHappyMemo = State (Map.Map (Integer, Integer) Bool)

isHappy :: Set.Set Integer -> Integer -> Integer -> IsHappyMemo Bool
isHappy _ _ 1 = return True
isHappy path base num = do
  memo <- get
  case Map.lookup (base, num) memo of
    Just r -> return r
    Nothing -> do
      r <- calc
      when (num < 1000) . modify $ Map.insert (base, num) r
      return r
  where
    calc
      | num `Set.member` path = return False
      | otherwise = isHappy (Set.insert num path) base nxt
    nxt =
      sum . map ((^ (2::Int)) . (`mod` base)) .
      takeWhile (not . (== 0)) . iterate (`div` base) $ num

solve1 :: [Integer] -> IsHappyMemo Integer
solve1 bases =
  fmap snd .
  (`runStateT` 2) .
  runMaybeT .
  forever $ do
    (`when` mzero) . isJust =<<
      runMaybeT (mapM_ f bases)
    lift $ modify (+ 1)
  where
    f base = do
      cur <- lift . lift $ get
      happy <- lift . lift . lift $ isHappy Set.empty base cur
      unless happy mzero

solve :: [String] -> String
solve =
  concat .
  (`evalState` Map.empty) .
  mapM f .
  zip [1 :: Integer ..]
  where
    f (idx, prob) = do
      s <- solve1 . map read . words $ prob
      return $ "Case #" ++ show idx ++ ": " ++ show s ++ "\n"

main :: IO ()
main =
  getContents >>=
  putStr . solve . tail . lines

Other contestants using Haskell did have nicer solutions, but solved the problem differently. My question is about small iterative improvements to my code.

Answer 1

Your solution is certainly awkward in its use (and abuse) of monads:

• It is usual to build monads piecemeal by stacking several transformers
• It is less usual, but still happens sometimes, to stack several states
• It is very unusual to stack several Maybe transformers
• It is even more unusual to use MaybeT to interrupt a loop

Your code is a bit too pointless :

(`when` mzero) . isJust =<<
   runMaybeT (mapM_ f bases)

instead of the easier to read

let isHappy = isJust $ runMaybeT (mapM_ f bases)
when isHappy mzero

Focusing now on function solve1, let us simplify it. An easy way to do so is to remove the inner MaybeT monad. Instead of a forever loop which breaks when a happy number is found, you can go the other way around and recurse only if the number is not happy.

Moreover, you don't really need the State monad either, do you ? One can always replace the state with an explicit argument.

Applying these ideas solve1 now looks much better:

solve1 :: [Integer] -> IsHappyMemo Integer
solve1 bases = go 2 where
  go i = do happyBases <- mapM (\b -> isHappy Set.empty b i) bases
            if and happyBases
              then return i
              else go (i+1)

I would be more han happy with that code. The rest of your solution is fine. One thing that bothers me is that you throw away the memo cache for every subproblem. Is there a reason for that?

solve :: [String] -> String
 solve =
    concat .
    (`evalState` Map.empty) .
    mapM f .
   zip [1 :: Integer ..]
  where
    f (idx, prob) = do
      s <- solve1 . map read . words $ prob
      return $ "Case #" ++ show idx ++ ": " ++ show s ++ "\n"

Wouldn't your solution be more efficient if you reused it instead ?

solve :: [String] -> String
solve cases = (`evalState` Map.empty) $ do
   solutions <- mapM f (zip [1 :: Integer ..] cases)
   return (unlines solutions)
  where
    f (idx, prob) = do
      s <- solve1 . map read . words $ prob
      return $ "Case #" ++ show idx ++ ": " ++ show s

Answer 2

The Monad* classes exist to remove the need for repeated lifting. If you change your signatures like this:

type IsHappyMemo = Map.Map (Integer, Integer) Bool

isHappy :: MonadState IsHappyMemo m => Set.Set Integer -> Integer -> Integer -> m Bool

This way you can remove most of the 'lift's. However, the longest sequence of lifts cannot be removed, since it is a State monad inside a StateT, so using the MonadState type class will give you the outer StateT, where you need tot get to the inner State. You could wrap your State monad in a newtype and make a MonadHappy class, similar to the existing monad classes.