How to write fold_left in CPS?

Question

Even though it's already tail-recursive, It's still interesting to see a CPS version of it.

Here's the standart fold_left:

let rec myFoldLeft f acc list =
  match list with
  | [] -> acc
  | h :: tl -> myFoldLeft f (f h acc) tl
;;

Here's an example of fold_right

let rec myFoldRightCPS f acc list cont =
  match list with
  | [] ->
      cont acc
  | h :: tl ->
      myFoldRightCPS f acc tl (fun t -> cont @@ f t h)

I'm using OCaml for this one

Answer 1

The function cps_fold_left needs to get an accumulator, a folding function, and a function to call with the result of the fold.

Also, the folding functions is also in CPS style, meaning that it accepts an accumulator value, a new value, and a continuation to call to send its result.

As we will see in the implementation, you can even accept two continuations in the folding function: the continuation to call when you want to continue folding, and a continuation to call if you want to stop folding.

For example if you fold a multiplication, the result is zero as soon as you encounter a zero, there is no need to iterate over the rest of the list.

let cps_mult a v stop next =
  if a = 0 
  then (stop 0) 
  else (next (a * v))

val cps_mult : int -> int -> (int -> 'a) -> (int -> 'a) -> 'a = <fun>

If you call the function with an accumulator that is equal to zero, you call stop, otherwise, you call next.

If you want to write cps_fold_left, you need to remember that you don't call the function recursively, you will call the given fold function, for example cps_mult, with arguments that are functions to call on completion.

let rec cps_fold_left fold acc xs stop =
match xs with
| [] -> stop acc
| x::xs ->
  let next acc = cps_fold_left fold acc xs stop in
  fold acc x stop next

You can test is as follows:

cps_fold_left cps_mult 1 [3;4] (fun x -> x);;
- : int = 12

Lazy iterators

Having the user-provided folding function be in CPS form itself is useful (e.g. you can short-circuit iteration). Also, notice that the result of cps_fold_left is also the result of fold, so you can return a value from fold, and by doing so you can interrupt the iteration.

Here for example I define a cursor as being either a result (R) or a cursor (C) which is an intermediate result and function of no argument that produces a next cursor.

type 'a cursor = R of 'a | C of 'a * (unit -> 'a cursor)

I can write lazy_mult as follows:

let lazy_mult a v _ next = C (a, (fun () -> next (a * v)));;

For now I don't care about the stop case, but this can be encapsulated in the cursor too. Then I can call the same cps_left_fold function to now have a lazy stream of values:

# let cursor = cps_fold_left lazy_mult 1 [2; 3; 4] (fun x -> R x);;
val cursor : int cursor = C (1, <fun>)

Thanks to the following iter function

let iter = function
| R v -> R v
| C (_, f) -> f ();;
val iter : 'a cursor -> 'a cursor = <fun>

I can iter cursor:

# iter cursor;;
- : int cursor = C (2, <fun>)

And with that, you can iterate over two lists concurrently, something that is hard to do with plain higher-order functions like List.fold_left.

Remark

But can you write the version that doesn't require any special operations like cps_mult?

Using cps_fold_left as define above, it is possible to write a lift function that takes a regular operation like ( * ) and turn it into a CPS style function:

let lift f = (fun acc value _stop next -> next @@ f acc value);;
val lift : ('a -> 'b -> 'c) -> 'a -> 'b -> 'd -> ('c -> 'e) -> 'e = <fun>

let cps_mult = lift ( * );;

Then either you can change cps_fold_left so that it calls lift itself (and it only accept regular functions), or you let the user call it manually.

The other option I see is to just call fold acc x, get a new value, and directly call cps_fold_left with the same continuation. That would be identical to fold_left except that you eventually call a continuation on the result, and that could be refactored as cont @@ fold_left fold acc xs. That doesn't seem useful to me, I mean there is no need to write the function in CPS style in that case.

Or, as in your example, the computation could happen in a fresh continuation, but I don't see a simple way to do that without allocating data to apply the function in the right order.

In your example of myFoldRightCPS, a new continuation is created at each level of recursion to make sure that the rightmost element is applied first to the initial accumulator value, and then the result is passed to the previous value, etc. until the function is applied to the first element.

let rec myFoldRightCPS f acc list cont =
  match list with
  | [] ->
      cont acc
  | h :: tl ->
      myFoldRightCPS f acc tl (fun t -> cont @@ f t h)

This is useful because it makes myFoldRightCPS tail-recursive, unlike the direct implementation, and then continuations calling themselves are recursive too, and also tail-recursive. So from that point of view there is an advantage in using the CPS style, which is that the call stack usage is constant w.r.t. input size. The code still allocates data for the intermediate closures, but that's expected.

For the fold left operation, there is no need to delay the computation of the folding operation, and in fact once you delay it you have to reverse the order of operation to obtain the original behavior.