A possible implementation for a cartesian product function of the Enum module.

_tags

<*>: package(batteries)

command.sh

ocamlbuild -use-ocamlfind enum_cartesian_product.byte -- >> output.txt

enum_cartesian_product.ml

open Batteries_uni;;

(* Those function could be put inside the other one (or left outside and added to Enum?) and they need new names *)
type enum_length =
  | Integer of int
  | Infinity
;;

let get_enum_length e =
  try
    Integer (Enum.count e)
  with
  | _ ->
    Infinity
;;

let string_of_enum_length =
  function
    | Infinity -> "+oo"
    | Integer n -> string_of_int n
;;

let enum_length_product a b =
  match a, b with
    | Integer 0, _
    | _, Integer 0 -> Integer 0
    | Infinity, _
    | _, Infinity -> Infinity
    | Integer x, Integer y -> Integer (x*y)
;;

type 'a dllist = {
  mutable first: 'a Dllist.t;
  mutable current: 'a Dllist.t;
  mutable last: 'a Dllist.t
};;

let make_dllist el =
  let node = Dllist.create el in
  {
    first = node;
    current = node;
    last = node
  }
;;

let rec carthesian_product left_enum_arg right_enum_arg =
  (*
  ** Clones the enums.
  ** It is needed to be able to clone the enum and it's the behaviour I'd expect from that function (not modifying the original enums).
  *)
  let left_enum = Enum.clone left_enum_arg in
  let right_enum = Enum.clone right_enum_arg in
  (*
  ** This is the tricky part.
  ** You want to enumerate the carthesian product of two countable and possibly infinite sets.
  ** To do that, I imagine a huge grid where:
  **   - the origin is in the top left corner
  **   - the x-axis is indexed by the elements of left_enum
  **   - the y-axis is indexed by the elements of right_enum
  **   - each square if the grid contains the element (x, y)
  ** Now, to list all the element of the grid, I'll simply list all the ascending diagonals from left to right.
  ** To list this diagonals, I start with a square somewhere on the grid that I will represent by its projection on the x and the y axes, both being dllists.
  ** And I'll apply this algorithm:
  **
  ** - initialize the square so that is contains only (0, 0) (so the dllists are [0] and [0])
  ** - while my square is not reduced to a point (the dllists aren't empty)
  **   - list the diagonal of the square
  **   - if: I can extend the square by adding a line under it and a column on its right, I do it
  **     else if: I can move the square towards the right to downward, I do it (it can't be both because otherwise, I would be able to extend it)
  **     else: I shrink the square by removing its top line and its left-most column
  ** 
  ** This is more or less the algorithm used here. If you got the idea of how the square moves, you should be able to understand what's going on in the code.
  ** 
  *)
  (* I need one element to initialize the dllists so I need a special treatment if one of the enum is empty *)
  match Enum.get left_enum, Enum.get right_enum with
    | None, _
    | _, None -> Enum.empty ()
    | Some x0, Some y0 -> begin
      let x_list = make_dllist x0 in
      let y_list = make_dllist y0 in
      (* This counter is used for clone and count because I couldn't find another way to keep track of where I was in the enum *)
      let i = ref 0 in

      let next () =
        incr i;
        if !i = 1 then
          (x0, y0)
        else begin
          (* enumerate the diagonal if we aren't done yet *)
          if x_list.current != x_list.last then begin
            x_list.current <- Dllist.next x_list.current;
            y_list.current <- Dllist.prev y_list.current;
          (* if we're do move the square *)
          end else begin
            (match Enum.get left_enum, Enum.get right_enum with
              (* if we can extend, extend *)
              | Some xn, Some yn -> begin
                x_list.last <- Dllist.append x_list.last xn;
                y_list.last <- Dllist.append y_list.last yn
              end
              (* if we can't extend, try to move to the right *)
              (* note that we must append the node before dropping the other one or it will fail if the dllist contains a single node and all the elements will appear twice *)
              | Some xn, None -> begin
                x_list.last <- Dllist.append x_list.last xn;
                x_list.first <- Dllist.drop x_list.first
              end
              (* or towards the bottom *)
              | None, Some yn -> begin
                y_list.last <- Dllist.append y_list.last yn;
                y_list.first <- Dllist.drop y_list.first
              end
              (* and if we can't move, shrink *)
              | None, None -> begin
                (* or throw No_more_elements if we can't shrink anymore, which means both sets are finite and we're at the bottom right corner in a 1x1 square *)
                if x_list.first == x_list.last then
                  raise Enum.No_more_elements
                else begin
                  x_list.first <- Dllist.drop x_list.first;
                  y_list.first <- Dllist.drop y_list.first
                end
              end
            );
            x_list.current <- x_list.first;
            y_list.current <- y_list.last
          end;
          (Dllist.get x_list.current, Dllist.get y_list.current)
        end
      in

      let count () =
        match enum_length_product (get_enum_length left_enum_arg) (get_enum_length right_enum_arg) with
          | Integer n -> n - !i
          | Infinity -> raise Enum.Infinite_enum
      in

      let clone () =
        carthesian_product left_enum_arg right_enum_arg
        |> Enum.skip !i
      in

      Enum.make
        ~next: next
        ~count: count
        ~clone: clone
    end

(*
** Another approach would be too keep the whole list of elements of both enum in memory and simply remember what nodes you consider as "first" and a "last" node and ignore the rest.
** The consequences would be:
**   + Multiple enum of the same carthesian product (that you got from cloning one) would be able to share the same dllists
**   + The clone function wouldn't have to skip like it does, it would simply take get a local copy of the current "first" and "last" node in each ddlist
**   - The memory used would tend to +infty in the case where only one of the enum is infinite while in this implementation it doesn't
**   ? If the clone can't be improved (I'm convinced it can but it's getting late so it doesn't mean anything), it will improve the perfs of that method a lot
*)

let print_enum string_of_thing e =
  let e2 = Enum.clone e in
  print_string ((string_of_enum_length (get_enum_length e2)) ^ " < ");
  Enum.iter (fun i -> print_string ((string_of_thing i) ^ " ")) (Enum.take 10 (e2));
  if not (Enum.peek e2 = None) then
    print_string "... "
  else ();
  print_string ">";
  print_newline ()
;;

let string_of_int_pair (a, b) =
  "(" ^ (string_of_int a) ^ ", " ^ (string_of_int b) ^ ")"
;;

let empty_enum = Enum.empty ();;
let finite_enum_1 = 0 -- 0;;
let finite_enum_2 = 0 -- 1;;
let finite_enum_3 = 0 -- 2;;
let infinite_enum =
  let rec init n =
    let i = ref n in
    Enum.make
      ~next: (fun () -> let j = !i in incr i; j)
      ~count: (fun () -> raise Enum.Infinite_enum)
      ~clone: (fun () -> init !i)
  in
  init 0
;;

let enums = [empty_enum; finite_enum_1; finite_enum_2; finite_enum_3; infinite_enum];;

List.iter (fun e -> print_enum string_of_int e) enums;;
print_newline ();;

List.iter (fun e1 ->
  List.iter (fun e2 ->
    print_enum string_of_int e1;
    print_enum string_of_int e2;
    print_enum string_of_int_pair (carthesian_product e1 e2);
    print_newline ()
  ) enums
) enums;;

output.txt

0 < >
1 < 0 >
2 < 0 1 >
3 < 0 1 2 >
+oo < 0 1 2 3 4 5 6 7 8 9 ... >

0 < >
0 < >
0 < >

0 < >
1 < 0 >
0 < >

0 < >
2 < 0 1 >
0 < >

0 < >
3 < 0 1 2 >
0 < >

0 < >
+oo < 0 1 2 3 4 5 6 7 8 9 ... >
0 < >

1 < 0 >
0 < >
0 < >

1 < 0 >
1 < 0 >
1 < (0, 0) >

1 < 0 >
2 < 0 1 >
2 < (0, 0) (0, 1) >

1 < 0 >
3 < 0 1 2 >
3 < (0, 0) (0, 1) (0, 2) >

1 < 0 >
+oo < 0 1 2 3 4 5 6 7 8 9 ... >
+oo < (0, 0) (0, 1) (0, 2) (0, 3) (0, 4) (0, 5) (0, 6) (0, 7) (0, 8) (0, 9) ... >

2 < 0 1 >
0 < >
0 < >

2 < 0 1 >
1 < 0 >
2 < (0, 0) (1, 0) >

2 < 0 1 >
2 < 0 1 >
4 < (0, 0) (0, 1) (1, 0) (1, 1) >

2 < 0 1 >
3 < 0 1 2 >
6 < (0, 0) (0, 1) (1, 0) (0, 2) (1, 1) (1, 2) >

2 < 0 1 >
+oo < 0 1 2 3 4 5 6 7 8 9 ... >
+oo < (0, 0) (0, 1) (1, 0) (0, 2) (1, 1) (0, 3) (1, 2) (0, 4) (1, 3) (0, 5) ... >

3 < 0 1 2 >
0 < >
0 < >

3 < 0 1 2 >
1 < 0 >
3 < (0, 0) (1, 0) (2, 0) >

3 < 0 1 2 >
2 < 0 1 >
6 < (0, 0) (0, 1) (1, 0) (1, 1) (2, 0) (2, 1) >

3 < 0 1 2 >
3 < 0 1 2 >
9 < (0, 0) (0, 1) (1, 0) (0, 2) (1, 1) (2, 0) (1, 2) (2, 1) (2, 2) >

3 < 0 1 2 >
+oo < 0 1 2 3 4 5 6 7 8 9 ... >
+oo < (0, 0) (0, 1) (1, 0) (0, 2) (1, 1) (2, 0) (0, 3) (1, 2) (2, 1) (0, 4) ... >

+oo < 0 1 2 3 4 5 6 7 8 9 ... >
0 < >
0 < >

+oo < 0 1 2 3 4 5 6 7 8 9 ... >
1 < 0 >
+oo < (0, 0) (1, 0) (2, 0) (3, 0) (4, 0) (5, 0) (6, 0) (7, 0) (8, 0) (9, 0) ... >

+oo < 0 1 2 3 4 5 6 7 8 9 ... >
2 < 0 1 >
+oo < (0, 0) (0, 1) (1, 0) (1, 1) (2, 0) (2, 1) (3, 0) (3, 1) (4, 0) (4, 1) ... >

+oo < 0 1 2 3 4 5 6 7 8 9 ... >
3 < 0 1 2 >
+oo < (0, 0) (0, 1) (1, 0) (0, 2) (1, 1) (2, 0) (1, 2) (2, 1) (3, 0) (2, 2) ... >

+oo < 0 1 2 3 4 5 6 7 8 9 ... >
+oo < 0 1 2 3 4 5 6 7 8 9 ... >
+oo < (0, 0) (0, 1) (1, 0) (0, 2) (1, 1) (2, 0) (0, 3) (1, 2) (2, 1) (3, 0) ... >

In the description of the algorithm, there is a "to" that should be an "or". I'll edit later.

Line 70.