javascript - 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

yes, yes, and yes (part 2)

So I believe this answer gets closer to the core of your question – can we make any recursive program stack-safe? Even if recursion isn't in tail position? Even if the host language doesn't have tail-call elimination? Yes. Yes. And yes – with one small requirement...

The end of my first answer talked about loop as a sort of evaluator and then described a rough idea of how it would be implemented. The theory sounded good but I wanted to make sure the technique works in practice. So here we go!

---

non-trivial program

Fibonacci is great for this. The binary recursion implementation builds a big recursive tree and neither recursive call is in tail position. If we can get this program right, we can have reasonable confidence we implemented loop correctly.

And here's that small requirement: You cannot call a function to recur. Instead of f (x), you will write call (f, x) –

const add = (a = 0, b = 0) =>
  a + b

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

fib (10)
// => 55

But these call and recur functions are nothing special. They only create ordinary JS objects –

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

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

So in this program, we have a call that depends on two recurs. Each recur has the potential to spawn yet another call and additional recurs. A non-trivial problem indeed, but in reality we're just dealing with a well-defined recursive data structure.

---

writing loop

If loop is going to process this recursive data structure, it'll be easiest if we can write loop as a recursive program. But aren't we just going to run into a stack-overflow somewhere else then? Let's find out!

// loop : (unit -> 'a expr) -> 'a
const loop = f =>
{ // aux1 : ('a expr, 'a -> 'b) -> 'b 
  const aux1 = (expr = {}, k = identity) =>
    expr.type === recur
      ? // todo: when given { type: recur, ... }
  : expr.type === call
      ? // todo: when given { type: call, ... }
  : k (expr) // default: non-tagged value; no further evaluation necessary

  return aux1 (f ())
}

So loop takes a function to loop, f. We expect f to return an ordinary JS value when the computation is completed. Otherwise return either call or recur to grow the computation.

These todos are somewhat trivial to fill in. Let's do that now –

// loop : (unit -> 'a expr) -> 'a
const loop = f =>
{ // aux1 : ('a expr, 'a -> 'b) -> 'b 
  const aux1 = (expr = {}, k = identity) =>
    expr.type === recur
      ? aux (expr.values, values => aux1 (f (...values), k))
  : expr.type === call
      ? aux (expr.values, values => aux1 (expr.f (...values), k))
  : k (expr)

  // aux : (('a expr) array, 'a array -> 'b) -> 'b
  const aux = (exprs = [], k) =>
    // todo: implement me

  return aux1 (f ())
}

So intuitively, aux1 (“auxiliary one”) is the magic wand we wave over one expression, expr, and the result comes back in the continuation. In other words –

// evaluate expr to get the result
aux1 (expr, result => ...)

To evaluate recur or call, we must first evaluate the corresponding values. We wish we could write something like –

// can't do this!
const r =
  expr.values .map (v => aux1 (v, ...))

return k (expr.f (...r))

What would the continuation ... be? We can't call aux1 in .map like that. Instead, we need another magic wand that can take an array of expressions, and pass the resulting values to its continuation; such as aux –

// evaluate each expression and get all results as array
aux (expr.values, values => ...)

---

meat & potatoes

Ok, this is the probably the toughest part of the problem. For each expression in the input array, we have to call aux1 and chain the continuation to the next expression, finally passing the values to the user-supplied continuation, k –

// aux : (('a expr) array, 'a array -> 'b) -> 'b
const aux = (exprs = [], k) =>
  exprs.reduce
    ( (mr, e) =>
        k => mr (r => aux1 (e, x => k ([ ...r, x ])))
    , k => k ([])
    )
    (k)

We won't end up using this, but it helps to see what we're doing in aux expressed as an ordinary reduce and append –

// cont : 'a -> ('a -> 'b) -> 'b
const cont = x =>
  k => k (x)

// append : ('a array, 'a) -> 'a array
const append = (xs, x) =>
  [ ...xs, x ]

// lift2 : (('a, 'b) -> 'c, 'a cont, 'b cont) -> 'c cont
const lift2 = (f, mx, my) =>
  k => mx (x => my (y => k (f (x, y))))

// aux : (('a expr) array, 'a array -> 'b) -> 'b
const aux = (exprs = [], k) =>
  exprs.reduce
    ( (mr, e) =>
        lift2 (append, mr, k => aux1 (e, k))
    , cont ([])
    )

Putting it all together we get –

// identity : 'a -> 'a
const identity = x =>
  x

// loop : (unit -> 'a expr) -> 'a
const loop = f =>
{ // aux1 : ('a expr, 'a -> 'b) -> 'b 
  const aux1 = (expr = {}, k = identity) =>
    expr.type === recur
      ? aux (expr.values, values => aux1 (f (...values), k))
  : expr.type === call
      ? aux (expr.values, values => aux1 (expr.f (...values), k))
  : k (expr)

  // aux : (('a expr) array, 'a array -> 'b) -> 'b
  const aux = (exprs = [], k) =>
    exprs.reduce
      ( (mr, e) =>
          k => mr (r => aux1 (e, x => k ([ ...r, x ])))
      , k => k ([])
      )
      (k)

  return aux1 (f ())
}

Time for a little celebration –

fib (10)
// => 55

But only a little –

fib (30)
// => RangeError: Maximum call stack size exceeded

---

your original problem

Before we attempt to fix loop, let's revisit the program in your question, foldr, and see how it's expressed using loop, call, and recur –

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

And how does it work?

// small : number array
const small =
  [ 1, 2, 3 ]

// large : number array
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

Okay, it works but for small but it blows up the stack for large. But this is what we expected, right? After all, loop is just an ordinary recursive function, bound for an inevitable stack-overflow... right?

Before we go on, verify the results to this point in your own browser –

// call : (* -> 'a expr, *) -> 'a expr
const call = (f, ...values) =>
  ({ type: call, f, values })

// recur : * -> 'a expr
const recur = (...values) =>
  ({ type: recur, values })

// identity : 'a -> 'a
const identity = x =>
  x

// loop : (unit -> 'a expr) -> 'a
const loop = f =>
{ // aux1 : ('a expr, 'a -> 'b) -> 'b
  const aux1 = (expr = {}, k = identity) =>
    expr.type === recur
      ? aux (expr.values, values => aux1 (f (...values), k))
  : expr.type === call
      ? aux (expr.values, values => aux1 (expr.f (...values), k))
  : k (expr)

  // aux : (('a expr) array, 'a array -> 'b) -> 'b
  const aux = (exprs = [], k) =>
    exprs.reduce
      ( (mr, e) =>
          k => mr (r => aux1 (e, x => k ([ ...r, x ])))
      , k => k ([])
      )
      (k)

  return aux1 (f ())
}

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

// foldr : (('b, 'a) -> 'b, 'b, 'a array) -> 'b
const foldr = (f, init, xs = []) =>
  loop
    ( (i = 0) =>
        i >= xs.length
          ? init
          : call (f, recur (i + 1), xs[i])
    )

// small : number array
const small =
  [ 1, 2, 3 ]

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

console .log (fib (10))
// 55

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 exc