Fix MPR#7259 by implementing switch branch sharing for flambda (#603)

10 files changed, 99 insertions, 19 deletions

diff --git a/Changes b/Changes
index 2742ca6e49..8b144f8ed2 100644
--- a/Changes
+++ b/Changes
@@ -1382,6 +1382,10 @@ OCaml 4.04.0 (4 Nov 2016):
   it now calls (nano)sleep for 0 seconds as in (< 4.03) versions.
   (Hannes Mehnert, review by Damien Doligez)
 
+- PR#7259 and GPR#603: flambda does not collapse pattern matching
+  in some cases
+  (Pierre Chambart, report by Reed Wilson, review by Mark Shinwell)
+
 - PR#7260: GADT + subtyping compile time crash
   (Jacques Garrigue, report by Nicolas Ojeda Bar)
 
diff --git a/asmcomp/closure.ml b/asmcomp/closure.ml
index 1154778834..e86ecb6bac 100644
--- a/asmcomp/closure.ml
+++ b/asmcomp/closure.ml
@@ -28,6 +28,7 @@ module Storer =
       type t = lambda
       type key = lambda
       let make_key =  Lambda.make_key
+      let compare_key = Pervasives.compare
     end)
 
 (* Auxiliaries for compiling functions *)
diff --git a/asmcomp/cmmgen.ml b/asmcomp/cmmgen.ml
index 0c4b3c1247..4b6739087e 100644
--- a/asmcomp/cmmgen.ml
+++ b/asmcomp/cmmgen.ml
@@ -1443,6 +1443,7 @@ module StoreExp =
       let make_key = function
         | Cexit (i,[]) -> Some i
         | _ -> None
+      let compare_key = Pervasives.compare
     end)
 
 module SwitcherBlocks = Switch.Make(SArgBlocks)
diff --git a/asmcomp/un_anf.ml b/asmcomp/un_anf.ml
index 940dccc727..b0440959a4 100644
--- a/asmcomp/un_anf.ml
+++ b/asmcomp/un_anf.ml
@@ -332,8 +332,7 @@ let let_bound_vars_that_can_be_moved ident_info (clam : Clambda.ulambda) =
       let_stack := []
     | Ustaticfail (static_exn, args) ->
       ignore_int static_exn;
-      ignore_ulambda_list args;
-      let_stack := []
+      examine_argument_list args
     | Ucatch (static_exn, idents, body, handler) ->
       ignore_int static_exn;
       ignore_ident_list idents;
diff --git a/bytecomp/bytegen.ml b/bytecomp/bytegen.ml
index ec02835256..274ad1c431 100644
--- a/bytecomp/bytegen.ml
+++ b/bytecomp/bytegen.ml
@@ -443,6 +443,7 @@ let is_immed n = immed_min <= n && n <= immed_max
 module Storer =
   Switch.Store
     (struct type t = lambda type key = lambda
+      let compare_key = Pervasives.compare
       let make_key = Lambda.make_key end)
 
 (* Compile an expression.
diff --git a/bytecomp/matching.ml b/bytecomp/matching.ml
index d559549592..25013a6578 100644
--- a/bytecomp/matching.ml
+++ b/bytecomp/matching.ml
@@ -467,6 +467,7 @@ module StoreExp =
     (struct
       type t = lambda
       type key = lambda
+      let compare_key = Pervasives.compare
       let make_key = Lambda.make_key
     end)
 
diff --git a/bytecomp/switch.ml b/bytecomp/switch.ml
index 07c0721dd5..6dea87e4d5 100644
--- a/bytecomp/switch.ml
+++ b/bytecomp/switch.ml
@@ -27,12 +27,13 @@ exception Not_simple
 module type Stored = sig
   type t
   type key
+  val compare_key : key -> key -> int
   val make_key : t -> key option
 end
 
 module Store(A:Stored) = struct
   module AMap =
-    Map.Make(struct type t = A.key let compare = Pervasives.compare end)
+    Map.Make(struct type t = A.key let compare = A.compare_key end)
 
   type intern =
       { mutable map : (bool * int)  AMap.t ;
diff --git a/bytecomp/switch.mli b/bytecomp/switch.mli
index efc345ad6f..0b1da8946b 100644
--- a/bytecomp/switch.mli
+++ b/bytecomp/switch.mli
@@ -42,6 +42,7 @@ exception Not_simple
 module type Stored = sig
   type t
   type key
+  val compare_key : key -> key -> int
   val make_key : t -> key option
 end
 
diff --git a/middle_end/flambda_utils.ml b/middle_end/flambda_utils.ml
index efd4608280..f5f44b3690 100644
--- a/middle_end/flambda_utils.ml
+++ b/middle_end/flambda_utils.ml
@@ -734,18 +734,93 @@ let substitute_read_symbol_field_for_variables
   in
   Flambda_iterators.map_toplevel f (fun v -> v) expr
 
-(* CR-soon mshinwell: implement this so that sharing can occur in
-   matches.  Should probably leave this for the first release. *)
-type sharing_key = unit
-let make_key _ = None
-
-module Switch_storer =
-  Switch.Store
-    (struct
-      type t = Flambda.t
-      type key = sharing_key
-      let make_key = make_key
-    end)
+module Switch_storer = Switch.Store (struct
+  type t = Flambda.t
+
+  (* An easily-comparable subset of [Flambda.t]: currently this only
+     supports that required to share switch branches. *)
+  type key =
+    | Var of Variable.t
+    | Let of Variable.t * key_named * key
+    | Static_raise of Static_exception.t * Variable.t list
+  and key_named =
+    | Symbol of Symbol.t
+    | Const of Flambda.const
+    | Prim of Lambda.primitive * Variable.t list
+    | Expr of key
+
+  exception Not_comparable
+
+  let rec make_expr_key (expr : Flambda.t) : key =
+    match expr with
+    | Var v -> Var v
+    | Let { var; defining_expr; body; } ->
+      Let (var, make_named_key defining_expr, make_expr_key body)
+    | Static_raise (e, args) -> Static_raise (e, args)
+    | _ -> raise Not_comparable
+  and make_named_key (named:Flambda.named) : key_named =
+    match named with
+    | Symbol s -> Symbol s
+    | Const c -> Const c
+    | Expr e -> Expr (make_expr_key e)
+    | Prim (prim, args, _dbg) -> Prim (prim, args)
+    | _ -> raise Not_comparable
+
+  let make_key expr =
+    match make_expr_key expr with
+    | exception Not_comparable -> None
+    | key -> Some key
+
+  let compare_key e1 e2 =
+    (* The environment [env] maps variables bound in [e2] to the corresponding
+       bound variables in [e1]. Every variable to compare in [e2] must have an
+       equivalent in [e1], otherwise the comparison wouldn't have gone
+       past the [Let] binding.  Hence [Variable.Map.find] is safe here. *)
+    let compare_var env v1 v2 =
+      match Variable.Map.find v2 env with
+      | exception Not_found ->
+        (* The variable is free in the expression [e2], hence we can
+           compare it with [v1] directly. *)
+        Variable.compare v1 v2
+      | bound ->
+        Variable.compare v1 bound
+    in
+    let rec compare_expr env (e1 : key) (e2 : key) : int =
+      match e1, e2 with
+      | Var v1, Var v2 ->
+        compare_var env v1 v2
+      | Var _, (Let _| Static_raise _) -> -1
+      | (Let _| Static_raise _), Var _ ->  1
+      | Let (v1, n1, b1), Let (v2, n2, b2) ->
+        let comp_named = compare_named env n1 n2 in
+        if comp_named <> 0 then comp_named
+        else
+          let env = Variable.Map.add v2 v1 env in
+          compare_expr env b1 b2
+      | Let _, Static_raise _ -> -1
+      | Static_raise _, Let _ ->  1
+      | Static_raise (sexn1, args1), Static_raise (sexn2, args2) ->
+        let comp_sexn = Static_exception.compare sexn1 sexn2 in
+        if comp_sexn <> 0 then comp_sexn
+        else Misc.Stdlib.List.compare (compare_var env) args1 args2
+    and compare_named env (n1:key_named) (n2:key_named) : int =
+      match n1, n2 with
+      | Symbol s1, Symbol s2 -> Symbol.compare s1 s2
+      | Symbol _, (Const _ | Expr _ | Prim _) -> -1
+      | (Const _ | Expr _ | Prim _), Symbol _ ->  1
+      | Const c1, Const c2 -> compare c1 c2
+      | Const _, (Expr _ | Prim _) -> -1
+      | (Expr _ | Prim _), Const _ ->  1
+      | Expr e1, Expr e2 -> compare_expr env e1 e2
+      | Expr _, Prim _ -> -1
+      | Prim _, Expr _ ->  1
+      | Prim (prim1, args1), Prim (prim2, args2) ->
+        let comp_prim = Pervasives.compare prim1 prim2 in
+        if comp_prim <> 0 then comp_prim
+        else Misc.Stdlib.List.compare (compare_var env) args1 args2
+    in
+    compare_expr Variable.Map.empty e1 e2
+end)
 
 let fun_vars_referenced_in_decls
       (function_decls : Flambda.function_declarations) ~backend =
diff --git a/middle_end/flambda_utils.mli b/middle_end/flambda_utils.mli
index b644bd961f..37196c06c9 100644
--- a/middle_end/flambda_utils.mli
+++ b/middle_end/flambda_utils.mli
@@ -54,10 +54,6 @@ val can_be_merged : Flambda.t -> Flambda.t -> bool
 
 val description_of_toplevel_node : Flambda.t -> string
 
-(** Sharing key, used for coalescing switch cases. *)
-type sharing_key
-val make_key : Flambda.t -> sharing_key option
-
 (* Given an expression, freshen all variables within it, and form a function
    whose body is the resulting expression.  The variables specified by
    [params] will become the parameters of the function; the closure will be