Collatz Memoization in Haskell

gistfile1.hs

{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE TypeOperators #-}

import Data.Function ( fix )
import Data.MemoTrie        -- from package MemoTrie
import Data.MemoCombinators -- from package data-memocombinators

main :: IO ()
main =
  -- without caching
  -- 4 seconds
  -- test $ fix collatz

  -- with caching of each call individually using MemoTrie
  -- 30 seconds
  -- test $ fixmemotrie collatz

  -- with global caching using MemoTrie
  -- 31 seconds
  -- test $ untrie collatzTrie

  -- with global caching using memocombinators
  -- 3 seconds
  test $ fixmemocomb (arrayRange (1,10^6)) collatz

test :: (Int -> [Int]) -> IO ()
test collatz = print $ maxLengthBetween collatz 1 (10^6)

maxLengthBetween :: (Int -> [Int]) -> Int -> Int -> Int
maxLengthBetween collatz i j = maximum $ map (length . collatz) [i..j]

collatz :: (Int -> [Int]) -> (Int -> [Int])
collatz collatz 1             = [1]
collatz collatz n | even n    = n : collatz (n`div`2)
                  | otherwise = n : collatz (3*n + 1)

fixmemotrie :: HasTrie a => ((a -> b) -> (a -> b)) -> (a -> b)
fixmemotrie = fixmemocomb memo

fixmemocomb :: Memo a -> ((a -> b) -> (a -> b)) -> (a -> b)
fixmemocomb memo f = fix (memo . f)

collatzTrie :: Int :->: [Int]
collatzTrie = trie (collatz (untrie collatzTrie))