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) (         


        

Answer 1

hidden powers (part 3)

In our last answer, we made it possible to write foldr using natural expression and the computation remained stack-safe even though the recursive call is not in tail position -

// 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])
    )

This is made possible because loop is effectively an evaluator for the call and recur expressions. But something surprising happened over the last day. It dawned on me that loop has a lot more potential just beneath the surface...

---

first-class continuations

Stack-safe loop is made possible by use of continuation-passing style and I realised that we can reify the continuation and make it available to the loop user: you -

// shift : ('a expr -> 'b expr) -> 'b expr
const shift = (f = identity) =>
  ({ type: shift, f })

// reset : 'a expr -> 'a
const reset = (expr = {}) =>
  loop (() => expr)

const loop = f =>
{ const aux1 = (expr = {}, k = identity) =>
  { switch (expr.type)
    { case recur: // ...
      case call: // ...

      case shift:
        return call
          ( aux1
          , expr.f (x => run (aux1 (x, k)))
          , identity
          )

      default: // ...
    }
  }

  const aux = // ...

  return run (aux1 (f ()))
}

---

examples

In this first example we capture the continuation add(3, ...) (or 3 + ?) in k -

reset
  ( call
      ( add
      , 3
      , shift (k => k (k (1)))
      )
  )

// => 7

We call apply k to 1 and then apply its result to k again -

//        k(?)  = (3 + ?)
//    k (k (?)) = (3 + (3 + ?))
//          ?   = 1
// -------------------------------
// (3 + (3 + 1))
// (3 + 4)
// => 7

The captured continuation can be arbitrarily deep in an expression. Here we capture the continuation (1 + 10 * ?) -

reset
  ( call
      ( add
      , 1
      , call
          ( mult
          , 10
          , shift (k => k (k (k (1))))
          )
      )
  )

// => 1111

Here we'll apply the continuation k three (3) times to an input of 1 -

//       k (?)   =                     (1 + 10 * ?)
//    k (k (?))  =           (1 + 10 * (1 + 10 * ?))
// k (k (k (?))) = (1 + 10 * (1 + 10 * (1 + 10 * ?)))
//          ?    = 1
// ----------------------------------------------------
// (1 + 10 * (1 + 10 * (1 + 10 * 1)))
// (1 + 10 * (1 + 10 * (1 + 10)))
// (1 + 10 * (1 + 10 * 11))
// (1 + 10 * (1 + 110))
// (1 + 10 * 111)
// (1 + 1110)
// => 1111

So far we've been capturing a continuation, k, and then applying it, k (...). Now watch what happens when we use k in a different way -

// r : ?
const r =
  loop
    ( (x = 10) =>
        shift (k => ({ value: x, next: () => k (recur (x + 1))}))
    )

r
// => { value: 10, next: [Function] }

r.next()
// => { value: 11, next: [Function] }

r.next()
// => { value: 11, next: [Function] }

r.next().next()
// => { value: 12, next: [Function] }

A wild stateless iterator appeared! Things are starting to get interesting...

---

harvest and yield

JavaScript generators allow us to produce a lazy stream of values using yield keyword expressions. However when a JS generator is advanced, it is permanently modified -

const gen = function* ()
{ yield 1
  yield 2
  yield 3
}

const iter = gen ()

console.log(Array.from(iter))
// [ 1, 2, 3 ]

console.log(Array.from(iter))
// [] // <-- iter already exhausted!

iter is impure and produces a different output for Array.from each time. This means that JS iterators cannot be shared. If you want to use the iterator in more than one place, you must recompute gen entirely each time -

console.log(Array.from(gen()))
// [ 1, 2, 3 ]

console.log(Array.from(gen()))
// [ 1, 2, 3 ]

As we saw with the shift examples, we can re-use the same continuation many times or save it and call it at a later time. We can effectively implement our own yield but without these pesky limitations. We'll call it stream below -

// emptyStream : 'a stream
const emptyStream =
  { value: undefined, next: undefined }

// stream : ('a, 'a expr) -> 'a stream
const stream = (value, next) =>
  shift (k => ({ value, next: () => k (next) }))

So now we can write our own lazy streams like -

// numbers : number -> number stream
const numbers = (start = 0) =>
  loop
    ( (n = start) =>
        stream (n, recur (n + 1))
    )

// iter : number stream
const iter =
  numbers (10)

iter
// => { value: 10, next: [Function] }

iter.next()
// => { value: 11, next: [Function] }

iter.next().next()
// => { value: 12, next: [Function] }

---

higher-order stream functions

stream constructs an iterator where value is the current value and next is a function that produce the next value. We can write higher-order functions like filter which take a filtering function, f, and an input iterator, iter, and produce a new lazy stream -

// filter : ('a -> boolean, 'a stream) -> 'a stream
const filter = (f = identity, iter = {}) =>
  loop
    ( ({ value, next } = iter) =>
        next
          ? f (value)
            ? stream (value, recur (next ()))
            : recur (next ())
          : emptyStream
    )

const odds =
  filter (x => x & 1 , numbers (1))

odds
// { value: 1, next: [Function] }

odds.next()
// { value: 3, next: [Function] }

odds.next().next()
// { value: 5, next: [Function] }

We'll write take to limit the infinite stream to 20,000 elements and then convert the stream to an array using toArray -

// take : (number, 'a stream) -> 'a stream
const take = (n = 0, iter = {}) =>
  loop
    ( ( m = n
      , { value, next } = iter
      ) =>
        m && next
          ? stream (value, recur (m - 1, next ()))
          : emptyStream
    )

// toArray : 'a stream -> 'a array
const toArray = (iter = {}) =>
  loop
    ( ( r = []
      , { value, next } = iter
      ) =>
        next
          ? recur (push (r, value), next ())
          : r
    )

toArray (take (20000, odds))
// => [ 1, 3, 5, 7, ..., 39999 ]

This is just a start. There are many other stream operations and optimisations we could make to enhance usability and performance.

---

higher-order continuations

With first-class continuations available to us, we can easily make new and interesting kinds of computation possible. Here's a famous "ambiguous" operator, amb, for representing non-deterministic computations -

// amb : ('a array) -> ('a array) expr
const amb = (xs = []) =>
  shift (k => xs .flatMap (x => k (x)))

Intuitively, amb allows you to evaluate an ambiguous expression – one that may return no results, [], or one that returns many, [ ... ] -

// pythag : (number, number, number) -> boolean
const pythag = (a, b, c) =>
  a ** 2 + b ** 2 === c ** 2

// solver : number array -> (number array) array
const solver = (guesses = []) =>
  reset
    ( call
        ( (a, b, c) =>
            pythag (a, b, c) 
              ? [ [ a, b, c ] ] // <-- possible result
              : []              // <-- no result
        , amb (guesses)
        , amb (guesses)
        , amb (guesses)
      )
    )

solver ([ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10 ])
// => [ [ 3, 4, 5 ], [ 4, 3, 5 ], [ 6, 8, 10 ], [ 8, 6, 10 ] ]

And amb is used again here to write product -

// product : (* 'a array) -> ('a array) array
const product = (...arrs) =>
  loop
    ( ( r = []
      , i = 0
      ) =>
        i >= arrs.length
          ? [ r ]
          : call
              ( x => recur ([ ...r, x ], i + 1)
              , amb (arrs [i])
              )
    )


product([ 0, 1 ], [ 0, 1 ], [ 0, 1 ])
// [ [0,0,0], [0,0,1], [0,1,0], [0,1,1], [1,0,0], [1,0,1], [1,1,0], [1,1,1] ]

product([ 'J', 'Q', 'K', 'A' ], [ '♡', '♢', '♤', '♧' ])
// [ [ J, ♡ ], [ J, ♢ ], [ J, ♤ ], [ J, ♧ ]
// , [ Q, ♡ ], [ Q, ♢ ], [ Q, ♤ ], [ Q, ♧ ]
// , [ K, ♡ ], [ K, ♢ ], [ K, ♤ ], [ K, ♧ ]
// , [ A, ♡ ], [ A, ♢ ], [ A, ♤ ], [ A, ♧ ]
// ]

---

full circle

To keep this answer relevant to the post, we'll rewrite foldr using first-class continuations. Of course no one would write foldr like this, but we want to demonstrate that our continuations are robust and complete -

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

foldr (add, "z", "abcefghij")
// => "zjihgfedcba"


foldr (add, "z", "abcefghij".repeat(2000))
// => RangeError: Maximum call stack size exceeded

This is precisely the "deferred overflow" we talked about in the first answer. But since we have full control of the continuations here, we can chain them in a safe way. Simply replace comp above with compExpr and everything works as intended -

// compExpr : ('b expr -> 'c expr, 'a expr -> 'b expr) -> 'a expr -> 'c expr
const compExpr = (f, g) =>
  x => call (f, call (g, x))

foldr (add, "z", "abcefghij".repeat(2000))
// => "zjihgfecbajihgfecbajihgf....edcba"

---

code demonstration

Expand the snippet below to verify the results in your own browser -

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

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

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

// shift : ('a expr -> 'b expr) -> 'b expr
const shift = (f = identity) =>
  ({ type: shift, f })

// reset : 'a expr -> 'a
const reset = (expr = {}) =>
  loop (() => expr)

// amb : ('a array) -> ('a array) expr
const amb = (xs = []) =>
  shift (k => xs .flatMap (x => k (x)))

// add : (number, number) -> number
const add = (x = 0, y = 0) =>
  x + y

// mult : (number, number) -> number
const mult = (x = 0, y = 0) =>
  x * y

// loop : (unit -> 'a expr) -> 'a
const loop = f =>
{ // aux1 : ('a expr, 'a -> 'b) -> 'b
  const aux1 = (expr = {}, k = identity) =>
  { switch (expr.type)
    { case recur:
        return call (aux, f, expr.values, k)
      case call:
        return call (aux, expr.f, expr.values, k)
      case shift:
          return call
            ( aux1
            , expr.f (x => run (aux1 (x, k)))
            , identity
            )
      default:
        return call (k, expr)
    }
  }

  // aux : (* -> 'a, (* expr) array, 'a -> 'b) -> 'b
  const aux = (f, exprs = [], k) =>
  { switch (exprs.length)
    { case 0:
        return call (aux1, f (), k) // nullary continuation
      case 1:
        return call
          ( aux1
          , exprs[0]
          , x => call (aux1, f (x), k) // unary
          )
      case 2:
        return call
          ( aux1
          , exprs[0]
          , x =>
            call
              ( aux1
              , exprs[1]
              , y => call (aux1, f (x, y), k) // binary
              )
          )
      case 3: // ternary ...
      case 4: // quaternary ...
      default: // variadic
        return call
          ( exprs.reduce
              ( (mr, e) =>
                  k => call (mr, r => call (aux1, e, x => call (k, [ ...r, x ])))
              , k => call (k, [])
              )
          , values => call (aux1, f (...values), k)
          )
    }
  }

  return run (aux1 (f ()))
}

// run : * -> *
const run = r =>
{ while (r && r.type === call)
    r = r.f (...r.values)
  return r
}

// example1 : number
const example1 =
  reset
    ( call
        ( add
        , 3
        , shift (k => k (k (1)))
        )
    )

// example2 : number
const example2 =
  reset
    ( call
        ( add
        , 1
        , call
            ( mult
            , 10
            , shift (k => k (k (1)))
            )
        )
    )

// emptyStream : 'a stream
const emptyStream =
  { value: undefined, next: undefined }

// stream : ('a, 'a expr) -> 'a stream
const stream = (value, next) =>
  shift (k => ({ value, next: () => k (next) }))

// numbers : number -> number stream
const numbers = (start = 0) =>
  loop
    ( (n = start) =>
        stream (n, recur (n + 1))
    )

// filter : ('a -> boolean, 'a stream) -> 'a stream
const filter = (f = identity, iter = {}) =>
  loop
    ( ({ value, next } = iter) =>
        next
          ? f (value)
            ? stream (value, recur (next ()))
            : recur (next ())
          : emptyStream
    )

// odds : number stream
const odds =
  filter (x => x & 1 , numbers (1))

// take : (number, 'a stream) -> 'a stream
const take = (n = 0, iter = {}) =>
  loop
    ( ( m = n
      , { value, next } = iter
      ) =>
        m && next
          ? stream (value, recur (m - 1, next ()))
          : emptyStream
    )

// toArray : 'a stream -> 'a array
const toArray = (iter = {}) =>
  loop
    ( ( r = []
      , { value, next } = iter
      ) =>
        next
          ? recur ([ ...r, value ], next ())
          : r
    )

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

// pythag : (number, number, number) -> boolean
const pythag = (a, b, c) =>
  a ** 2 + b ** 2 === c ** 2

// solver : number array -> (number array) array
const solver = (guesses = []) =>
  reset
    ( call
        ( (a, b, c) =>
            pythag (a, b, c)
              ? [ [ a, b, c ] ] // <-- possible result
              : []              // <-- no result
        , amb (guesses)
        , amb (guesses)
        , amb (guesses)
      )
    )

// product : (* 'a array) -> ('a array) array
const product = (...arrs) =>
  loop
    ( ( r = []
      , i = 0
      ) =>
        i >= arrs.length
          ? [ r ]
          : call
              ( x => recur ([ ...r, x ], i + 1)
              , amb (arrs [i])
              )
    )

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

// compExpr : ('b expr -> 'c expr, 'a expr -> 'b expr) -> 'a expr -> 'c expr
const compExpr = (f, g) =>
  x => call (f, call (g, x))

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

// log : (string, 'a) -> unit
const log = (label, x) =>
  console.log(label, JSON.stringify(x))

log("example1:", example1)
// 7

log("example2:", example2)
// 1111

log("odds", JSON.stringify (toArray (take (100, odds))))
// => [ 1, 3, 5, 7, ..., 39999 ]

log("solver:", solver ([ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10 ]))
// => [ [ 3, 4, 5 ], [ 4, 3, 5 ], [ 6, 8, 10 ], [ 8, 6, 10 ] ]

log("product:", product([ 0, 1 ], [ 0, 1 ], [ 0, 1 ]))
// [ [0,0,0], [0,0,1], [0,1,0], [0,1,1], [1,0,0], [1,0,1], [1,1,0], [1,1,1] ]

log("product:", product([ 'J', 'Q', 'K', 'A' ], [ '♡', '♢', '♤', '♧' ]))
// [ [ J, ♡ ], [ J, ♢ ], [ J, ♤ ], [ J, ♧ ]
// , [ Q, ♡ ], [ Q, ♢ ], [ Q, ♤ ], [ Q, ♧ ]
// , [ K, ♡ ], [ K, ♢ ], [ K, ♤ ], [ K, ♧ ]
// , [ A, ♡ ], [ A, ♢ ], [ A, ♤ ], [ A, ♧ ]
// ]

log("foldr:", foldr (add, "z", "abcefghij".repeat(2000)))
// "zjihgfecbajihgfecbajihgf....edcba"

remarks

This was my first time implementing first-class continuations in any language and it was a truly eye-opening experience I wanted to share with others. We got all of this for adding two simple functions shift and reset -

// shift : ('a expr -> 'b expr) -> 'b expr
const shift = (f = identity) =>
  ({ type: shift, f })

// reset : 'a expr -> 'a
const reset = (expr = {}) =>
  loop (() => expr)

And adding the corresponding pattern-match in our loop evaluator -

// ...
case shift:
  return call
    ( aux1
    , expr.f (x => run (aux1 (x, k)))
    , identity
    )

Between stream and amb alone, this is an enormous amount of potential. It makes me wonder just how fast we could make loop such that we could use this in a practical setting.