characterization of subset of Monads that respect tail calls?

[Tom Ellis]

On Tue, Sep 27, 2022 at 05:24:43PM -0400, A S wrote:
> The definition tailRecM1, while it typechecks, will not always
> operate in constant stack space

Correct.

> As for tailRecM2, unfortunately this precludes many useful and nontrivial
> MonadRec instances such as Effect and Aff, which admit no such distributing
> natural transformation (if this is what you actually meant for the first
> argument to be).

Correct, the distribution property is too strong to capture all
MonadRec instances.  However, it does capture Maybe, Either e and
Identity.  To capture (r ->)/Reader we need something more general
(tailRecM3) which is sadly significantly less pleasant, but unless I
am much mistaken it captures all instances of tailRecM that use
tailRec.  Maybe it's too unpleasant to be used in practice as an
alternative of a direct implementation of tailRecM.

Tom


{-# LANGUAGE LambdaCase #-}

module Rec where

import Data.Functor.Identity (Identity(Identity))

data Step a b = Loop a | Done b

onDone :: (b -> b') -> Step a b -> Step a b'
onDone f = \case
  Loop a -> Loop a
  Done b -> Done (f b)

tailRec :: (a -> Step a b) -> a -> b
tailRec f = go . f
  where
  go (Loop a) = go (f a)
  go (Done b) = b

-- | This works for all monads but is not necessarily stack safe.
tailRecM1 :: Monad m => (a -> m (Step a b)) -> a -> m b
tailRecM1 f a = do
  f a >>= \case
    Loop a' -> tailRecM1 f a'
    Done b -> pure b

-- | This works for all monads with a particular distribution property
-- and is stack safe through the use of @tailRec at .
tailRecM3 ::
  Monad m =>
  ((a -> m (Step a b)) -> m (a -> Step a l)) ->
  (l -> m b) ->
  (a -> m (Step a b)) ->
  a ->
  m b
tailRecM3 distribute finish f a = do
  g <- distribute f
  finish (tailRec g a)

tailRecReader :: (a -> r -> Step a b) -> a -> r -> b
tailRecReader = tailRecM3 (\f r a -> f a r) (\b _ -> b)

tailRecDistribute ::
  Monad m =>
  (m (Step a b) -> Step a (m b)) ->
  (a -> m (Step a b)) ->
  a ->
  m b
tailRecDistribute distribute = tailRecM3 (\f -> pure (\a -> distribute
  (f a))) id

tailRecMaybe :: (a -> Maybe (Step a b)) -> a -> Maybe b
tailRecMaybe = tailRecDistribute $ \case
  Nothing -> Done Nothing
  Just j -> onDone pure j

tailRecEither :: (a -> Either e (Step a b)) -> a -> Either e b
tailRecEither = tailRecDistribute $ \case
  Left e -> Done (Left e)
  Right r -> onDone pure r

tailRecIdentity :: (a -> Identity (Step a b)) -> a -> Identity b
tailRecIdentity = tailRecDistribute $ \case
  Identity i -> onDone pure i