How to adapt trampolines to Continuation Passing Style?

Question

Here is a naive implementation of a right fold:

const foldr = f => acc => ([x, ...xs]) =>
  x === undefined
    ? acc 
    : f(x) (foldkr(f) (acc) (xs));

This is non-tail recursion and hence we cannot apply a trampoline. One approach would be to make the algorithm iterative and use a stack to mimick the function call stack.

Another approch would be to transform the recursion into CPS:

const Cont = k => ({runCont: k});

const foldkr = f => acc => ([x, ...xs]) =>
  Cont(k =>
    x === undefined
      ? k(acc)
      : foldkr(f) (acc) (xs)
          .runCont(acc_ => k(f(x) (acc_))));

This is still naive, because it is insanely slow. Here is a less memory consuming version:

const foldkr = f => acc => xs => {
  const go = i =>
    Cont(k =>
      i === xs.length
        ? k(acc)
        : go(i + 1)
            .runCont(acc_ => k(f(xs[i]) (acc_))));

  return go(0);
};

The recursive call is now in tail position hence we should be able to apply a trampoline of our choice:

const loop = f => {
  let step = f();

  while (step && step.type === recur)
    step = f(...step.args);

  return step;
};

const recur = (...args) =>
  ({type: recur, args});

const foldkr = f => acc => xs =>
  loop((i = 0) => 
    Cont(k =>
      i === xs.length
        ? k(acc)
        : recur(i + 1)
            .runCont(acc_ => k(f(xs[i]) (acc_)))));

This doesn't work, because the trampoline call is inside the continuation and thus lazily evaluated. How must the trampoline be adapted so that it works with CPS?

Answer 1

tail calls first (part 1)

First write the loop such that it recurs in tail position

const foldr = (f, init, xs = []) =>
  loop
    ( ( i = 0
      , k = identity
      ) =>
        i >= xs.length 
          ? k (init)
          : recur
              ( i + 1
              , r => k (f (r, xs[i]))
              )
   )

Given two inputs, small and large, we test foldr -

const small =
  [ 1, 2, 3 ]

const large =
  Array.from (Array (2e4), (_, n) => n + 1)

foldr ((a, b) => `(${a}, ${b})`, 0, small)
// => (((0, 3), 2), 1)

foldr ((a, b) => `(${a}, ${b})`, 0, large)
// => RangeError: Maximum call stack size exceeded

But it uses a trampoline, why does it fail for large? The short answer is because we built a huge deferred computation, k ...

loop
  ( ( i = 0
    , k = identity // base computation
    ) =>
      // ...
      recur // this gets called 20,000 times
        ( i + 1
        , r => k (f (r, xs[i])) // create new k, deferring previous k
        )
  )

In the terminating condition, we finally call k(init) which fires off the stack of deferred computations, 20,000 function calls deep, which triggers the stack-overflow.

Before reading on, expand the snippet below to make sure we're on the same page -

const identity = x =>
  x
  
const loop = f =>
{ let r = f ()
  while (r && r.recur === recur)
    r = f (...r.values)
  return r
}

const recur = (...values) =>
  ({ recur, values })

const foldr = (f, init, xs = []) =>
  loop
    ( ( i = 0
      , k = identity
      ) =>
        i >= xs.length 
          ? k (init)
          : recur
              ( i + 1
              , r => k (f (r, xs[i]))
              )
   )

const small =
  [ 1, 2, 3 ]

const large =
  Array.from (Array (2e4), (_, n) => n + 1)

console.log(foldr ((a, b) => `(${a}, ${b})`, 0, small))
// (((0, 3), 2), 1)

console.log(foldr ((a, b) => `(${a}, ${b})`, 0, large))
// RangeError: Maximum call stack size exceeded

---

deferred overflow

The problem we're seeing here is the same one you might encounter if you were to compose(...) or pipe(...) 20,000 functions together -

// build the composition, then apply to 1
foldl ((r, f) => (x => f (r (x))), identity, funcs) (1)

Or similar using comp -

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

// build the composition, then apply to 1
foldl (comp, identity, funcs) 1

Sure, foldl is stack-safe and it can compose 20,000 functions, but as soon as you call the massive composition, you risk blowing the stack. Now compare that to -

// starting with 1, fold the list; apply one function at each step
foldl ((r, f) => f (r), 1, funcs)

... which does not blow the stack because the computations are not deferred. Instead the result from one step overwrites the result from the previous step until the final step is reached.

In fact, when we write -

r => k (f (r, xs[i]))

Another way to see this is -

comp (k, r => f (r, xs[i]))

This should highlight exactly where the problem is.

---

possible solution

One simple remedy is to add a separate call tag that flattens the deferred computation in the trampoline. So instead of calling a function directly like f (x), we'll write call (f, x) -

const call = (f, ...values) =>
  ({ call, f, values })

const foldr = (f, init, xs = []) =>
  loop
    ( ( i = 0
      , k = identity
      ) =>
        i >= xs.length 
          // k (init) rewrite as
          ? call (k, init)
          : recur
              ( i + 1
              // r => k (f (r, xs[i])) rewrite as
              , r => call (k, f (r, xs[i]))
              )
   )

We modify the trampoline to act on call-tagged values -

const loop = f =>
{ let r = f ()
  while (r)
    if (r.recur === recur)
      r = f (...r.values)
    else if (r.call === call)
      r = r.f (...r.values)
    else
      break
  return r
}

Finally, we see that the large input no longer overflows the stack -

foldr ((a, b) => `(${a}, ${b})`, 0, small)
// => (((0, 3), 2), 1)

foldr ((a, b) => `(${a}, ${b})`, 0, large)
// => (Press "Run snippet" below see results ...)

const identity = x =>
  x
  
const loop = f =>
{ let r = f ()
  while (r)
    if (r.recur === recur)
      r = f (...r.values)
    else if (r.call === call)
      r = r.f (...r.values)
    else
      break
  return r
}

const recur = (...values) =>
  ({ recur, values })
  
const call = (f, ...values) =>
  ({ call, f, values })

const foldr = (f, init, xs = []) =>
  loop
    ( ( i = 0
      , k = identity
      ) =>
        i >= xs.length 
          ? call (k, init)
          : recur
              ( i + 1
              , r => call (k, f (r, xs[i]))
              )
   )
   
const small =
  [ 1, 2, 3 ]

const large =
  Array.from (Array (2e4), (_, n) => n + 1)

console.log(foldr ((a, b) => `(${a}, ${b})`, 0, small))
// (((0, 3), 2), 1)

console.log(foldr ((a, b) => `(${a}, ${b})`, 0, large))
// (Press "Run snippet" to see results ...)

---

wups, you built your own evaluator

Above, recur and call appear to be magic functions. But in reality, recur and call create simple objects { ... } and loop is doing all of the work. In this way, loop is a type of evaluator that accepts recur and call expressions. The one down-side to this solution is that we expect the caller always to use recur or call in tail position, otherwise the loop will return an incorrect result.

This is different than the Y-combinator which reifies the recursion mechanism as a parameter, and is not limited to a tail-only position, such as recur here -

const Y = f => f (x => Y (f) (x))

const fib = recur => n =>
  n < 2
    ? n
    : recur (n - 1) + recur (n - 2) // <-- non-tail call supported
    
console .log (Y (fib) (30))
// => 832040

The one down-side to Y is, of course, because you control recursion by calling a function, you are still stack-unsafe just like all other functions in JS. The result is a stack-overflow -

console .log (Y (fib) (100))
// (After a long time ...)
// RangeError: Maximum call stack size exceeded

So would it be possible to support recur in non-tail position and remain stack-safe? Sure, a sufficiently clever loop should be able evaluate recursive expressions -

const fib = (init = 0) =>
  loop
    ( (n = init) =>
        n < 2
          ? n
          : call
              ( (a, b) => a + b
              , recur (n - 1)
              , recur (n - 2)
              ) 
    )

fib (30)
// expected: 832040

loop becomes a CPS tail-recursive function for evaluating the input expressions call, recur, etc. Then we put loop on a trampoline. loop effectively becomes an evaluator for our custom language. Now you can forget all about the stack – your only limitation now is memory!

Alternatively -

const fib = (n = 0) =>
  n < 2
    ? n
    : call
        ( (a, b) => a + b
        , call (fib, n - 1)
        , call (fib, n - 2)
        )

loop (fib (30))
// expected: 832040

In this related Q&A, I write a normal-order evaluator for untyped lambda calculus in JavaScript. It shows how you can write programs that are freed from the implementation effects (evaluation strategy, stack model, etc) of the host language. There we're using Church-encoding, here were using call and recur, but the technique is the same.

Years back, I wrote a stack-safe variation using the technique I described above. I'll see if I can ressurrect it and later make it available in this answer. For now, I'll leave the loop evaluator as an exercise for the reader.

PART 2 added: loop evaluator

---

alternative solution

In this related Q&A, we build a stack-safe continuation monad.