generate all possible expressions from a list of basic values

gen_expr.ml

open Sequence.Infix

type op =
  | Add
  | Mult
  | Div
  | Minus
  | Int of int

let pp_op out = function
  | Int i -> Format.fprintf out "%d" i
  | Add -> Format.pp_print_string out "+"
  | Mult -> Format.pp_print_string out "*"
  | Div -> Format.pp_print_string out "/"
  | Minus -> Format.pp_print_string out "-"

let rec pp_expr out l = match l with
  | [] -> ()
  | a :: l -> Format.fprintf out "%a.%a" pp_op a pp_expr l

let eval (e:op list): int =
  (* evaluation with a stack *)
  let rec aux e st = match e, st with
    | [], [i] -> i
    | Int i :: e', _ -> aux e' (i::st)
    | Add :: e', a :: b :: st' -> aux e' ((a+b)::st')
    | Mult :: e', a :: b :: st' -> aux e' ((a*b)::st')
    | Div :: e', a :: b :: st' -> aux e' ((a/b)::st')
    | Minus :: e', a :: b :: st' -> aux e' ((a-b)::st')
    | _ -> assert false
  in
  aux e []

(* generate permutations *)
let rec perms (l:'a list): 'a list Sequence.t = match l with
  | [] -> Sequence.return []
  | x :: xs -> Sequence.flatMap (insert x) (perms xs)

and insert x l = match l with
  | [] -> Sequence.return [x]
  | y :: tail ->
    Sequence.cons
      (x :: l)
      (Sequence.map (fun tail' -> y :: tail') (insert x tail))

let insert_ops (l:int list) : op list Sequence.t =
  let pick_op = Sequence.of_list [Add; Mult; Div; Minus] in
  (* l: remaining integers to use
     n: current height of stack *)
  let rec aux n l res = match l with
    | [] when n=1 -> Sequence.return (List.rev res)
    | [] ->
      assert (n>=2);
      Sequence.flatMap (fun o -> aux (n-1) l (o::res)) pick_op
    | a :: tail ->
      (* can we insert an operator right here? *)
      let s1 =
        if n>=2
        then Sequence.flatMap (fun o -> aux (n-1) l (o::res)) pick_op
        else Sequence.empty
      (* push [a] now *)
      and s2 =
        aux (n+1) tail (Int a::res)
      in
      Sequence.append s2 s1
  in
  aux 0 l []

(* take a list of ints, and make all possible expressions that use
   every int exactly once *)
let all_ops (l:int list): op list Sequence.t =
   Sequence.flatMap insert_ops (perms l)

let () =
  let n =
    if Array.length Sys.argv < 1 then 5 else int_of_string Sys.argv.(1)
  in
  Printf.printf "base list: [1…%d]\n" n;
  let l = Array.init n succ |> Array.to_list in
  all_ops l
  |> Sequence.iter
    (fun e ->
       try Format.printf "@[<2>@[%a@]@ = %d@]@." pp_expr e (eval e)
       with Division_by_zero ->
         Format.printf "@[<2>@[%a@]@ = div_by_zero@]@." pp_expr e
    )