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.

Insulator

Javascript doesn't have TCO i thought it does? – Chymotrypsin

@JonasW. Vendors are reluctant to implement it and it's not clear when and if it will happen. – Insulator

afaik TCO only applies to recursive tail calls, wich you don't have in your code snippet. Anyways, that's why I don't like to dive that deep into FP. It's too abstract. It'l take me some time to wrap my head around Cont and to apply the differences in the linked answer to your repeat function. – Librettist

@Librettist oops, right, that was a mistake. I edited my question, thanks. – Insulator

@Thomas, afaik tail call optimization is about safely reusing stack frames but recursive functions aren't the only benefactors. If TCO only on direct, by-name recursive calls but failed for anonymous recursion and indirect/co-recursion cases, it'd be useless – I never got around to finishing this answer but my goal was to show that implementing TCO is simply about managing a stack of calls, and it could be used for any type of function; not just recursive ones. – Pill

@ftor, in that linked answer, there's a section labeled Continuation passing style with trampoline. I think this one might be interesting to you. If I understand your question correctly, the "trick" of Cont working is the explicit use of runCont which calls trampoline – Pill

Answer 1

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)

Hoag

I implemented repeat and applied it to a sync and async f both returning continuations. The former works fine but the latter does not. Is this a consequence of the flattening step, ie. that Cont loses its lazyness in conjunction with chainRec? – Insulator

Cont never was meant to work asynchronously, we always expected it to return the result of the last continuation. The problem is that you cannot unwrap an asynchronous continuation, we cannot get the result by calling it with id and use that in a loop. For asynchronous continuations, chainRec = f => x => f(chainRec(f), of, x); should work, as asynchronous calls don't overflow the stack anyway. – Hoag

Mixing async/sync functions is the reason for me to examine the Cont type in the first place. But right, it was not part of this question. – Insulator

Well I guess you could write a Cont type that dynamically checks whether the callback is invoked synchronously or not, but I'm not sure it would be efficient – Hoag

Answer 2

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!

Pill

That's a lot of food for thought. Thank you for explaining your approach in detail so that I have a chance to understand it. I'll come back to your answer... – Insulator

The reason why I have difficulties decoding such functional encodings is a lack of understanding of functional (data) structures. When I work with []/{} I immediately have a mental model of the tree structures they are forming and it is easy to apply traverse/map/flatten operations to them. Partially applied functions now build up tree-like structures too, but I am not able to create a corresponding mental model. I guess these computational structures are more complex than nested []/{}, because the function type is more complex. This is the missing building block I have to work on. – Insulator

@naomik I think you made the same mistake as ftor in the implementation of repeat. The chainRec callback actually needs to return a continuation, not just a loop or done value. (It works though because you simplified chainRec right in the beginning so that loop and done actually do create continuations) – Hoag

@Hoag thanks for catching that. The chainRec type signature is quite an eye bender, but it makes perfect sense now. I love that we landed on the same implementation in end, but I'm sorry to have ruined your reader exercise! No way to mark spoilers in code snippets yet :( – Pill

@naomik But of(x)(done) is still the same as done(x)? It would need to be of(x).map(done) or of(done(x)). – Hoag

@Hoag I guess I've only managed to confuse myself further. If loop and done are mapping functions, doesn't that completely invalidate chainRec = f => x => f(chainRec(f), of, x)? – Pill

@naomik Right. Dammit, I need to look further into that. I guess it needs a join/flatten step. – Hoag

@Hoag the chainRec type signature is starting to feel a little unnatural - it seems like it’s describing more of a mapRec process... – Pill

@ftor, arrays and objects can be nested infinitely because references from parent to child are simply memory addresses. Accessing a nested value or property is a matter of traversing the memory addresses. When the lambda is used, accessing a deep value means utilizing a stack frame for each level of nesting. To achieve infinite "nesting" with lambda, we use call to structure a sequence of deferred lambdas, each existing in memory with a result pointing to the next call, waiting to be applied by the trampoline – note call uses an infinitely nestable type ({}) to make this possible – Pill

@naomik I am mostly aware of this. What is giving me a hard time is that with function types we seem to have nesting (e.g. f(g(h(...)))), we have sequencing (e.g. x => y => z => ...) and we have combinations of both k => f(...) (x => g(...) (k)). I think it is harder to reason about these functional structures than about ordinary data structures. – Insulator

@naomik On second thought we have both properties with data structures too. So probably there isn't really a difference on the semantic level, only at the technical, as you described. That's exactly what I want to take a closer look at. – Insulator

@ftor functional structures are so cool!. Thanks again for a fun exercise, even though I still mostly failed the entire thing lol. – Pill

@naomik Yeah, partly because they offer real pattern matching without anything having to implement Eq. It's all about continuation passing :D – Insulator

I think I am beginning to understand. Since the call stack in Javascript is insufficient for deeply nested computations you replace it with your own nested data structure. If everything is in tail position you use a trampoline to eliminate nestings. Essentially, you build your own tail call elimination. That is wonderful. – Insulator

@ftor, bingo! Instead of relying on the stack, we rely on the heap (memory) to store our sequence of lambdas. Now we can essentially nest until we run out of memory :D – Pill