Control.Monad.Fix

Description

Monadic fixpoints, used for desugaring of {-# LANGUAGE RecursiveDo #-}.

Consider the generalized version of so-called repmin (replace with minimum) problem: accumulate elements of a container into a Monoid and modify each element using the final accumulator.

repmin
  :: (Functor t, Foldable t, Monoid b)
  => (a -> b) -> (a -> b -> c) -> t a -> t c
repmin f g as = fmap (`g` foldMap f as) as

The naive implementation as above makes two traversals. Can we do better and achieve the goal in a single pass? It's seemingly impossible, because we would have to know the future, but lazy evaluation comes to the rescue:

import Data.Traversable (mapAccumR)

repmin
  :: (Traversable t, Monoid b)
  => (a -> b) -> (a -> b -> c) -> t a -> t c
repmin f g as =
  let (b, cs) = mapAccumR (\acc a -> (f a <> acc, g a b)) mempty as in cs

How can we check that repmin indeed traverses only once? Let's run it on an infinite input:

>>> import Data.Monoid (All(..))
>>> take 3 $ repmin All (const id) ([True, True, False] ++ undefined)
[All {getAll = False},All {getAll = False},All {getAll = False}]

So far so good, but can we generalise g to return a monadic value a -> b -> m c? The following does not work, complaining that b is not in scope:

import Data.Traversable (mapAccumM)

repminM
  :: (Traversable t, Monoid b, Monad m)
  => (a -> b) -> (a -> b -> m c) -> t a -> m (t c)
repminM f g as = do
  (b, cs) <- mapAccumM (\acc a -> (f a <> acc,) <$> g a b) mempty as
  pure cs

To solve the riddle, let's rewrite repmin via fix:

repmin
  :: (Traversable t, Monoid b)
  => (a -> b) -> (a -> b -> c) -> t a -> t c
repmin f g as = snd $ fix $
  \(b, cs) -> mapAccumR (\acc a -> (f a <> acc, g a b)) mempty as

Now we can replace fix with mfix to obtain the solution:

repminM
  :: (Traversable t, Monoid b, MonadFix m)
  => (a -> b) -> (a -> b -> m c) -> t a -> m (t c)
repminM f g as = fmap snd $ mfix $
  \(~(b, cs)) -> mapAccumM (\acc a -> (f a <> acc,) <$> g a b) mempty as

For example,

>>> import Data.Monoid (Sum(..))
>>> repminM Sum (\a b -> print a >> pure (a + getSum b)) [3, 5, 2]
3
5
2
[13,15,12]

Incredibly, GHC is capable to do this transformation automatically, when {-# LANGUAGE RecursiveDo #-} is enabled. Namely, the following implementation of repminM works (note mdo instead of do):

{-# LANGUAGE RecursiveDo #-}

repminM
  :: (Traversable t, Monoid b, MonadFix m)
  => (a -> b) -> (a -> b -> m c) -> t a -> m (t c)
repminM f g as = mdo
  (b, cs) <- mapAccumM (\acc a -> (f a <> acc,) <$> g a b) mempty as
  pure cs

Further reading:

• GHC User’s Guide, The recursive do-notation.
• Haskell Wiki, MonadFix.
• Levent Erkök, Value recursion in monadic computations, Oregon Graduate Institute, 2002.
• Levent Erkök, John Launchbury, A recursive do for Haskell, Haskell '02, 29-37, 2002.
• Richard S. Bird, Using circular programs to eliminate multiple traversals of data, Acta Informatica 21, 239-250, 1984.
• Jasper Van der Jeugt, Lazy layout, 2023.

>>> let fac n = if n <= 1 then 1 else n * fac (n-1) in fac 5
120

>>> fix (\rec n -> if n <= 1 then 1 else n * rec (n-1)) 5
120

>>> take 10 $ fix (0:)
[0,0,0,0,0,0,0,0,0,0]

>>> map (fix (\rec n -> if n < 2 then n else rec (n - 1) + rec (n - 2))) [1..10]
[1,1,2,3,5,8,13,21,34,55]

fix f = let x = f x in x

fix f = f (fix f)