Tail call optimization with a function returning a tuple

Question

I have a simple function that splits a list at an index:

let rec split_at ls i =
  match i with
  | 0 -> ([], ls)
  | _ ->
    match ls with
    | [] -> raise Not_found
    | h::t ->
      match split_at t (i - 1) with
      | (left, right) -> ((h :: left), right)

Is there a way to get the OCaml compiler to optimize this function to use constant stack space?

I have tried using @tail_mod_cons but it doesn't work. I understand that the call is not really in tail position, but it feels like it should be optimizable.

Answer 1

The function split_at can be written in a partial tail_mod_cons way if we split the construction of the new prefix from the part of the function returning the suffix using a reference:

let[@tail_mod_cons] rec split_at r ls i =
  match i with
  | 0 -> r := ls; []
  | _ ->
    match ls with
    | [] -> raise Not_found
    | h::t ->
       h:: (split_at[@tailcall]) r t (i - 1)

let split_at ls i =
  let r = ref [] in
  let l = split_at r ls i in
  l, !r

Answer 2

The way I understand it, "tail mod cons" works by passing an incomplete constructor into which the called function should place its answer. So to make the optimization work you have to be able to put your problem into a form for which this is a solution.

Maybe it would work if you split the problem into two parts. The first part duplicates the first n elements of the list. The second part returns all but the first n elements of the list.

The second part is trivial to implement tail recursively. And it seems like you should be able to duplicate a list using "tail mod cons".

Answer 3

Firstly, let's clean up your function by pattern matching on a tuple of i and ls and using a local let binding rather than that last match expression.

let rec split_at ls i =
  match i, ls with
  | 0, _    -> ([], ls)
  | _, []   -> raise Not_found 
  | _, h::t -> 
    let (left, right) = split_at t (i - 1) in
    (h::left, right)

As Jeffrey says, the cons (::) is not in tail call position, so tail_mod_cons does nothing for you. If we try to use it, we'll get a warning to that effect:

Lines 1-7, characters 33-20:
Warning 71 [unused-tmc-attribute]: This function is marked @tail_mod_cons
but is never applied in TMC position.

As also hinted, though, it's trivial to let your brain modify this for tail-recursion using an accumulator.

let split_at lst n =
  let rec aux lst n first_part =
    match n, lst with  
    | 0, _    -> (List.rev first_part, lst)
    | _, []   -> raise Not_found
    | _, h::t -> aux t (n - 1) (h::first_part)
  in
  aux lst n []