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

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)

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!