Efficient implementation of List.split

Learning
1

Small functional programming puzzle: Give a function split_list which accepts a list l and returns (List.take n l, List.drop n l), in constant stack space (or at least stack space strictly less than O(|l|)) and minimizing the overall number of pointer traversals.

It seems to be just outside what the tail_mod_cons annotation is intended to handle: List.take is tail recursive modulo cons, and List.drop is certainly tail recursive, but it seems trickier to combine them.

2

Here’s my solution, time and space both in O(n):

let split_list n l =
  let[@tail_mod_cons] rec aux r n = function
    | [] -> []
    | x :: xs as l -> if n > 0 then x :: aux r (n-1) xs else (r := l; [])
  in
  let r = ref [] in
  let u = aux r n l in
  u, !r

Cheers,
Nicolas

2 Likes
3

My solution would be (to use an explicit accumulator acc, but no mutable data):

let split n l =
  let[@tail_mod_cons] rec aux n l acc =
    match n, l with
    | 0, tl -> List.rev acc, tl
    | _, [] -> List.rev acc, []
    | n, x::l -> aux (n - 1) l (x :: acc)
  in
  aux n l []
1 Like
4

@nojb Great, thank you!

@hannes Nice, that’s the first solution that came to mind for me as well. Even though the performance is $O(n)$, the number of times we have to traverse the tail pointer of a cons cell is 2n, which is what sparked the question, I am curious to find solutions where the number of allocations is $n+O(1)$ rather than $2n+O(1)$.

5

It depends what one means by “mutable data”: tail modulo cons is inherently an imperative technique, mutating blocks after they are created.

Cheers,
Nicolas

1 Like
6

I get what you’re saying. But wouldn’t that imply that all logic programming is imperative, because unification variables are allocated as unset and then later are mutated to take on a value through unification?
e.g.

append(Hd :: Tl, L, Hd :: Z) :- append(Tl, L, Z).
append([], L, L).
7

I think @nojb is being provocative :–) @hannes you don’t need that [@tail_mod_cons] annotation.

1 Like
8

I actually meant my comment in a very precise sense: at runtime, tail modulo cons mutates blocks after creating them and, in particular, triggers the write barrier. Not that it matters for the question at hand, though…

Cheers,
Nicolas

9

The tail mod cons annotation tells the compiler to transform the function to use destination-passing style (DPS). Unlike some other transformations like CPS, DPS is not something that we can safely expose to the user at the moment (it involves manipulating values with holes, and without a linear type system that is very hard to make safe).

Here is an example implementation, using a custom (unsafe) List_dps module:

module List_dps : sig
  type 'a t
  val mk_single : 'a -> 'a t
  val set_tl : 'a t -> 'a t -> unit
  val freeze : 'a t -> 'a list
end = struct
type empty = private E
type non_empty = private N

type ('elt, 'is_empty) mut_list =
  | Empty : (_, empty) mut_list
  | Non_empty : { hd : 'a; mutable tl : 'a any_mut_list } -> ('a, non_empty) mut_list
and 'elt any_mut_list = Any : ('a, _) mut_list -> 'a any_mut_list [@@unboxed]

type 'a t = ('a, non_empty) mut_list

let mk_single hd =
  Non_empty { hd; tl = Any Empty }

let set_tl (l : 'a t) (v : 'a t) =
  match l with
  | Non_empty r -> r.tl <- Any v

external freeze : ('a, _) mut_list -> 'a list = "%opaque"
end

let split_list n l =
  let rec dps_split (dst : _ List_dps.t) n l =
    match n, l with
    | 0, _ | _, [] -> l
    | _, hd :: tl ->
      let new_dst = List_dps.mk_single hd in
      List_dps.set_tl dst new_dst;
      dps_split new_dst (n-1) tl
  in
  if n <= 0 then [], l
  else match l with
    | [] -> [], []
    | hd :: tl ->
      let dst = List_dps.mk_single hd in
      let tail = dps_split dst (n-1) tl in
      List_dps.freeze dst, tail

With more complex types, we could probably make it safer (we should encode that a List_dps.t must be set exactly once on all code paths, and add a way to track whether there are any remaining holes in the list so that freeze can only be called if it gets a proof that no holes remain), but I suspect it would be more interesting to try to extend the OCaml language with native, safe DPS support (I don’t think it’s going to be easy though).

1 Like
10

I don’t think it can be written directly in OxCaml yet (split_list has to be written in C), but if you can guarantee that the input list won’t be used again, OxCaml can give you a type-safe version where the only allocation is the pair for the result (which can also be eliminated…):

# split_list;;
- : int -> 'a list @ unique -> 'a list * 'a list = <fun>
# split_list 2 [1; 2; 3];;
- : int list * int list = ([1; 2], [3])
# let l = [1; 2; 3];;
val l : int list = [1; 2; 3]
# split_list 2 l;;
Error: This value is aliased but is expected to be unique.

(O(n) time; O(1) for both stack and heap space))