Hash-based value numbering to eliminate redundant expressions

value_numbering.ml

type stmts = stmt ref list

and stmt = Assign of var * expr

and expr =
  | Int of int
  | Val of var
  | Binop of expr * expr (* e1 + e2 *)

and var = string

module Vars = Set.Make (struct
  type t = var
  let compare = Stdlib.compare
end)

let rec string_of_expr = function
  | Int i -> string_of_int i
  | Val x -> x
  | Binop (e1, e2) ->
    string_of_expr e1 ^ " + " ^ string_of_expr e2

let string_of_stmt = function
  | Assign (x, e) -> x ^ " := " ^ string_of_expr e

let value_numbers : (var, int) Hashtbl.t = Hashtbl.create 10

let available_exprs : (int, Vars.t) Hashtbl.t = Hashtbl.create 10

let gen_number init =
  let count = ref init in
  fun () ->
    let c = !count in
    incr count; c

let next_value_number = gen_number 1

let rec value_number = function
  | Int i -> (
      let s = string_of_int i in
      match Hashtbl.find_opt value_numbers s with
      | Some n -> n
      | None ->
        let n = next_value_number () in
        Hashtbl.add value_numbers s n; n
    )
  | Val x -> (
      match Hashtbl.find_opt value_numbers x with
      | Some n -> n
      | None ->
        let n = next_value_number () in
        Hashtbl.add value_numbers x n; n
    )
  | Binop (e1, e2) -> (
      let n1 = value_number e1 in
      let n2 = value_number e2 in
      let s = string_of_expr (
          if n1 < n2 then Binop (Int n1, Int n2)
          else Binop (Int n2, Int n1)
        )
      in
      match Hashtbl.find_opt value_numbers s with
      | Some n -> n
      | None ->
        let n = next_value_number () in
        Hashtbl.add value_numbers s n; n
    )

let available vn =
  match Hashtbl.find_opt available_exprs vn with
  | Some xs ->
    (* Expression e is redundant *)
    Vars.choose_opt xs
  | None -> None

let update x vn =
  (* Is this the first assignment to x? *)
  begin match Hashtbl.find_opt value_numbers x with
    | Some m -> (
        match Hashtbl.find_opt available_exprs m with
        | Some xs ->
          assert (Vars.mem x xs);
          Hashtbl.replace available_exprs m (Vars.remove x xs);
          Hashtbl.replace value_numbers x vn
        | None ->
          (* x only appeared in expressions *)
          Hashtbl.replace value_numbers x vn
      )
    | None ->
      Hashtbl.add value_numbers x vn
  end;
  match Hashtbl.find_opt available_exprs vn with
  | Some xs -> Hashtbl.replace available_exprs vn (Vars.add x xs)
  | None -> Hashtbl.add available_exprs vn (Vars.singleton x)

let visit stmt =
  match !stmt with
  | Assign (x, e) ->
    let vn = value_number e in
    let _ = match e, available vn with
      | Binop _, Some y -> stmt := Assign (x, Val y)
      | _ -> ()
    in
    update x vn

let ( >> ) f g x = g (f x)

let print_stmts =
  List.iter (( ! ) >> string_of_stmt >> print_endline)

let print_value_numbers () =
  print_endline "Value numbers:";
  Hashtbl.iter (Printf.printf "%s |-> %d\n") value_numbers

let print_available_exprs () =
  print_endline "Available expressions:";
  Hashtbl.iter (fun vn e ->
      let xs = String.concat ", " (Vars.elements e) in
      Printf.printf "%d |-> {%s}\n" vn xs
    ) available_exprs

(* A few examples *)

let t1 = List.map ref [
  (* a := 4 *)
  Assign ("a", Int 4);
  (* b := 5 *)
  Assign ("b", Int 5);
  (* c := a + b *)
  Assign ("c", Binop (Val "a", Val "b"));
  (* d := 5 *)
  Assign ("d", Int 5);
  (* e := a + d *)
  Assign ("e", Binop (Val "a", Val "d"));
]

let t2 = List.map ref [
  (* a := 1 *)
  Assign ("a", Int 1);
  (* b := 2 *)
  Assign ("b", Int 2);
  (* c := a + b *)
  Assign ("c", Binop (Val "a", Val "b"));
  (* b := 3 *)
  Assign ("b", Int 3);
  (* d := a + b *)
  Assign ("d", Binop (Val "a", Val "b"));
]

let t3 = List.map ref [
  (* a := x + y *)
  Assign ("a", Binop (Val "x", Val "y"));
  (* b := x + y *)
  Assign ("b", Binop (Val "x", Val "y"));
  (* a := 1 *)
  Assign ("a", Int 1);
  (* c := y + x *)
  Assign ("c", Binop (Val "y", Val "x"));
  (* b := 2 *)
  Assign ("b", Int 2);
  (* c := 3 *)
  Assign ("c", Int 3);
  (* d := x + y *)
  Assign ("d", Binop (Val "x", Val "y"));
]

let test stmts =
  print_stmts stmts;
  List.iter visit stmts;
  print_newline ();
  print_value_numbers ();
  print_newline ();
  print_available_exprs ();
  print_newline ();
  print_endline "After optimization:";
  print_stmts stmts

let () =
  print_endline "\nExample 1:";
  test t1;
  print_endline "\nExample 2:";
  test t2;
  print_endline "\nExample 3:";
  test t3

utop # use "value_numbering.ml";;
(...)

Example 1:
a := 4
b := 5
c := a + b
d := 5
e := a + d

Value numbers:
a |-> 1
d |-> 2
4 |-> 1
e |-> 3
5 |-> 2
b |-> 2
1 + 2 |-> 3
c |-> 3

Available expressions:
2 |-> {b, d}
3 |-> {c, e}
1 |-> {a}

After optimization:
a := 4
b := 5
c := a + b
d := 5
e := c

Example 2:
a := 1
b := 2
c := a + b
b := 3
d := a + b

Value numbers:
a |-> 4
d |-> 8
4 |-> 1
e |-> 3
5 |-> 2
b |-> 7
3 |-> 7
2 |-> 5
4 + 5 |-> 6
4 + 7 |-> 8
1 |-> 4
1 + 2 |-> 3
c |-> 6

Available expressions:
6 |-> {c}
2 |-> {}
8 |-> {d}
7 |-> {b}
3 |-> {e}
5 |-> {}
4 |-> {a}
1 |-> {}

After optimization:
a := 1
b := 2
c := a + b
b := 3
d := a + b

Example 3:
a := x + y
b := x + y
a := 1
c := y + x
b := 2
c := 3
d := x + y

Value numbers:
9 + 10 |-> 11
a |-> 4
d |-> 11
4 |-> 1
e |-> 3
5 |-> 2
y |-> 10
b |-> 5
3 |-> 7
2 |-> 5
x |-> 9
4 + 5 |-> 6
4 + 7 |-> 8
1 |-> 4
1 + 2 |-> 3
c |-> 7

Available expressions:
6 |-> {}
2 |-> {}
8 |-> {}
7 |-> {c}
3 |-> {e}
5 |-> {b}
4 |-> {a}
11 |-> {d}
1 |-> {}

After optimization:
a := x + y
b := a
a := 1
c := b
b := 2
c := 3
d := x + y

Compare this implementation with an implementation for programs in SSA form, which is simpler and more effective because no expression ever becomes inaccessible (like x + y in Example 3 above): https://github.com/aprell/compiler-potpourri/blob/master/ssa/value_numbering.ml