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