flambda-backend: Add layout on Lregion (#1107)

bytecomp/bytegen.ml

@@ -229,7 +229,7 @@ let rec size_of_lambda env = function
   | Lprim (Pduprecord (Record_float, size), _, _) -> RHS_floatblock size
   | Levent (lam, _) -> size_of_lambda env lam
   | Lsequence (_lam, lam') -> size_of_lambda env lam'
-  | Lregion lam -> size_of_lambda env lam
+  | Lregion (lam, _) -> size_of_lambda env lam
   | _ -> RHS_nonrec
 
 (**** Merging consecutive events ****)
@@ -1019,7 +1019,7 @@ let rec comp_expr env exp sz cont =
       end
   | Lifused (_, exp) ->
       comp_expr env exp sz cont
-  | Lregion exp ->
+  | Lregion (exp, _) ->
       comp_expr env exp sz cont
 
 (* Compile a list of arguments [e1; ...; eN] to a primitive operation.

lambda/lambda.ml

@@ -445,7 +445,7 @@ type lambda =
       * region_close * alloc_mode * scoped_location * layout
   | Levent of lambda * lambda_event
   | Lifused of Ident.t * lambda
-  | Lregion of lambda
+  | Lregion of lambda * layout
 
 and lfunction =
   { kind: function_kind;
@@ -655,7 +655,7 @@ let make_key e =
     | Lsend (m,e1,e2,es,pos,mo,_loc,layout) ->
         Lsend (m,tr_rec env e1,tr_rec env e2,tr_recs env es,pos,mo,Loc_unknown,layout)
     | Lifused (id,e) -> Lifused (id,tr_rec env e)
-    | Lregion e -> Lregion (tr_rec env e)
+    | Lregion (e,layout) -> Lregion (tr_rec env e,layout)
     | Lletrec _|Lfunction _
     | Lfor _ | Lwhile _
 (* Beware: (PR#6412) the event argument to Levent
@@ -754,7 +754,7 @@ let shallow_iter ~tail ~non_tail:f = function
       tail e
   | Lifused (_v, e) ->
       tail e
-  | Lregion e ->
+  | Lregion (e, _) ->
       f e
 
 let iter_head_constructor f l =
@@ -836,7 +836,7 @@ let rec free_variables = function
   | Lifused (_v, e) ->
       (* Shouldn't v be considered a free variable ? *)
       free_variables e
-  | Lregion e ->
+  | Lregion (e, _) ->
       free_variables e
 
 and free_variables_list set exprs =
@@ -1041,8 +1041,8 @@ let subst update_env ?(freshen_bound_variables = false) s input_lam =
     | Lifused (id, e) ->
         let id = try Ident.Map.find id l with Not_found -> id in
         Lifused (id, subst s l e)
-    | Lregion e ->
-        Lregion (subst s l e)
+    | Lregion (e, layout) ->
+        Lregion (subst s l e, layout)
   and subst_list s l li = List.map (subst s l) li
   and subst_decl s l (id, exp) = (id, subst s l exp)
   and subst_case s l (key, case) = (key, subst s l case)
@@ -1140,8 +1140,8 @@ let shallow_map ~tail ~non_tail:f = function
       Levent (tail l, ev)
   | Lifused (v, e) ->
       Lifused (v, tail e)
-  | Lregion e ->
-      Lregion (f e)
+  | Lregion (e, layout) ->
+      Lregion (f e, layout)
 
 let map f =
   let rec g lam = f (shallow_map ~tail:g ~non_tail:g lam) in

lambda/lambda.mli

@@ -373,7 +373,7 @@ type lambda =
              * region_close * alloc_mode * scoped_location * layout
   | Levent of lambda * lambda_event
   | Lifused of Ident.t * lambda
-  | Lregion of lambda
+  | Lregion of lambda * layout
 
 and lfunction = private
   { kind: function_kind;

lambda/matching.ml

@@ -3662,7 +3662,7 @@ let rec map_return f = function
   | ( Lvar _ | Lmutvar _ | Lconst _ | Lapply _ | Lfunction _ | Lsend _ | Lprim _
     | Lwhile _ | Lfor _ | Lassign _ | Lifused _ ) as l ->
       f l
-  | Lregion l -> Lregion (map_return f l)
+  | Lregion (l, layout) -> Lregion (map_return f l, layout)
 
 (* The 'opt' reference indicates if the optimization is worthy.
 

lambda/printlambda.ml

@@ -808,7 +808,7 @@ let rec lam ppf = function
       end
   | Lifused(id, expr) ->
       fprintf ppf "@[<2>(ifused@ %a@ %a)@]" Ident.print id lam expr
-  | Lregion expr ->
+  | Lregion (expr, _) ->
       fprintf ppf "@[<2>(region@ %a)@]" lam expr
 
 and sequence ppf = function

lambda/simplif.ml

@@ -95,8 +95,8 @@ let rec eliminate_ref id = function
       Levent(eliminate_ref id l, ev)
   | Lifused(v, e) ->
       Lifused(v, eliminate_ref id e)
-  | Lregion e ->
-      Lregion(eliminate_ref id e)
+  | Lregion (e, layout) ->
+      Lregion(eliminate_ref id e, layout)
 
 (* Simplification of exits *)
 
@@ -184,7 +184,7 @@ let simplify_exits lam =
   | Lsend(_k, m, o, ll, _, _, _, _) -> List.iter (count ~try_depth) (m::o::ll)
   | Levent(l, _) -> count ~try_depth l
   | Lifused(_v, l) -> count ~try_depth l
-  | Lregion l -> count ~try_depth l
+  | Lregion (l, _) -> count ~try_depth l
 
   and count_default ~try_depth sw = match sw.sw_failaction with
   | None -> ()
@@ -329,7 +329,7 @@ let simplify_exits lam =
       List.map (simplif ~try_depth) ll, pos, mode, loc, layout)
   | Levent(l, ev) -> Levent(simplif ~try_depth l, ev)
   | Lifused(v, l) -> Lifused (v,simplif ~try_depth l)
-  | Lregion l -> Lregion (simplif ~try_depth l)
+  | Lregion (l, layout) -> Lregion (simplif ~try_depth l, layout)
   in
   simplif ~try_depth:0 lam
 
@@ -462,7 +462,7 @@ let simplify_lets lam =
   | Levent(l, _) -> count bv l
   | Lifused(v, l) ->
       if count_var v > 0 then count bv l
-  | Lregion l ->
+  | Lregion (l, _) ->
       count bv l
 
   and count_default bv sw = match sw.sw_failaction with
@@ -613,7 +613,7 @@ let simplify_lets lam =
   | Levent(l, ev) -> Levent(simplif l, ev)
   | Lifused(v, l) ->
       if count_var v > 0 then simplif l else lambda_unit
-  | Lregion l -> Lregion (simplif l)
+  | Lregion (l, layout) -> Lregion (simplif l, layout)
   in
   simplif lam
 
@@ -704,7 +704,7 @@ let rec emit_tail_infos is_tail lambda =
       emit_tail_infos is_tail lam
   | Lifused (_, lam) ->
       emit_tail_infos is_tail lam
-  | Lregion lam ->
+  | Lregion (lam, _) ->
       emit_tail_infos is_tail lam
 and list_emit_tail_infos_fun f is_tail =
   List.iter (fun x -> emit_tail_infos is_tail (f x))
@@ -745,7 +745,7 @@ let split_default_wrapper ~id:fun_id ~kind ~params ~return ~body
       ->
         let wrapper_body, inner = aux ((optparam, id) :: map) add_region rest in
         Llet(Strict, k, id, def, wrapper_body), inner
-    | Lregion rest -> aux map true rest
+    | Lregion (rest, _) -> aux map true rest
     | _ when map = [] -> raise Exit
     | body ->
         (* Check that those *opt* identifiers don't appear in the remaining
@@ -778,7 +778,7 @@ let split_default_wrapper ~id:fun_id ~kind ~params ~return ~body
           ) Ident.Map.empty inner_params new_ids
         in
         let body = Lambda.rename subst body in
-        let body = if add_region then Lregion body else body in
+        let body = if add_region then Lregion (body, return) else body in
         let inner_fun =
           lfunction ~kind:(Curried {nlocal=0})
             ~params:(List.map (fun id -> id, Lambda.layout_top) new_ids)
@@ -902,7 +902,7 @@ let simplify_local_functions lam =
     | Lfunction lf ->
         check_static lf;
         function_definition lf
-    | Lregion lam -> region lam
+    | Lregion (lam, _) -> region lam
     | lam ->
         Lambda.shallow_iter ~tail ~non_tail lam
   and non_tail lam =

lambda/tmc.ml

@@ -665,9 +665,9 @@ let rec choice ctx t =
     | Lifused (x, lam) ->
         let+ lam = choice ctx ~tail lam in
         Lifused (x, lam)
-    | Lregion lam ->
+    | Lregion (lam, layout) ->
         let+ lam = choice ctx ~tail lam in
-        Lregion lam
+        Lregion (lam, layout)
 
   and choice_apply ctx ~tail apply =
     let exception No_tmc in

lambda/translcore.ml

@@ -126,7 +126,7 @@ let may_allocate_in_region lam =
        | None | Some Alloc_heap ->
           List.iter loop args
        end
-    | Lregion _body ->
+    | Lregion (_body, _layout) ->
        (* [_body] might do local allocations, but not in the current region *)
        ()
     | Lwhile {wh_cond_region=false} -> raise Exit
@@ -147,7 +147,7 @@ let may_allocate_in_region lam =
     | exception Exit -> true
   end
 
-let maybe_region lam =
+let maybe_region get_layout lam =
   let rec remove_tail_markers = function
     | Lapply ({ap_region_close = Rc_close_at_apply} as ap) ->
        Lapply ({ap with ap_region_close = Rc_normal})
@@ -158,9 +158,15 @@ let maybe_region lam =
        Lambda.shallow_map ~tail:remove_tail_markers ~non_tail:Fun.id lam
   in
   if not Config.stack_allocation then lam
-  else if may_allocate_in_region lam then Lregion lam
+  else if may_allocate_in_region lam then Lregion (lam, get_layout ())
   else remove_tail_markers lam
 
+let maybe_region_layout layout lam =
+  maybe_region (fun () -> layout) lam
+
+let maybe_region_exp exp lam =
+  maybe_region (fun () -> Typeopt.layout exp.exp_env exp.exp_type) lam
+
 (* Push the default values under the functional abstractions *)
 (* Also push bindings of module patterns, since this sound *)
 
@@ -589,10 +595,15 @@ and transl_exp0 ~in_new_scope ~scopes e =
       let cond = transl_exp ~scopes wh_cond in
       let body = transl_exp ~scopes wh_body in
       Lwhile {
-        wh_cond = if wh_cond_region then maybe_region cond else cond;
+        wh_cond =
+          if wh_cond_region then
+            maybe_region_layout layout_int cond
+          else cond;
         wh_cond_region;
         wh_body = event_before ~scopes wh_body
-                    (if wh_body_region then maybe_region body else body);
+                    (if wh_body_region then
+                       maybe_region_layout layout_unit body
+                     else body);
         wh_body_region;
       }
   | Texp_arr_comprehension (body, blocks) ->
@@ -613,7 +624,9 @@ and transl_exp0 ~in_new_scope ~scopes e =
         for_to = transl_exp ~scopes for_to;
         for_dir;
         for_body = event_before ~scopes for_body
-                     (if for_region then maybe_region body else body);
+                     (if for_region then
+                        maybe_region_layout layout_unit body
+                      else body);
         for_region;
       }
   | Texp_send(expr, met, pos) ->
@@ -773,7 +786,9 @@ and transl_exp0 ~in_new_scope ~scopes e =
                             ~loc:(of_location ~scopes e.exp_loc)
                             ~mode:alloc_heap
                             ~region:true
-                            ~body:(maybe_region (transl_exp ~scopes e))
+                            ~body:(maybe_region_layout
+                                     Lambda.layout_lazy_contents
+                                     (transl_exp ~scopes e))
          in
           Lprim(Pmakeblock(Config.lazy_tag, Mutable, None, alloc_heap), [fn],
                 of_location ~scopes e.exp_loc)
@@ -1201,7 +1216,7 @@ and transl_function ~scopes e param cases partial warnings region curry =
   in
   let attr = default_function_attribute in
   let loc = of_location ~scopes e.exp_loc in
-  let body = if region then maybe_region body else body in
+  let body = if region then maybe_region_layout return body else body in
   let lam = lfunction ~kind ~params ~return ~body ~attr ~loc ~mode ~region in
   Translattribute.add_function_attributes lam e.exp_loc e.exp_attributes
 
@@ -1237,7 +1252,7 @@ and transl_let ~scopes ?(add_regions=false) ?(in_structure=false)
       | {vb_pat=pat; vb_expr=expr; vb_attributes=attr; vb_loc} :: rem ->
           let lam = transl_bound_exp ~scopes ~in_structure pat expr in
           let lam = Translattribute.add_function_attributes lam vb_loc attr in
-          let lam = if add_regions then maybe_region lam else lam in
+          let lam = if add_regions then maybe_region_exp expr lam else lam in
           let mk_body = transl rem in
           fun body ->
             Matching.for_let ~scopes pat.pat_loc lam pat body_kind (mk_body body)
@@ -1255,7 +1270,7 @@ and transl_let ~scopes ?(add_regions=false) ?(in_structure=false)
         let lam =
           Translattribute.add_function_attributes lam vb_loc vb_attributes
         in
-        let lam = if add_regions then maybe_region lam else lam in
+        let lam = if add_regions then maybe_region_exp expr lam else lam in
         begin match transl_exp_mode expr, lam with
         | Alloc_heap, _ -> ()
         | Alloc_local, Lfunction _ -> ()
@@ -1539,7 +1554,7 @@ and transl_letop ~scopes loc env let_ ands param case partial warnings =
     in
     let attr = default_function_attribute in
     let loc = of_location ~scopes case.c_rhs.exp_loc in
-    let body = maybe_region body in
+    let body = maybe_region_layout return body in
     lfunction ~kind ~params ~return ~body ~attr ~loc
               ~mode:alloc_heap ~region:true
   in
@@ -1560,17 +1575,17 @@ and transl_letop ~scopes loc env let_ ands param case partial warnings =
    that can only return global values *)
 
 let transl_exp ~scopes exp =
-  maybe_region (transl_exp ~scopes exp)
+  maybe_region_exp exp (transl_exp ~scopes exp)
 
 let transl_let ~scopes ?in_structure rec_flag pat_expr_list =
   transl_let ~scopes ~add_regions:true ?in_structure rec_flag pat_expr_list
 
 let transl_scoped_exp ~scopes exp =
-  maybe_region (transl_scoped_exp ~scopes exp)
+  maybe_region_exp exp (transl_scoped_exp ~scopes exp)
 
 let transl_apply
       ~scopes ?tailcall ?inlined ?specialised ?position ?mode fn args loc =
-  maybe_region (transl_apply
+  maybe_region_layout Lambda.layout_top (transl_apply
       ~scopes ?tailcall ?inlined ?specialised ?position ?mode fn args loc)
 
 (* Error report *)

middle_end/closure/closure.ml

@@ -1411,7 +1411,7 @@ let rec close ({ backend; fenv; cenv ; mutable_vars; kinds; catch_env } as env)
       close env lam
   | Lifused _ ->
       assert false
-  | Lregion lam ->
+  | Lregion (lam, _) ->
       let ulam, approx = close env lam in
       region ulam, approx
 

middle_end/flambda/closure_conversion.ml

@@ -571,7 +571,7 @@ let rec close t env (lam : Lambda.lambda) : Flambda.t =
        or by completely removing it (replacing by unit). *)
     Misc.fatal_error "[Lifused] should have been removed by \
         [Simplif.simplify_lets]"
-  | Lregion body ->
+  | Lregion (body, _) ->
     Region (close t env body)
 
 (** Perform closure conversion on a set of function declarations, returning a