Restore Cmm unboxing behaviour inside regions (ocaml-flambda#1285)

Co-authored-by: Mark Shinwell <mshinwell@pm.me>

backend/cmm_helpers.ml

@@ -672,7 +672,7 @@ let test_bool dbg cmm =
 
 let box_float dbg m c = Cop (Calloc m, [alloc_float_header m dbg; c], dbg)
 
-let unbox_float dbg =
+let rec unbox_float dbg =
   map_tail ~kind:Vfloat (function
     | Cop (Calloc _, [Cconst_natint (hdr, _); c], _)
       when Nativeint.equal hdr float_header
@@ -682,6 +682,20 @@ let unbox_float dbg =
       match Cmmgen_state.structured_constant_of_sym s with
       | Some (Uconst_float x) -> Cconst_float (x, dbg) (* or keep _dbg? *)
       | _ -> Cop (Cload (Double, Immutable), [cmm], dbg))
+    | Cregion e as cmm -> (
+      (* It is valid to push unboxing inside a Cregion except when the extra
+         unboxing logic pushes a tail call out of tail position *)
+      match
+        map_tail ~kind:Vfloat
+          (function
+            | Cop (Capply (_, Rc_close_at_apply), _, _) -> raise Exit
+            | Ctail e -> Ctail (unbox_float dbg e)
+            | e -> unbox_float dbg e)
+          e
+      with
+      | e -> Cregion e
+      | exception Exit -> Cop (Cload (Double, Immutable), [cmm], dbg))
+    | Ctail e -> Ctail (unbox_float dbg e)
     | cmm -> Cop (Cload (Double, Immutable), [cmm], dbg))
 
 (* Complex *)
@@ -1409,7 +1423,7 @@ let alloc_matches_boxed_int bi ~hdr ~ops =
     && String.equal sym caml_int64_ops
   | (Pnativeint | Pint32 | Pint64), _, _ -> false
 
-let unbox_int dbg bi =
+let rec unbox_int dbg bi =
   let default arg =
     if size_int = 4 && bi = Primitive.Pint64
     then split_int64_for_32bit_target arg dbg
@@ -1458,6 +1472,20 @@ let unbox_int dbg bi =
             Ctuple
               [natint_const_untagged dbg low; natint_const_untagged dbg high]
       | _ -> default cmm)
+    | Cregion e as cmm -> (
+      (* It is valid to push unboxing inside a Cregion except when the extra
+         unboxing logic pushes a tail call out of tail position *)
+      match
+        map_tail ~kind:Vint
+          (function
+            | Cop (Capply (_, Rc_close_at_apply), _, _) -> raise Exit
+            | Ctail e -> Ctail (unbox_int dbg bi e)
+            | e -> unbox_int dbg bi e)
+          e
+      with
+      | e -> Cregion e
+      | exception Exit -> default cmm)
+    | Ctail e -> Ctail (unbox_int dbg bi e)
     | cmm -> default cmm)
 
 let make_unsigned_int bi arg dbg =

backend/cmmgen.ml

@@ -394,8 +394,7 @@ let rec is_unboxed_number_cmm = function
           No_unboxing
       end
     | Cexit _ | Cop (Craise _, _, _) -> No_result
-    | Cregion _ | Ctail _ -> No_unboxing
-    | Csequence (_, a)
+    | Csequence (_, a) | Cregion a | Ctail a
     | Clet (_, _, a) | Cphantom_let (_, _, a) | Clet_mut (_, _, _, a) ->
       is_unboxed_number_cmm a
     | Cconst_int _

ocaml/asmcomp/cmm_helpers.ml

@@ -598,7 +598,7 @@ let test_bool dbg cmm =
 
 let box_float dbg m c = Cop(Calloc m, [alloc_float_header m dbg; c], dbg)
 
-let unbox_float dbg =
+let rec unbox_float dbg =
   map_tail
     ~kind:Vfloat
     (function
@@ -612,6 +612,20 @@ let unbox_float dbg =
           | _ ->
               Cop(Cload (Double, Immutable), [cmm], dbg)
           end
+      | Cregion e as cmm -> (
+        (* It is valid to push unboxing inside a Cregion except when the extra
+           unboxing logic pushes a tail call out of tail position *)
+        match
+          map_tail ~kind:Vfloat
+            (function
+              | Cop (Capply (_, Rc_close_at_apply), _, _) -> raise Exit
+              | Ctail e -> Ctail (unbox_float dbg e)
+              | e -> unbox_float dbg e)
+            e
+        with
+        | e -> Cregion e
+        | exception Exit -> Cop (Cload (Double, Immutable), [cmm], dbg))
+      | Ctail e -> Ctail (unbox_float dbg e)
       | cmm -> Cop(Cload (Double, Immutable), [cmm], dbg)
     )
 
@@ -1146,7 +1160,7 @@ let alloc_matches_boxed_int bi ~hdr ~ops =
         && String.equal sym caml_int64_ops
   | (Pnativeint | Pint32 | Pint64), _, _ -> false
 
-let unbox_int dbg bi =
+let rec unbox_int dbg bi =
   let default arg =
     if size_int = 4 && bi = Primitive.Pint64 then
       split_int64_for_32bit_target arg dbg
@@ -1198,6 +1212,20 @@ let unbox_int dbg bi =
           | _ ->
               default cmm
           end
+      | Cregion e as cmm -> (
+        (* It is valid to push unboxing inside a Cregion except when the extra
+           unboxing logic pushes a tail call out of tail position *)
+        match
+          map_tail ~kind:Vint
+            (function
+              | Cop (Capply (_, Rc_close_at_apply), _, _) -> raise Exit
+              | Ctail e -> Ctail (unbox_int dbg bi e)
+              | e -> unbox_int dbg bi e)
+            e
+        with
+        | e -> Cregion e
+        | exception Exit -> default cmm)
+      | Ctail e -> Ctail (unbox_int dbg bi e)
       | cmm ->
           default cmm
     )

ocaml/asmcomp/cmmgen.ml

@@ -368,6 +368,8 @@ let is_unboxed_number_cmm ~strict cmm =
         | _ ->
             notify No_unboxing
         end
+    | Cregion e | Ctail e ->
+        aux e
     | l ->
         if not (Cmm.iter_shallow_tail aux l) then
           notify No_unboxing

ocaml/testsuite/tests/typing-local/regression_cmm_unboxing.ml

@@ -1,4 +1,6 @@
-(* TEST *)
+(* TEST
+   * native
+*)
 
 (* Regression test for a bad interaction between Cmm unboxing
    and regions present after inlining *)
@@ -21,3 +23,43 @@ let () =
   Printf.printf "%Ld %Ld\n"
     (h 42L false f)
     (h 42L true f)
+
+
+let[@inline always] f a b =
+  let _ = opaque_local (local_ ref 42) in
+  if a < b then
+    a +. (a *. b)
+  else
+    a *. b
+
+let[@inline never] g a b =
+  let n = f a b in
+  assert (n = 3.)
+
+let () =
+  let prebefore = Gc.minor_words () in
+  let before = Gc.minor_words () in
+  g 1. 2.;
+  let after = Gc.minor_words () in
+  Printf.printf "%.0f words\n"
+    (after -. before -. (before -. prebefore))
+
+
+let[@inline always] f a b =
+  let _ = opaque_local (local_ ref 42) in
+  if a < b then
+    Int64.add a (Int64.mul a b)
+  else
+    Int64.mul a b
+
+let[@inline never] g a b =
+  let n = f a b in
+  assert (n = 3L)
+
+let () =
+  let prebefore = Gc.minor_words () in
+  let before = Gc.minor_words () in
+  g 1L 2L;
+  let after = Gc.minor_words () in
+  Printf.printf "%.0f words\n"
+    (after -. before -. (before -. prebefore));

ocaml/testsuite/tests/typing-local/regression_cmm_unboxing.reference

@@ -1 +1,3 @@
 44 3
+0 words
+0 words