How to implement a stack-safe chainRec operator for the continuation monad?

Question

I am currently experimenting with the continuation monad. Cont is actually useful in Javascript, because it abstracts from the callback pattern.

When we deal with monadic recursion, there is always the risk of a stack overflow, because the recursive call isn't in tail position:

const chain = g => f => k =>
  g(x => f(x) (k));

const of = x => k =>
  k(x);
  
const id = x =>
  x;

const inc = x =>
  x + 1;

const repeat = n => f => x => 
  n === 0
    ? of(x)
    : chain(of(f(x))) (repeat(n - 1) (f));

console.log(
  repeat(1e6) (inc) (0) (id) // stack overflow
);

However, even if we are able to transform some cases into tail recursion we are still doomed, because Javascript doesn't have TCO. Consequently we have to fall back to a loop at some point.

puresrcipt has a MonadRec typeclass with a tailRecM operator that enables tail recursive monadic computations for some monads. So I tried to implement chainRec in Javascript mainly according to the fantasy land spec:

const chain = g => f => k => g(x => f(x) (k));
const of = x => k => k(x);
const id = x => x;

const Loop = x =>
  ({value: x, done: false});

const Done = x =>
  ({value: x, done: true});

const chainRec = f => x => {
  let step = f(Loop, Done, x);

  while (!step.done) {
    step = f(Loop, Done, step.value);
  }

  return of(step.value);
};

const repeat_ = n => f => x => 
  chainRec((Loop, Done, [n, x]) => n === 0 ? Done(x) : Loop([n - 1, f(x)])) ([n, x]);

console.log(
  repeat_(1e6) (n => n + 1) (0) (id) // 1000000
);

This works, but it looks a lot like cheating, because it seems to bypass the monadic chaining and thus Cont's context. In this case the context is just "the rest of the computation", ie. function composition in reverse and as a result the expected value is returned. But does it work for any monad?

To make it clear what I mean take a look at the following code snippet from this outstanding answer:

const Bounce = (f,x) => ({ isBounce: true, f, x })

const Cont = f => ({
  _runCont: f,
  chain: g =>
    Cont(k =>
      Bounce(f, x =>
        Bounce(g(x)._runCont, k)))
})

// ...

const repeat = n => f => x => {
  const aux = (n,x) =>
    n === 0 ? Cont.of(x) : Cont.of(f(x)).chain(x => aux(n - 1, x))
  return runCont(aux(n,x), x => x)
}

Here chain is somehow incorporated into the recursive algorithm, that is the monadic effect can occur. Unfortunately, I cannot decipher this operator or reconcile it with the stack-unsafe version (Bounce(g(x)._runCont, k) seems to be the f(x) (k) portion, though).

Ultimately, my question is if I messed up the implementation of chainRecor misunderstood the FL spec or both or none of it?

---

[EDIT]

Both given answers are very helpful by looking at the problem from different perspectives and deserve to be accepted. Since I can only accept one - hey stackoverflow, the world isn't that simple!!! - I won't accept any.

Answer 1

with best wishes,

I think this might be what you're looking for,

const chainRec = f => x =>
  f ( chainRec (f)
    , of
    , x
    )

Implementing repeat is just as you have it – with two exceptions (thanks @Bergi for catching this detail). 1, loop and done are the chaining functions, and so the chainRec callback must return a continuation. And 2, we must tag a function with run so cont knows when we can safely collapse the stack of pending continuations – changes in bold

const repeat_ = n => f => x =>
  chainRec
    ((loop, done, [n, x]) =>
       n === 0
         ? of (x) (done)                // cont chain done
         : of ([ n - 1, f (x) ]) (loop) // cont chain loop
    ([ n, x ])

const repeat = n => f => x =>
  repeat_ (n) (f) (x) (run (identity))

But, if you're using chainRec as we have here, of course there's no reason to define the intermediate repeat_. We can define repeat directly

const repeat = n => f => x =>
  chainRec
    ((loop, done, [n, x]) =>
       n === 0
         ? of (x) (done)
         : of ([ n - 1, f (x) ]) (loop)
    ([ n, x ])
    (run (identity))

Now for it to work, you just need a stack-safe continuation monad – cont (f) constructs a continuation, waiting for action g. If g is tagged with run, then it's time to bounce on the trampoline. Otherwise constructor a new continuation that adds a sequential call for f and g

// not actually stack-safe; we fix this below
const cont = f => g =>
  is (run, g)
    ? trampoline (f (g))
    : cont (k =>
        call (f, x =>
          call (g (x), k)))

const of = x =>
  cont (k => k (x))

Before we go further, we'll verify things are working

const TAG =
  Symbol ()

const tag = (t, x) =>
  Object.assign (x, { [TAG]: t })
  
const is = (t, x) =>
  x && x [TAG] === t

// ----------------------------------------

const cont = f => g =>
  is (run, g)
    ? trampoline (f (g))
    : cont (k =>
        call (f, x =>
          call (g (x), k)))
  
const of = x =>
  cont (k => k (x))

const chainRec = f => x =>
  f ( chainRec (f)
    , of
    , x
    )
  
const run = x =>
  tag (run, x)
  
const call = (f, x) =>
  tag (call, { f, x })  

const trampoline = t =>
{
  let acc = t
  while (is (call, acc))
    acc = acc.f (acc.x)
  return acc
}

// ----------------------------------------

const identity = x =>
  x
  
const inc = x =>
  x + 1

const repeat = n => f => x =>
  chainRec
    ((loop, done, [n, x]) =>
       n === 0
         ? of (x) (done)
         : of ([ n - 1, f (x) ]) (loop))
    ([ n, x ])
    (run (identity))
      
console.log (repeat (1e3) (inc) (0))
// 1000

console.log (repeat (1e6) (inc) (0))
// Error: Uncaught RangeError: Maximum call stack size exceeded

where's the bug?

The two implementations provided contain a critical difference. Specifically, it's the g(x)._runCont bit that flattens the structure. This task is trivial using the JS Object encoding of Cont as we can flatten by simply reading the ._runCont property of g(x)

const Cont = f =>
  ({ _runCont: f
   , chain: g =>
       Cont (k =>
         Bounce (f, x =>
          // g(x) returns a Cont, flatten it
          Bounce (g(x)._runCont, k))) 
   })

In our new encoding, we're using a function to represent cont, and unless we provide another special signal (like we did with run), there's no way to access f outside of cont once it's been partially applied – look at g (x) below

const cont = f => g =>
  is (run, g)
    ? trampoline (f (g))
    : cont (k =>
        call (f, x =>
          // g (x) returns partially-applied `cont`, how to flatten?
          call (g (x), k))) 

Above, g (x) will return a partially-applied cont, (ie cont (something)), but this means that the entire cont function can nest infinitely. Instead of cont-wrapped something, we only want something.

At least 50% of the time I spent on this answer has been coming up with various ways to flatten partially-applied cont. This solution isn't particularly graceful, but it does get the job done and highlights precisely what needs to happen. I'm really curious to see what other encodings you might find – changes in bold

const FLATTEN =
  Symbol ()

const cont = f => g =>
  g === FLATTEN
    ? f
    : is (run, g)
      ? trampoline (f (g))
      : cont (k =>
          call (f, x =>
            call (g (x) (FLATTEN), k)))

all systems online, captain

With the cont flattening patch in place, everything else works. Now see chainRec do a million iterations…

const TAG =
  Symbol ()

const tag = (t, x) =>
  Object.assign (x, { [TAG]: t })
  
const is = (t, x) =>
  x && x [TAG] === t

// ----------------------------------------

const FLATTEN =
  Symbol ()

const cont = f => g =>
  g === FLATTEN
    ? f
    : is (run, g)
      ? trampoline (f (g))
      : cont (k =>
          call (f, x =>
            call (g (x) (FLATTEN), k)))
  
const of = x =>
  cont (k => k (x))

const chainRec = f => x =>
  f ( chainRec (f)
    , of
    , x
    )
  
const run = x =>
  tag (run, x)
  
const call = (f, x) =>
  tag (call, { f, x })  

const trampoline = t =>
{
  let acc = t
  while (is (call, acc))
    acc = acc.f (acc.x)
  return acc
}

// ----------------------------------------

const identity = x =>
  x
  
const inc = x =>
  x + 1

const repeat = n => f => x =>
  chainRec
    ((loop, done, [n, x]) =>
       n === 0
         ? of (x) (done)
         : of ([ n - 1, f (x) ]) (loop))
    ([ n, x ])
    (run (identity))
      
console.log (repeat (1e6) (inc) (0))
// 1000000

evolution of cont

When we introduced cont in the code above, it's not immediately obvious how such an encoding was derived. I hope to shed some light on that. We start with how we wish we could define cont

const cont = f => g =>
  cont (comp (g,f))

const comp = (f, g) =>
  x => f (g (x))

In this form, cont will endlessly defer evaluation. The only available thing we can do is apply g which always creates another cont and defers our action. We add an escape hatch, run, which signals to cont that we don't want to defer any longer.

const cont = f => g =>
  is (run, g)
    ? f (g)
    : cont (comp (g,f))

const is = ...

const run = ...

const square = x =>
  of (x * x)

of (4) (square) (square) (run (console.log))
// 256

square (4) (square) (run (console.log))
// 256

Above, we can begin to see how cont can express beautiful and pure programs. However in an environment without tail-call elimination, this still allows programs to build deferred functions sequences that exceed the evaluator's stack limit. comp directly chains functions, so that's out of the picture. Instead we'll sequence the functions using a call mechanism of our own making. When the program signals run, we collapse the stack of calls using trampoline.

Below, we arrive at the form we had before the flatten fix was applied

const cont = f => g =>
  is (run, g)
    ? trampoline (f (g))
    : cont (comp (g,f))
    : cont (k =>
        call (f, x =>
          call (g (x), k)))

const trampoline = ...

const call = ...

wishful thinking

Another technique we were using above is one of my favorites. When I write is (run, g), I don't know how I'm going to represent is or run right away, but I can figure it out later. I use the same wishful thinking for trampoline and call.

I point this out because it means I can keep all of that complexity out of cont and just focus on its elementary structure. I ended up with a set of functions that gave me this "tagging" behavior

// tag contract
// is (t, tag (t, value)) == true

const TAG =
  Symbol ()

const tag = (t, x) =>
  Object.assign (x, { [TAG]: t })

const is = (t, x) =>
  x && x [TAG] === t

const run = x =>
  tag (run, x)

const call = (f, x) =>
  tag (call, { f, x })

Wishful thinking is all about writing the program you want and making your wishes come true. Once you fulfill all of your wishes, your program just magically works!

Answer 2

Did I mess up the implementation of chainRec, or misunderstood the FantasyLand spec, or both or none of it?

Probably both, or at least the first. Notice that the type should be

chainRec :: ChainRec m => ((a -> c, b -> c, a) -> m c, a) -> m b

wherein m is Cont and c is your Done/Loop wrapper over a or b:

chainRec :: ((a -> DL a b, b -> DL a b, a) -> Cont (DL a b), a) -> Cont b

But your chainRec and repeat implementations don't use continations at all!

If we implement just that type, without the requirement that it should need constant stack space, it would look like

const chainRec = f => x => k =>
  f(Loop, Done, x)(step =>
    step.done
      ? k(step.value) // of(step.value)(k)
      : chainRec(f)(step.value)(k)
  );

or if we drop even the lazyness requirement (similar to transforming chain from g => f => k => g(x => f(x)(k)) to just g => f => g(f) (i.e. g => f => k => g(x => f(x))(k))), it would look like

const chainRec = f => x =>
  f(Loop, Done, x)(step =>
    step.done
      ? of(step.value)
      : chainRec(f)(step.value)
  );

or even dropping Done/Loop

const join = chain(id);
const chainRec = f => x => join(f(chainRec(f), of, x));

(I hope I'm not going out on a limb too far with that, but it perfectly presents the idea behind ChainRec)

With the lazy continuation and the non-recursive trampoline, we would however write

const chainRec = f => x => k => {
  let step = Loop(x);
  do {
    step = f(Loop, Done, step.value)(id);
//                                  ^^^^ unwrap Cont
  } while (!step.done)
  return k(step.value); // of(step.value)(k)
};

The loop syntax (initialise step with an f call, do/while instead of do) doesn't really matter, yours is fine as well but the important part is that f(Loop, Done, v) returns a continuation.

I'll leave the implementation of repeat as an exercise to the reader :D
(Hint: it might become more useful and also easier to get right if you have the repeated function f already use continuations)