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

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

Quick attempt, seeems to work:

let rec quicksort list cont =
    match list with
    | [] -> cont []
    | element::[] -> cont [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``
            (fun x -> quicksort ``elements larger or equal to pivot`` (fun y -> cont (x @ [pivot] @ y)))

> quicksort [2; 6; 3; 8; 5; 1; 9; 4] id;;
val it : int list = [1; 2; 3; 4; 5; 6; 8; 9]

Edit:

Of course, this code is highly inefficient. I hope nobody will use it in real code. The code was not difficult to write, but continuations might be difficult to read and can be error-prone (it's easy to forget a call to cont). If you want to play more, you can write a continuation monad (Brian wrote a blog post about it).

Continuation monad (stolen from here) can also be used (usually makes code more readable):

type ContinuationMonad() =
    // ma -> (a -> mb) -> mb
    member this.Bind (m, f) = fun c -> m (fun a -> f a c)
    // a -> ma
    member this.Return x = fun k -> k x
    // ma -> ma
    member this.ReturnFrom m = m
let cont = ContinuationMonad()

// Monadic definition of QuickSort
// it's shame F# doesn't allow us to use generic monad code
// (we need to use 'cont' monad here)
// otherwise we could run the same code as Identity monad, for instance
// producing direct (non-cont) behavior
let rec qsm = function
     |[]    -> cont.Return []
     |x::xs -> cont {
        let l,r = List.partition ((>=)x) xs
        let! ls = qsm l 
        let! rs = qsm r
        return (ls @ x :: rs) }

// Here we run our cont with id
let qs xs = qsm xs id     

printf "%A" (qs [2;6;3;8;5;1;9;4])