Fix duplication of code in Cmmgen (#1370)

* Fix duplicates in Cmmgen when handling switches with no default and not all cases

* Improve handling of incomplete Lambda switches in Flambda

* Add test (for reference) and changes

* Fix nitpick

* Cleanup: split the switch stores to reflect usage

* Improve compilation of incomplete switches with flambda

* Split switch stores into context-aware and -unaware versions

* Credit reviewers

* Go back to a single Switch.t_store type

12 files changed, 136 insertions, 38 deletions

diff --git a/Changes b/Changes
index f89fe37a76..ddd8fa2b20 100644
--- a/Changes
+++ b/Changes
@@ -10,6 +10,10 @@ be mentioned in the 4.06 section below instead of here.)
 
 ### Type system:
 
+- GPR#1370: Fix code duplication in Cmmgen
+  (Vincent Laviron, with help from Pierre Chambart,
+   reviews by Gabriel Scherer and Luc Maranget)
+
 ### Standard library:
 
 ### Other libraries:
diff --git a/asmcomp/closure.ml b/asmcomp/closure.ml
index e86ecb6bac..5f8b9e627a 100644
--- a/asmcomp/closure.ml
+++ b/asmcomp/closure.ml
@@ -1258,13 +1258,13 @@ and close_switch fenv cenv cases num_keys default =
   (* First default case *)
   begin match default with
   | Some def when ncases < num_keys ->
-      assert (store.act_store def = 0)
+      assert (store.act_store () def = 0)
   | _ -> ()
   end ;
   (* Then all other cases *)
   List.iter
     (fun (key,lam) ->
-     index.(key) <- store.act_store lam)
+     index.(key) <- store.act_store () lam)
     cases ;
 
   (*  Explicit sharing with catch/exit, as switcher compilation may
diff --git a/asmcomp/cmmgen.ml b/asmcomp/cmmgen.ml
index ad1e58536e..079254d70d 100644
--- a/asmcomp/cmmgen.ml
+++ b/asmcomp/cmmgen.ml
@@ -1475,6 +1475,30 @@ end
 
 (* cmm store, as sharing as normally been detected in previous
    phases, we only share exits *)
+(* Some specific patterns can lead to switches where several cases
+   point to the same action, but this action is not an exit (see GPR#1370).
+   The addition of the index in the action array as context allows to
+   share them correctly without duplication. *)
+module StoreExpForSwitch =
+  Switch.CtxStore
+    (struct
+      type t = expression
+      type key = int option * int
+      type context = int
+      let make_key index expr =
+        let continuation =
+          match expr with
+          | Cexit (i,[]) -> Some i
+          | _ -> None
+        in
+        Some (continuation, index)
+      let compare_key (cont, index) (cont', index') =
+        match cont, cont' with
+        | Some i, Some i' when i = i' -> 0
+        | _, _ -> Pervasives.compare index index'
+    end)
+
+(* For string switches, we can use a generic store *)
 module StoreExp =
   Switch.Store
     (struct
@@ -1495,10 +1519,10 @@ let transl_int_switch loc arg low high cases default = match cases with
 | [] -> assert false
 | _::_ ->
     let store = StoreExp.mk_store () in
-    assert (store.Switch.act_store default = 0) ;
+    assert (store.Switch.act_store () default = 0) ;
     let cases =
       List.map
-        (fun (i,act) -> i,store.Switch.act_store act)
+        (fun (i,act) -> i,store.Switch.act_store () act)
         cases in
     let rec inters plow phigh pact = function
       | [] ->
@@ -2720,10 +2744,10 @@ and transl_switch loc env arg index cases = match Array.length cases with
 | 1 -> transl env cases.(0)
 | _ ->
     let cases = Array.map (transl env) cases in
-    let store = StoreExp.mk_store () in
+    let store = StoreExpForSwitch.mk_store () in
     let index =
       Array.map
-        (fun j -> store.Switch.act_store cases.(j))
+        (fun j -> store.Switch.act_store j cases.(j))
         index in
     let n_index = Array.length index in
     let inters = ref []
diff --git a/asmcomp/flambda_to_clambda.ml b/asmcomp/flambda_to_clambda.ml
index b9d2beb9b7..54ac0befb6 100644
--- a/asmcomp/flambda_to_clambda.ml
+++ b/asmcomp/flambda_to_clambda.ml
@@ -411,13 +411,28 @@ and to_clambda_switch t env cases num_keys default =
     if Numbers.Int.Set.cardinal num_keys = 0 then 0
     else Numbers.Int.Set.max_elt num_keys + 1
   in
-  let index = Array.make num_keys 0 in
   let store = Flambda_utils.Switch_storer.mk_store () in
-  begin match default with
-  | Some def when List.length cases < num_keys -> ignore (store.act_store def)
-  | _ -> ()
+  let default_action =
+    match default with
+    | Some def when List.length cases < num_keys ->
+      store.act_store () def
+    | _ -> -1
+  in
+  let index = Array.make num_keys default_action in
+  let smallest_key = ref num_keys in
+  List.iter
+    (fun (key, lam) ->
+      index.(key) <- store.act_store () lam;
+      smallest_key := min key !smallest_key
+    )
+    cases;
+  if !smallest_key < num_keys then begin
+    let action = ref index.(!smallest_key) in
+    Array.iteri
+      (fun i act ->
+         if act >= 0 then action := act else index.(i) <- !action)
+      index
   end;
-  List.iter (fun (key, lam) -> index.(key) <- store.act_store lam) cases;
   let actions = Array.map (to_clambda t env) (store.act_get ()) in
   match actions with
   | [| |] -> [| |], [| |]  (* May happen when [default] is [None]. *)
diff --git a/bytecomp/bytegen.ml b/bytecomp/bytegen.ml
index 6368cdcaf7..2579b8819f 100644
--- a/bytecomp/bytegen.ml
+++ b/bytecomp/bytegen.ml
@@ -776,13 +776,13 @@ let rec comp_expr env exp sz cont =
       let act_consts = Array.make sw.sw_numconsts 0
       and act_blocks = Array.make sw.sw_numblocks 0 in
       begin match sw.sw_failaction with (* default is index 0 *)
-      | Some fail -> ignore (store.act_store fail)
+      | Some fail -> ignore (store.act_store () fail)
       | None      -> ()
       end ;
       List.iter
-        (fun (n, act) -> act_consts.(n) <- store.act_store act) sw.sw_consts;
+        (fun (n, act) -> act_consts.(n) <- store.act_store () act) sw.sw_consts;
       List.iter
-        (fun (n, act) -> act_blocks.(n) <- store.act_store act) sw.sw_blocks;
+        (fun (n, act) -> act_blocks.(n) <- store.act_store () act) sw.sw_blocks;
 (* Compile and label actions *)
       let acts = store.act_get () in
 (*
diff --git a/bytecomp/matching.ml b/bytecomp/matching.ml
index 25a819cc8f..fb2495d774 100644
--- a/bytecomp/matching.ml
+++ b/bytecomp/matching.ml
@@ -1830,10 +1830,10 @@ let share_actions_tree sw d =
   let d =
     match d with
     | None -> None
-    | Some d -> Some (store.Switch.act_store_shared d) in
+    | Some d -> Some (store.Switch.act_store_shared () d) in
 (* Store all other actions *)
   let sw =
-    List.map  (fun (cst,act) -> cst,store.Switch.act_store act) sw in
+    List.map  (fun (cst,act) -> cst,store.Switch.act_store () act) sw in
 
 (* Retrieve all actions, including potential default *)
   let acts = store.Switch.act_get_shared () in
@@ -1957,14 +1957,14 @@ let share_actions_sw sw =
   | None -> None
   | Some fail ->
       (* Fail is translated to exit, whatever happens *)
-      Some (store.Switch.act_store_shared fail) in
+      Some (store.Switch.act_store_shared () fail) in
   let consts =
     List.map
-      (fun (i,e) -> i,store.Switch.act_store e)
+      (fun (i,e) -> i,store.Switch.act_store () e)
       sw.sw_consts
   and blocks =
     List.map
-      (fun (i,e) -> i,store.Switch.act_store e)
+      (fun (i,e) -> i,store.Switch.act_store () e)
       sw.sw_blocks in
   let acts = store.Switch.act_get_shared () in
   let hs,handle_shared = handle_shared () in
@@ -2032,7 +2032,7 @@ let as_interval_canfail fail low high l =
 
   let do_store _tag act =
 
-    let i =  store.act_store act in
+    let i =  store.act_store () act in
 (*
     eprintf "STORE [%s] %i %s\n" tag i (string_of_lam act) ;
 *)
@@ -2096,7 +2096,7 @@ let as_interval_nofail l =
     | [] ->
         [cur_low, cur_high, cur_act]
     | (i,act)::rem ->
-        let act_index = store.act_store act in
+        let act_index = store.act_store () act in
         if act_index = cur_act then
           i_rec cur_low i cur_act rem
         else
@@ -2110,9 +2110,9 @@ let as_interval_nofail l =
            cases (cf. switch.ml, make_switch).
            Hence, this action will be shared *)
         if some_hole rem then
-          store.act_store_shared act
+          store.act_store_shared () act
         else
-          store.act_store act in
+          store.act_store () act in
       assert (act_index = 0) ;
       i_rec i i act_index rem
   | _ -> assert false in
diff --git a/bytecomp/switch.ml b/bytecomp/switch.ml
index 2e37323921..b03982ddeb 100644
--- a/bytecomp/switch.ml
+++ b/bytecomp/switch.ml
@@ -16,11 +16,11 @@
 
 type 'a shared = Shared of 'a | Single of 'a
 
-type 'a t_store =
+type ('a, 'ctx) t_store =
     {act_get : unit -> 'a array ;
      act_get_shared : unit -> 'a shared array ;
-     act_store : 'a -> int ;
-     act_store_shared : 'a -> int ; }
+     act_store : 'ctx -> 'a -> int ;
+     act_store_shared : 'ctx -> 'a -> int ; }
 
 exception Not_simple
 
@@ -31,7 +31,13 @@ module type Stored = sig
   val make_key : t -> key option
 end
 
-module Store(A:Stored) = struct
+module type CtxStored = sig
+  include Stored
+  type context
+  val make_key : context -> t -> key option
+end
+
+module CtxStore(A:CtxStored) = struct
   module AMap =
     Map.Make(struct type t = A.key let compare = A.compare_key end)
 
@@ -52,7 +58,7 @@ module Store(A:Stored) = struct
       st.next <- i+1 ;
       i in
 
-    let store mustshare act = match A.make_key act with
+    let store mustshare ctx act = match A.make_key ctx act with
     | Some key ->
         begin try
           let (shared,i) = AMap.find key st.map in
@@ -86,6 +92,18 @@ module Store(A:Stored) = struct
      act_get = get; act_get_shared = get_shared; }
 end
 
+module Store(A:Stored) = struct
+  module Me =
+    CtxStore
+      (struct
+        include A
+        type context = unit
+        let make_key () = A.make_key
+      end)
+
+  let mk_store = Me.mk_store
+end
+
 
 
 module type S =
diff --git a/bytecomp/switch.mli b/bytecomp/switch.mli
index 2d0cfd7fcf..b4058c1784 100644
--- a/bytecomp/switch.mli
+++ b/bytecomp/switch.mli
@@ -31,11 +31,11 @@
 
 type 'a shared = Shared of 'a | Single of 'a
 
-type 'a t_store =
+type ('a, 'ctx) t_store =
     {act_get : unit -> 'a array ;
      act_get_shared : unit -> 'a shared array ;
-     act_store : 'a -> int ;
-     act_store_shared : 'a -> int ; }
+     act_store : 'ctx -> 'a -> int ;
+     act_store_shared : 'ctx -> 'a -> int ; }
 
 exception Not_simple
 
@@ -46,9 +46,20 @@ module type Stored = sig
   val make_key : t -> key option
 end
 
+module type CtxStored = sig
+  include Stored
+  type context
+  val make_key : context -> t -> key option
+end
+
+module CtxStore(A:CtxStored) :
+    sig
+      val mk_store : unit -> (A.t, A.context) t_store
+    end
+
 module Store(A:Stored) :
     sig
-      val mk_store : unit -> A.t t_store
+      val mk_store : unit -> (A.t, unit) t_store
     end
 
 (* Arguments to the Make functor *)
@@ -106,13 +117,13 @@ module Make :
           (int * int) ->
            Arg.act ->
            (int * int * int) array ->
-           Arg.act t_store ->
+           (Arg.act, _) t_store ->
            Arg.act
 
 (* Output test sequence, sharing tracked *)
      val test_sequence :
            Arg.act ->
            (int * int * int) array ->
-           Arg.act t_store ->
+           (Arg.act, _) t_store ->
            Arg.act
     end
diff --git a/middle_end/closure_conversion.ml b/middle_end/closure_conversion.ml
index 807889fc4a..e20fe3bfff 100755
--- a/middle_end/closure_conversion.ml
+++ b/middle_end/closure_conversion.ml
@@ -465,12 +465,19 @@ let rec close t env (lam : Lambda.lambda) : Flambda.t =
   | Lswitch (arg, sw, _loc) ->
     let scrutinee = Variable.create "switch" in
     let aux (i, lam) = i, close t env lam in
-    let zero_to_n = Numbers.Int.zero_to_n in
+    let nums sw_num cases default =
+      let module I = Numbers.Int in
+      match default with
+      | Some _ ->
+          I.zero_to_n (sw_num - 1)
+      | None ->
+          List.fold_left (fun set (i, _) -> I.Set.add i set) I.Set.empty cases
+    in
     Flambda.create_let scrutinee (Expr (close t env arg))
       (Switch (scrutinee,
-        { numconsts = zero_to_n (sw.sw_numconsts - 1);
+        { numconsts = nums sw.sw_numconsts sw.sw_consts sw.sw_failaction;
           consts = List.map aux sw.sw_consts;
-          numblocks = zero_to_n (sw.sw_numblocks - 1);
+          numblocks = nums sw.sw_numblocks sw.sw_blocks sw.sw_failaction;
           blocks = List.map aux sw.sw_blocks;
           failaction = Misc.may_map (close t env) sw.sw_failaction;
         }))
diff --git a/middle_end/flambda_utils.mli b/middle_end/flambda_utils.mli
index 37196c06c9..71aa46787c 100644
--- a/middle_end/flambda_utils.mli
+++ b/middle_end/flambda_utils.mli
@@ -165,7 +165,7 @@ val substitute_read_symbol_field_for_variables
 
 (** For the compilation of switch statements. *)
 module Switch_storer : sig
-  val mk_store : unit -> Flambda.t Switch.t_store
+  val mk_store : unit -> (Flambda.t, unit) Switch.t_store
 end
 
 (** Within a set of function declarations there is a set of function bodies,
diff --git a/testsuite/tests/misc/gpr1370.ml b/testsuite/tests/misc/gpr1370.ml
new file mode 100644
index 0000000000..c5d4777f1d
--- /dev/null
+++ b/testsuite/tests/misc/gpr1370.ml
@@ -0,0 +1,19 @@
+type t = A|B|C|D
+type s =
+  | G of t
+  | E of t
+  | H of t
+  | F of (unit list * t)
+  | I of t
+
+let r = ref 0
+
+let set x = r := x
+
+let f x =
+  match x with
+  | E B | F ([()], B) -> set 0
+  | E x | F ([()], x) when Sys.opaque_identity true -> set 1
+  | E _ -> set 2
+  | F _ -> set 3
+  | G _ | H _ | I _ -> set 4
diff --git a/testsuite/tests/misc/gpr1370.reference b/testsuite/tests/misc/gpr1370.reference
new file mode 100644
index 0000000000..e69de29bb2
--- /dev/null
+++ b/testsuite/tests/misc/gpr1370.reference