Compiler: simplify branch

Author: hhugo

compiler/lib/eval.ml

@@ -710,7 +710,30 @@ let the_cond_of info x =
       | _ -> Unknown)
     x
 
-let eval_branch update_branch info l =
+module Simple_block = struct
+  type t =
+    | Return of Var.t
+    | Expr of expr
+    | Branch of cont
+
+  let equal a b =
+    match a, b with
+    | Return a, Return b -> Var.equal a b
+    | Expr a, Expr b -> Poly.equal a b
+    | Branch c1, Branch c2 -> Poly.equal c1 c2
+    | Return _, (Expr _ | Branch _)
+    | Expr _, (Return _ | Branch _)
+    | Branch _, (Return _ | Expr _) -> false
+
+  let hash = function
+    | Return x -> Var.idx x
+    | Branch c1 -> Hashtbl.hash c1
+    | Expr x -> Hashtbl.hash x
+end
+
+module SBT = Hashtbl.Make (Simple_block)
+
+let eval_branch blocks update_branch info l =
   match l with
   | Cond (x, ftrue, ffalse) as b -> (
       match the_cond_of info x with
@@ -721,13 +744,43 @@ let eval_branch update_branch info l =
           incr update_branch;
           Branch ftrue
       | Unknown -> b)
-  | Switch (x, a) as b -> (
+  | Switch (x, a) -> (
       match the_cont_of info x a with
       | Some cont ->
           incr update_branch;
           Branch cont
-      | None -> b)
-  | _ as b -> b
+      | None ->
+          let t = SBT.create 18 in
+          Switch
+            ( x
+            , Array.map a ~f:(function
+                | pc, [] -> (
+                    let block = Code.Addr.Map.find pc blocks in
+                    let sb =
+                      match block with
+                      | { body = [] | [ Event _ ]; branch = Return x; params = [] } ->
+                          Some (Simple_block.Return x)
+                      | { body = [ Let (x, e) ] | [ Event _; Let (x, e) ]
+                        ; branch = Return x'
+                        ; params = []
+                        }
+                        when Var.equal x x' -> Some (Simple_block.Expr e)
+                      | { body = [] | [ Event _ ]; branch = Branch c1; params = [] } ->
+                          Some (Simple_block.Branch c1)
+                      | _ -> None
+                    in
+                    match sb with
+                    | None -> pc, []
+                    | Some sb -> (
+                        match SBT.find_opt t sb with
+                        | Some pc' when pc' <> pc ->
+                            incr update_branch;
+                            pc', []
+                        | Some _ | None ->
+                            SBT.add t sb pc;
+                            pc, []))
+                | cont -> cont) ))
+  | cont -> cont
 
 exception May_raise
 
@@ -808,7 +861,7 @@ let eval update_count update_branch inline_constant ~target info blocks =
           block.body
           ~f:(eval_instr update_count inline_constant ~target info)
       in
-      let branch = eval_branch update_branch info block.branch in
+      let branch = eval_branch blocks update_branch info block.branch in
       { block with Code.body; Code.branch })
     blocks
 

compiler/tests-compiler/cond.ml

@@ -104,10 +104,10 @@ let rip_relative_kind_beq x y =
     function rip_relative_kind_beq(x, y){
      switch(x){
        case 0:
-        switch(y){case 0: return 1;case 1: return 0;default: return 0;}
+        return 0 === y ? 1 : 0;
        case 1:
-        switch(y){case 0: return 0;case 1: return 1;default: return 0;}
-       default: switch(y){case 0: return 0;case 1: return 0;default: return 1;}
+        return 1 === y ? 1 : 0;
+       default: return 2 === y ? 1 : 0;
      }
     }
     //end