Implementing a tail recursive version of quicksort-like function in F#/OCaml

Question

Is it possible to implement a tail recursive version of the quick sort algorithm (via the continuation pattern)? And if it is, how would one implement it?

Normal (not optimized) version:

let rec quicksort list =
 match list with
 | [] -> []
 | element::[] -> [element]
 | pivot::rest -> let ``elements smaller than pivot``, ``elements larger or equal to pivot``= 
                    rest |> List.partition(fun element -> element < pivot)
                  quicksort ``elements smaller than pivot`` @ [pivot] @ quicksort ``elements larger or equal to pivot``

Answer 1

Direct style:

let rec quicksort list =
    match list with
    | [] -> []
    | [element] -> [element]
    | pivot::rest ->
        let left, right = List.partition (fun element -> element < pivot) rest in
        let sorted_left = quicksort left in
        let sorted_right = quicksort right in
        sorted_left @ [pivot] @ sorted_right

My first, naive translation is very similar to Laurent's version, except indented a bit weirdly to make apparent that calls with continuations are really a kind of binding:

let rec quicksort list cont =
    match list with
    | [] -> cont []
    | element::[] -> cont [element]
    | pivot::rest ->
        let left, right = List.partition (fun element -> element < pivot) rest in
        quicksort left (fun sorted_left ->
        quicksort right (fun sorted_right ->
        cont (sorted_left @ [pivot] @ sorted_right)))
let qsort li = quicksort li (fun x -> x)

Contrarily to Laurent, I find it easy to check that cont is not forgotten: CPS functions translated from direct style have the property that the continuation is used linearily, once and only once in each branch, in tail position. It is easy to check that no such call was forgotten.

But in fact, for most runs of quicksort (supposing you get a roughly logarithmic behavior because you're not unlucky or you shuffled the input first), the call stack is not an issue, as it only grows logarithmically. Much more worrying are the frequent calls to @ wich is linear in its left parameter. A common optimization technique is to define functions not as returning a list but as "adding input to an accumulator list":

let rec quicksort list accu =
    match list with
    | [] -> accu
    | element::[] -> element::accu
    | pivot::rest ->
        let left, right = List.partition (fun element -> element < pivot) rest in
        let sorted_right = quicksort right accu in
        quicksort left (pivot :: sorted_right)
let qsort li = quicksort li []

Of course this can be turned into CPS again:

let rec quicksort list accu cont =
    match list with
    | [] -> cont accu
    | element::[] -> cont (element::accu)
    | pivot::rest ->
        let left, right = List.partition (fun element -> element < pivot) rest in
        quicksort right accu (fun sorted_right ->
        quicksort left (pivot :: sorted_right) cont)
let qsort li = quicksort li [] (fun x -> x)    

Now a last trick is to "defunctionalize" the continuations by turning them into data structure (supposing the allocation of data structures is slightly more efficient than the allocation of a closure):

type 'a cont =
  | Left of 'a list * 'a * 'a cont
  | Return
let rec quicksort list accu cont =
    match list with
    | [] -> eval_cont cont accu
    | element::[] -> eval_cont cont (element::accu)
    | pivot::rest ->
        let left, right = List.partition (fun element -> element < pivot) rest in
        quicksort right accu (Left (left, pivot, cont))
and eval_cont = function
  | Left (left, pivot, cont) ->
    (fun sorted_right -> quicksort left (pivot :: sorted_right) cont)
  | Return -> (fun x -> x)
let qsort li = quicksort li [] Return

Finally, I chose the function .. fun style for eval_cont to make it apparent that those were just pieces of code from the CPS version, but the following version is probably better optimized by arity-raising:

and eval_cont cont accu = match cont with
  | Left (left, pivot, cont) ->
    quicksort left (pivot :: accu) cont
  | Return -> accu