Memo.hs

{-# LANGUAGE ViewPatterns #-}
module Memo where

import Data.Bits (testBit, setBit, finiteBitSize)

data Memo a b = Fork (Memo a b) b (Memo a b)
    deriving Show

type Bit = Bool
newtype Bits = Bits [Bit]

instance Show Bits where
    show (Bits b) = "[" ++ map (\x -> if x then '1' else '0') b ++ "]"

{-
Laws:
toBits . fromBits = fromBits . toBits = id
∀n
    zeros := replicate n False
    ∀x fromBits x = fromBits (x ++ zeros)
-}
class BitsBijective a where
    toBits :: a -> Bits
    fromBits :: Bits -> a

instance BitsBijective Bits where
    toBits = id
    fromBits = id

instance BitsBijective Int where
    toBits n = Bits $ map (testBit n) [0..finiteBitSize n - 1]
    fromBits (Bits b) = foldl setBit 0 $ map fst $ filter snd $ zip [0..] b

interleave :: Bits -> Bits -> Bits
interleave (Bits l) (Bits r) = Bits $ interleave' l r
    where interleave' []       [] = []
          interleave' []       ys = False : interleave' ys []
          interleave' (x : xs) ys = x : interleave' ys xs

uninterleave :: Bits -> (Bits, Bits)
uninterleave (Bits b) = (Bits l, Bits r)
    where (l, r) = uninterleave' b
          uninterleave' []       = ([], [])
          uninterleave' (x : xs) = let (rs, ls) = uninterleave' xs in (x : ls, rs)

instance (BitsBijective a, BitsBijective b) => BitsBijective (a, b) where
    toBits (toBits -> l, toBits -> r) = interleave l r
    fromBits (uninterleave -> (l, r)) = (fromBits l, fromBits r)

memo :: BitsBijective a => (a -> b) -> a -> b
memo f = readMemo
    where memo' l = Fork (memo' $ False : l) (f $ fromBits $ Bits $ reverse l) (memo' $ True : l)
          tree = memo' []
          readMemo x = followPath b tree
              where (Bits (ker -> b)) = toBits x
                    followPath []           (Fork _ y _) = y
                    followPath (False : xs) (Fork y _ _) = followPath xs y
                    followPath (True : xs)  (Fork _ _ y) = followPath xs y
          ker xs | null ones = []
                 | otherwise = take (1 + fst (last ones)) xs
                 where ones = filter snd $ zip [0..] xs

fib :: Int -> Int
fib = memo fib'
    where fib' 0 = 1    
          fib' 1 = 1
          fib' n = fib (n - 1) + fib (n - 2)