tailRecM

tailRecM.hs

{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE InstanceSigs #-}
--{-# LANGUAGE ScopedTypeVariables #-}
module TailRecM where

-- This only makes sense in Purescript (strict haskell dialect)    
-- But I am lacking a purescript runtime interpreter

import Control.Monad.Writer.Lazy
import Control.Monad.Trans.State.Lazy
import Data.Monoid (Product)
import Data.Functor.Identity

pow :: Int -> Int -> Int
pow n = go 1 
  where
    go acc 0 = acc
    go acc p = go (acc * n) (p - 1)

-- Will blow the stack
powWriter :: Int -> Int -> Writer (Product Int) ()
powWriter n = go
  where
    go :: Int -> Writer (Product Int) ()
    go 0 = pure ()
    go m = do
      tell (Product n)
      go (m - 1)

-- Generic tail-recursive function:
tailRec :: forall a b. (a -> Either a b) -> a -> b
tailRec f a = go (f a)
  where
    go (Left a)  = go (f a)
    go (Right b) = b
  
-- Pow stack-safe
powTR :: Int -> Int -> Int
powTR n p = tailRec go (1, p)
  where 
    go (acc, 0) = Right acc
    go (acc, m) = Left (acc * n, m - 1)

-- TailRec can be generalized to several monads
-- tailRec is a valid implementation for Id
class (Monad m) => MonadRec m where
  -- A valid implementation of MonadRec must guarantee that the stack usage of tailRecM 
  -- is at most a constant multiple of the stack usage off itself.
  tailRecM :: forall a b. (a -> m (Either a b)) -> a -> m b
  -- Not a valid instance. Doesn't fulfill the law
  tailRecM f a = f a >>= go 
    where
      go (Left a)  = f a >>= go 
      go (Right b) = return b


forever :: forall m a b. (MonadRec m) => m a -> m b
forever = undefined

-- Example of instances

instance (MonadRec m) => MonadRec (StateT s m) where
  tailRecM f a = StateT (\s -> tailRecM f' (a, s))
    where
      f' (a, s) = do
         (m, s1) <- runStateT (f a) s
         return $ case m of 
          Left a  -> Left (a, s1) 
          Right b -> Right (b, s1)

instance (Monoid w, MonadRec m) => MonadRec (WriterT w m) where
  tailRecM f a = WriterT $ tailRecM f' (a, mempty)
     where
      f' (a, w) = do
        (m, w1) <- runWriterT (f a)
        return $ case m of
          Left a -> Left (a, w1)
          Right b -> Right (b, w1)

instance MonadRec Identity where
  tailRecM f a = Identity $ go (f a)
    where
      go (Identity m) = case m of
        Left a -> go (f a)
        Right b -> b

powWriterTR :: Int -> Int -> Writer (Product Int) ()
powWriterTR n = tailRecM go
  where
    go :: Int -> Writer (Product Int) (Either Int ())
    go 0 = return (Right ())
    go p = do
      tell $ Product n
      return (Left $ p - 1)