Remove regions with exclaves

middle_end/flambda/flambda_iterators.ml

@@ -76,13 +76,15 @@ let map_snd_sharing f ((a, b) as cpl) =
   else
     (a, new_b)
 
-let map_subexpressions f f_named (tree:Flambda.t) : Flambda.t =
+let map_subexpressions_with_tail f f_named (tree:Flambda.t) : Flambda.t =
+  let f_tail v = f v ~tail:true in
+  let f_nontail v = f v ~tail:false in
   match tree with
   | Var _ | Apply _ | Assign _ | Send _ | Proved_unreachable
   | Static_raise _ -> tree
   | Let { var; defining_expr; body; _ } ->
     let new_named = f_named var defining_expr in
-    let new_body = f body in
+    let new_body = f_tail body in
     if new_named == defining_expr && new_body == body then
       tree
     else
@@ -91,22 +93,23 @@ let map_subexpressions f f_named (tree:Flambda.t) : Flambda.t =
     let new_defs =
       list_map_sharing (map_snd_sharing f_named) defs
     in
-    let new_body = f body in
+    let new_body = f_tail body in
     if new_defs == defs && new_body == body then
       tree
     else
       Let_rec (new_defs, new_body)
   | Let_mutable mutable_let ->
-    let new_body = f mutable_let.body in
+    let new_body = f_tail mutable_let.body in
     if new_body == mutable_let.body then
       tree
     else
       Let_mutable { mutable_let with body = new_body }
   | Switch (arg, sw) ->
-    let aux = map_snd_sharing (fun _ v -> f v) in
+    let aux = map_snd_sharing (fun _ v -> f_tail v) in
     let new_consts = list_map_sharing aux sw.consts in
     let new_blocks = list_map_sharing aux sw.blocks in
-    let new_failaction = may_map_sharing f sw.failaction in
+    let new_failaction = may_map_sharing f_tail sw.failaction
+    in
     if sw.failaction == new_failaction &&
        new_consts == sw.consts &&
        new_blocks == sw.blocks then
@@ -121,59 +124,66 @@ let map_subexpressions f f_named (tree:Flambda.t) : Flambda.t =
       in
       Switch (arg, sw)
   | String_switch (arg, sw, def, kind) ->
-    let new_sw = list_map_sharing (map_snd_sharing (fun _ v -> f v)) sw in
-    let new_def = may_map_sharing f def in
+    let new_sw = list_map_sharing (map_snd_sharing (fun _ v -> f_tail v)) sw in
+    let new_def = may_map_sharing f_tail def in
     if sw == new_sw && def == new_def then
       tree
     else
       String_switch(arg, new_sw, new_def, kind)
   | Static_catch (i, vars, body, handler, kind) ->
-    let new_body = f body in
-    let new_handler = f handler in
+    let new_body = f_tail body in
+    let new_handler = f_tail handler in
     if new_body == body && new_handler == handler then
       tree
     else
       Static_catch (i, vars, new_body, new_handler, kind)
   | Try_with(body, id, handler, kind) ->
-    let new_body = f body in
-    let new_handler = f handler in
+    let new_body = f_tail body in
+    let new_handler = f_tail handler in
     if body == new_body && handler == new_handler then
       tree
     else
       Try_with(new_body, id, new_handler, kind)
   | If_then_else(arg, ifso, ifnot, kind) ->
-    let new_ifso = f ifso in
-    let new_ifnot = f ifnot in
+    let new_ifso = f_tail ifso in
+    let new_ifnot = f_tail ifnot in
     if new_ifso == ifso && new_ifnot == ifnot then
       tree
     else
       If_then_else(arg, new_ifso, new_ifnot, kind)
   | While(cond, body) ->
-    let new_cond = f cond in
-    let new_body = f body in
+    let new_cond = f_nontail cond in
+    let new_body = f_nontail body in
     if new_cond == cond && new_body == body then
       tree
     else
       While(new_cond, new_body)
   | For { bound_var; from_value; to_value; direction; body; } ->
-    let new_body = f body in
+    let new_body = f_nontail body in
     if new_body == body then
       tree
     else
       For { bound_var; from_value; to_value; direction; body = new_body; }
   | Region body ->
-    let new_body = f body in
+    let new_body = f_tail body in
     if new_body == body then
       tree
     else
       Region new_body
   | Exclave body ->
-    let new_body = f body in
+    let new_body = f_tail body in
     if new_body == body then
       tree
     else
       Exclave new_body
 
+let map_subexpressions f f_named tree =
+  map_subexpressions_with_tail (fun v ~tail:_ -> f v) f_named tree
+
+let map_tail_subexpressions f tree =
+  let f v ~tail = if tail then f v else v in
+  map_subexpressions_with_tail f (fun _ named -> named) tree
+
 let iter_general = Flambda.iter_general
 
 let iter f f_named t = iter_general ~toplevel:false f f_named (Is_expr t)

middle_end/flambda/flambda_iterators.mli

@@ -35,6 +35,8 @@ val map_subexpressions
   -> Flambda.t
   -> Flambda.t
 
+val map_tail_subexpressions : (Flambda.t -> Flambda.t) -> Flambda.t -> Flambda.t
+
 (* CR-soon lwhite: add comment to clarify that these recurse unlike the
    ones above *)
 val iter

middle_end/flambda/inline_and_simplify.ml

@@ -394,6 +394,29 @@ let simplify_move_within_set_of_closures env r
                 let approx = A.value_closure value_set_of_closures move_to in
                 Move_within_set_of_closures move_within, ret r approx)
 
+let remove_exclaves (lam : Flambda.t) =
+  let rec remove lam ~depth : Flambda.t =
+    match (lam : Flambda.t) with
+    | Region body ->
+      let new_body = remove ~depth:(depth + 1) body in
+      if new_body == body then lam else Region new_body
+    | Exclave body ->
+      if depth = 0 then body
+      else
+        let new_body = remove ~depth:(depth - 1) body in
+        if new_body == body then lam else Exclave new_body
+    | Apply ({ reg_close = Rc_close_at_apply; _ } as apply) when depth = 0 ->
+      (* Can still be compiled as a tail call, so use [Rc_normal] rather than
+         [Rc_nontail] *)
+      Apply { apply with reg_close = Rc_normal }
+    | Send ({ reg_close = Rc_close_at_apply; _ } as send) when depth = 0 ->
+      (* Similar to [Apply] *)
+      Send { send with reg_close = Rc_normal }
+    | _ ->
+      Flambda_iterators.map_tail_subexpressions (remove ~depth) lam
+  in
+  remove lam ~depth:0
+
 (* Transform an expression denoting an access to a variable bound in
    a closure.  Variables in the closure ([project_var.closure]) may
    have been freshened since [expr] was constructed; as such, we
@@ -675,10 +698,18 @@ and simplify_set_of_closures original_env r
   set_of_closures, r, value_set_of_closures.freshening
 
 and mark_region_used_for_apply ~(reg_close : Lambda.region_close) ~(mode : Lambda.alloc_mode) r =
-  match reg_close, mode with
-  | (Rc_normal | Rc_nontail), Alloc_heap -> r
-  | Rc_close_at_apply, _
-  | _, Alloc_local -> R.set_region_used r
+  let r =
+    (* A close-at-apply tail call is effectively a small exclave *)
+    match reg_close with
+    | Rc_close_at_apply -> R.set_region_has_exclave r
+    | Rc_normal | Rc_nontail -> r
+  in
+  let r =
+    match mode with
+    | Alloc_local -> R.set_region_used r
+    | Alloc_heap -> r
+  in
+  r
 
 and simplify_apply env r ~(apply : Flambda.apply) : Flambda.t * R.t =
   let {
@@ -1460,11 +1491,11 @@ and simplify env r (tree : Flambda.t) : Flambda.t * R.t =
      let r = R.enter_region r in
      let body, r = simplify env r body in
      let use_inner_region = R.may_use_region r in
+     let has_exclave = R.region_has_exclave r in
      let r = R.leave_region r in
      if use_inner_region then Region body, r
-     else body, r
+     else if has_exclave then remove_exclaves body, r else body, r
   | Exclave body ->
-     let r = R.set_region_used r in
      let exclave, r = R.enter_exclave r in
      let body, r = simplify env r body in
      let r = R.leave_exclave r exclave in

middle_end/flambda/inline_and_simplify_aux.ml

@@ -437,7 +437,7 @@ let initial_inlining_toplevel_threshold ~round : Inlining_cost.Threshold.t =
     (unscaled * Inlining_cost.scale_inline_threshold_by)
 
 module Result = struct
-  type region = { may_be_used : bool; }
+  type region = { may_be_used : bool; has_exclave : bool }
 
   type t =
     { approx : Simple_value_approx.t;
@@ -448,7 +448,7 @@ module Result = struct
       regions : region list;
     }
 
-  let create_region () = { may_be_used = false; }
+  let create_region () = { may_be_used = false; has_exclave = false }
 
   let create () =
     { approx = Simple_value_approx.value_unknown Other;
@@ -490,17 +490,28 @@ module Result = struct
 
   let set_region_used t =
     match t.regions with
-    | _ :: regions ->
-        { t with regions = { may_be_used = true } :: regions }
-    | [] -> t
+    | region :: regions ->
+      { t with regions = { region with may_be_used = true } :: regions }
+    | [] ->
+      (* By rights this should be a fatal error, but currently
+         [Semantics_of_primitives.may_locally_allocate] has too many false
+         positives (including ccalls). *)
+      t
+
+  let set_region_has_exclave t =
+    match t.regions with
+    | region :: regions ->
+      { t with regions = { region with has_exclave = true } :: regions }
+    | [] -> no_current_region ()
 
   type exclave = { from_region : region }
 
   let enter_exclave t =
     match t.regions with
     | region :: regions ->
-        let exclave = { from_region = region } in
-        exclave, { t with regions = regions }
+      let region = { region with has_exclave = true } in
+      let exclave = { from_region = region } in
+      exclave, { t with regions = regions }
     | [] -> no_current_region ()
 
   let leave_exclave t { from_region } =
@@ -511,7 +522,11 @@ module Result = struct
     | region :: _ -> region
     | [] -> no_current_region ()
 
-  let may_use_region t = (current_region t).may_be_used
+  let may_use_region t =
+    (current_region t).may_be_used
+
+  let region_has_exclave t =
+    (current_region t).has_exclave
 
   let exit_scope_catch t i =
     { t with

middle_end/flambda/inline_and_simplify_aux.mli

@@ -321,6 +321,13 @@ module Result : sig
   (** Whether [set_region_used _] has been called *)
   val may_use_region : t -> bool
 
+  (** Mark that the nearest enclosing region has an exclave (either an actual
+      [exclave] expression or a close-on-apply tail call). *)
+  val set_region_has_exclave : t -> t
+
+  (** Whether [enter_exclave _] or [set_region_has_exclave] has been called. *)
+  val region_has_exclave : t -> bool
+
   (** Mark that we are moving up out of the scope of a static-catch block
       that catches the given static exception identifier.  This has the effect
       of removing the identifier from the [used_staticfail] set. *)