How to Adapt Trampolines to Continuation Passing Style

How to adapt trampolines to Continuation Passing Style?

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

Continuation-passing-style does not seem to make a difference in Clojure

Clojure has the trampoline function that can remove a lot of the confusing plumbing involved in this problem:

(defn sigma [n]
  (letfn [(sig [curr n]
            (if (<= n 1)
              curr
              #(sig (+ curr (Math/log n)) (dec n))))]
    (trampoline sig 0 n)))

(sigma 1000)
=> 5912.128178488164
(sigma 1500)
=> 9474.406184917756
(sigma 1e7) ;; might take a few seconds
=> 1.511809654875759E8

The function you pass to trampoline can either return a new function, in which case the trampoline continues "bouncing", or a non-function value which would be a "final" value. This example doesn't involve mutually recursive functions, but those are also doable with trampoline.

Does the continuation + tail recursion trick actually trade stack space for heap space?

There is a number of concepts here.

For a tail-recursive function, the compiler can optimize it into a loop and so it does not need any stack or heap space. You can rewrite your count function into a simple tail-recursive function by writing:

let rec count acc n = 
   if n = 0
      then acc
      else count (acc + 1) (n - 1)

This will be compiled into a method with a while loop that makes no recursive calls.

Continuations are generally needed when a function cannot be written as tail-recursive. Then you need to keep some state either on the stack or on the heap. Ignoring the fact that fib can be written more efficiently, the naïve recursive implementation would be:

let fib n = 
  if n <= 1 then 1
  else (fib (n-1)) + (fib (n-2))

This needs stack space to remember what needs to happen after the first recursive call returns the result (we then need to call the other recursive call and add the results). Using continuations, you can turn this into heap-allocated functions:

let fib n cont = 
  if n <= 1 then cont 1
  else fib (n-1) (fun r1 -> 
         fib (n-2) (fun r2 -> cont (r1 + r2))

This allocates one continuation (function value) for each recursive call, but it is tail-recursive so it will not exhaust the available stack space.

Attempting to use continuation passing style to avoid stack overflow with minimax algorithm

As you have undoubtedly seen in the "basic examples", the general idea is to take one extra parameter (the "continuation", usually denoted k), which is a function, and every time you would return a value, pass that value to the continuation instead. So, for example, to modify minimax this way, we'd get:

let rec minimax node depth alpha beta k =
    if depth = 0 || nodeIsTerminal node then
        k (heuristicValue node)
    else
        match node.PlayerToMakeNextChoice with
        | PlayerOneMakesNextChoice ->
            k (takeMax (getChildren node) depth alpha beta)
        | PlayerTwoMakesNextChoice ->
            k (takeMin (getChildren node) depth alpha beta)

And then the call site needs to be "turned inside out", so to speak, so instead of something like this:

let a = minimax ...
let b = f a
let c = g b
c

we would write something like this:

minimax ... (fun a ->
   let b = f a
   let c = g b
   c
)

See? a used to be a return value of minimax, but now a is the parameter of the continuation that is passed to minimax. The runtime mechanics is that, once minimax has finished running, it would call this continuation, and its result value would show up there as parameter a.

So, to apply this to your real code, we'd get this:

| firstChild :: remainingChildren ->
    minimax firstChild (depth - 1) alpha beta (fun minimaxResult ->
        let newAlpha = [alpha; minimaxResult] |> List.max

        if beta < newAlpha then newAlpha
        else takeMax remainingChildren depth newAlpha beta
    )

Ok, this is all well and good, but this is only half the job: we have rewritten minimax in CPS, but takeMin and takeMax are still recursive. Not good.

So let's do takeMax first. Same idea: add an extra parameter k and every time we would "return" a value, pass it to k instead:

and takeMax children depth alpha beta k =      
    match children with
    | [] -> k alpha
    | firstChild :: remainingChildren ->
        minimax firstChild (depth - 1) alpha beta (fun minimaxResult ->
            let newAlpha = [alpha; minimaxResult] |> List.max

            if beta < newAlpha then k newAlpha
            else takeMax remainingChildren depth newAlpha beta k
        )

And now, of course, I have to modify the call site correspondingly:

let minimax ... k =
    ...
    match node.PlayerToMakeNextChoice with
    | PlayerOneMakesNextChoice ->
        takeMax (getChildren node) depth alpha beta k

Wait, what just happened? I just said that every time I return a value I should pass it to k instead, but here I'm not doing it. Instead, I'm passing k itself to takeMax. Huh?

Well, the rule that "instead of returning pass to k" is only the first part of the approach. The second part is - "on every recursive call, pass k down the chain". This way, the original top-level k would travel down the whole recursive chain of calls, and be ultimately called by whatever function decides to stop the recursion.

---

Keep in mind that, although CPS helps with stack overflow, it does not free you from memory limits in general. All those intermediate values don't go on the stack anymore, but they have to go somewhere. In this case, every time we construct that lambda fun minimaxResult -> ..., that's a heap allocation. So all your intermediate values go on the heap.

There is a nice symmetry though: if the algorithm was truly tail-recursive, you'd be able to pass the original top-level continuation down the call chain without allocating any intermediate lambdas, and so you wouldn't require any heap memory.

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

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)

How to implement continuations in dynamic language like JavaScript?

One way to implement continuations in an interpreter is to make that interpreter use its own explicit stack for function calling/returning and parameter passing. If you use the stack of the host language, and that language doesn't have continuations, then things are difficult.

If you use your own explicit stack, you can turn it into a "spaghetti stack" for continuations. A "spaghetti stack" is very similar to common representations of lexical environments. It contains frames, which point to parent frames. Capturing a continuation amounts to retaining a pointer to such a frame, and some code. Resuming a continuation means, more or less, restoring the stack to that frame, and the execution to that point in the code. The interpreter for the language doesn't recurse. When the interpreted language nests or recurses, the interpreter iterates and just pushes and pops the explicit stack to keep track of state.

An alternative approach is to use a linear stack, but copy it when a continuation is made. To resume a continuation, you can restore the entire stack from the copied snapshot. This is useful for delimited continuations, which can avoid copying the entire stack, but only that part of it that is delimited (and restore it on top of the existing stack, rather than by replacement). I have implemented delimited continuations in a language that uses the underlying C stack. A segment of the C stack is memcpy-d out into an object that lives on the heap. When a continuation is restored, that saved stack segment is blasted on top of the current stack. Pointers have to be adjusted, of course, because the segment is now at a different address, and the "dangling cabling" has to be hooked up to integrate that stack segment into the stack properly.

If a language is treated by compilation to CPS (continuation passing style), then continuations pop out for free. Every function has an implicit hidden argument: the continuation it has received. Function returning is compiled into a call to this continuation. If a block of code in the function needs to compute the current continuation, it just has to compose a small local computational future (representable as a lambda) with that incoming continuation (the future computation that happens when the local part returns). Henry Baker wrote a paper based on the observation that under CPS, since no function ever returns (returns compile to tail calls to the continuation), old stack frames are never revisited. The stack can just be allowed to grow, and when it reaches a limit, it can be "rewound" back to the top. Chicken Scheme implements the concept; it's worth investigating if you're interested in continuations. Chicken Scheme compiles to C code that uses the regular C stack. However, the generated functions never return: they simulate returning by calling continuations, and so the stack grows. What is even more fascinating is that objects which we normally understand as dynamic material are allocated from the stack also. Since nothing returns, those objects are safe. Whenever a certain stack limit is reached, all of the objects on the stack which are still live are moved to the heap, and the stack pointer is rewound back to the top.

---

Related Topics