Ltail for lambda and use in dissect_letrec (#1313)

Author: chambart

middle_end/closure/closure.ml

@@ -987,6 +987,10 @@ let close_approx_var { fenv; cenv } id =
 let close_var env id =
   let (ulam, _app) = close_approx_var env id in ulam
 
+let compute_expr_layout kinds lambda =
+  let find_kind id = Ident.Map.find_opt id kinds in
+  compute_expr_layout find_kind lambda
+
 let rec close ({ backend; fenv; cenv ; mutable_vars; kinds; catch_env } as env) lam =
   let module B = (val backend : Backend_intf.S) in
   match lam with
@@ -1142,7 +1146,7 @@ let rec close ({ backend; fenv; cenv ; mutable_vars; kinds; catch_env } as env)
                                 _approx_res)), uargs)
         when nargs > List.length params_layout ->
           let nparams = List.length params_layout in
-          let args_kinds = List.map (Lambda.compute_expr_layout kinds) args in
+          let args_kinds = List.map (compute_expr_layout kinds) args in
           let args = List.map (fun arg -> V.create_local "arg", arg) uargs in
           (* CR mshinwell: Edit when Lapply has kinds *)
           let kinds =
@@ -1189,14 +1193,14 @@ let rec close ({ backend; fenv; cenv ; mutable_vars; kinds; catch_env } as env)
           warning_if_forced_inlined ~loc ~attribute "Unknown function";
           fail_if_probe ~probe "Unknown function";
           (Ugeneric_apply(ufunct, uargs,
-                          List.map (Lambda.compute_expr_layout kinds) args,
+                          List.map (compute_expr_layout kinds) args,
                           ap_result_layout, (pos, mode), dbg), Value_unknown)
       end
   | Lsend(kind, met, obj, args, pos, mode, loc, result_layout) ->
       let (umet, _) = close env met in
       let (uobj, _) = close env obj in
       let dbg = Debuginfo.from_location loc in
-      let args_layout = List.map (Lambda.compute_expr_layout kinds) args in
+      let args_layout = List.map (compute_expr_layout kinds) args in
       (Usend(kind, umet, uobj, close_list env args, args_layout, result_layout, (pos,mode), dbg),
        Value_unknown)
   | Llet(str, kind, id, lam, body) ->

middle_end/flambda2/from_lambda/dissect_letrec.ml

@@ -542,7 +542,7 @@ let rec prepare_letrec (recursive_set : Ident.Set.t)
       Printlambda.lambda lam
   [@@ocaml.warning "-fragile-match"]
 
-let dissect_letrec ~bindings ~body =
+let dissect_letrec ~bindings ~body ~free_vars_kind =
   let letbound = Ident.Set.of_list (List.map fst bindings) in
   let letrec =
     List.fold_right
@@ -575,19 +575,7 @@ let dissect_letrec ~bindings ~body =
         id, Lprim (Pccall desc, [size], Loc_unknown))
       letrec.blocks
   in
-  let real_body = body in
-  let bound_ids_freshening =
-    List.map (fun (bound_id, _) -> bound_id, Ident.rename bound_id) bindings
-    |> Ident.Map.of_list
-  in
-  let cont = next_raise_count () in
-  let body =
-    if not letrec.needs_region
-    then body
-    else
-      let args = List.map (fun (bound_id, _) -> Lvar bound_id) bindings in
-      Lstaticraise (cont, args)
-  in
+  let body = if not letrec.needs_region then body else Lexclave body in
   let effects_then_body = lsequence (letrec.effects, body) in
   let functions =
     match letrec.functions with
@@ -618,29 +606,32 @@ let dissect_letrec ~bindings ~body =
       with_preallocations letrec.consts
   in
   let substituted = Lambda.rename letrec.substitution with_constants in
-  let body_layout = Lambda.layout_top in
-  if not letrec.needs_region
-  then substituted
-  else
-    Lstaticcatch
-      ( Lregion (Lambda.rename bound_ids_freshening substituted, body_layout),
-        ( cont,
-          List.map
-            (fun (bound_id, _) -> bound_id, Lambda.layout_letrec)
-            bindings ),
-        real_body,
-        body_layout )
+  if letrec.needs_region
+  then
+    let body_layout =
+      let bindings =
+        Ident.Map.map (fun _ -> Lambda.layout_letrec)
+        @@ Ident.Map.of_list bindings
+      in
+      let free_vars_kind id : Lambda.layout option =
+        try Some (Ident.Map.find id bindings)
+        with Not_found -> free_vars_kind id
+      in
+      Lambda.compute_expr_layout free_vars_kind body
+    in
+    Lregion (substituted, body_layout)
+  else substituted
 
 type dissected =
   | Dissected of Lambda.lambda
   | Unchanged
 
-let dissect_letrec ~bindings ~body =
+let dissect_letrec ~bindings ~body ~free_vars_kind =
   let is_a_function = function _, Lfunction _ -> true | _, _ -> false in
   if List.for_all is_a_function bindings
   then Unchanged
   else
-    try Dissected (dissect_letrec ~bindings ~body)
+    try Dissected (dissect_letrec ~bindings ~body ~free_vars_kind)
     with Bug ->
       Misc.fatal_errorf "let-rec@.%a@." Printlambda.lambda
         (Lletrec (bindings, body))

middle_end/flambda2/from_lambda/dissect_letrec.mli

@@ -24,4 +24,7 @@ type dissected =
 (** [dissect_letrec] assumes that bindings have not been dissected yet. In
     particular, that no arguments of function call are recursive. *)
 val dissect_letrec :
-  bindings:(Ident.t * Lambda.lambda) list -> body:Lambda.lambda -> dissected
+  bindings:(Ident.t * Lambda.lambda) list ->
+  body:Lambda.lambda ->
+  free_vars_kind:(Ident.t -> Lambda.layout option) ->
+  dissected

middle_end/flambda2/from_lambda/lambda_to_flambda.ml

@@ -1242,7 +1242,13 @@ let rec cps acc env ccenv (lam : L.lambda) (k : cps_continuation)
    * in
    * cps_non_tail_simple acc env ccenv defining_expr k k_exn *)
   | Lletrec (bindings, body) -> (
-    match Dissect_letrec.dissect_letrec ~bindings ~body with
+    let free_vars_kind id =
+      let _, kind_with_subkind = CCenv.find_var ccenv id in
+      Some
+        (Flambda_kind.to_lambda
+           (Flambda_kind.With_subkind.kind kind_with_subkind))
+    in
+    match Dissect_letrec.dissect_letrec ~bindings ~body ~free_vars_kind with
     | Unchanged ->
       let function_declarations = cps_function_bindings env bindings in
       let body acc ccenv = cps acc env ccenv body k k_exn in

middle_end/flambda2/kinds/flambda_kind.ml

@@ -57,6 +57,19 @@ let region = Region
 
 let rec_info = Rec_info
 
+let to_lambda (t : t) : Lambda.layout =
+  match t with
+  | Value -> Pvalue Pgenval
+  | Naked_number Naked_immediate ->
+    Misc.fatal_error "Can't convert kind [Naked_immediate] to lambda layout"
+  | Naked_number Naked_float -> Punboxed_float
+  | Naked_number Naked_int32 -> Punboxed_int Pint32
+  | Naked_number Naked_int64 -> Punboxed_int Pint64
+  | Naked_number Naked_nativeint -> Punboxed_int Pnativeint
+  | Region -> Misc.fatal_error "Can't convert kind [Region] to lambda layout"
+  | Rec_info ->
+    Misc.fatal_error "Can't convert kind [Rec_info] to lambda layout"
+
 include Container_types.Make (struct
   type nonrec t = t
 

middle_end/flambda2/kinds/flambda_kind.mli

@@ -64,6 +64,8 @@ val is_value : t -> bool
 
 val is_naked_float : t -> bool
 
+val to_lambda : t -> Lambda.layout
+
 include Container_types.S with type t := t
 
 module Standard_int : sig

ocaml/lambda/lambda.ml

@@ -1491,37 +1491,40 @@ let primitive_result_layout (p : primitive) =
       layout_any_value
   | (Parray_to_iarray | Parray_of_iarray) -> layout_any_value
 
-let rec compute_expr_layout kinds lam =
-  match lam with
-  | Lvar id | Lmutvar id ->
-    begin
-      try Ident.Map.find id kinds
-      with Not_found ->
-        Misc.fatal_errorf "Unbound layout for variable %a" Ident.print id
-    end
-  | Lconst cst -> structured_constant_layout cst
-  | Lfunction _ -> layout_function
-  | Lapply { ap_result_layout; _ } -> ap_result_layout
-  | Lsend (_, _, _, _, _, _, _, layout) -> layout
-  | Llet(_, kind, id, _, body) | Lmutlet(kind, id, _, body) ->
-    compute_expr_layout (Ident.Map.add id kind kinds) body
-  | Lletrec(defs, body) ->
-    let kinds =
-      List.fold_left (fun kinds (id, _) -> Ident.Map.add id layout_letrec kinds)
-        kinds defs
-    in
-    compute_expr_layout kinds body
-  | Lprim(p, _, _) ->
-    primitive_result_layout p
-  | Lswitch(_, _, _, kind) | Lstringswitch(_, _, _, _, kind)
-  | Lstaticcatch(_, _, _, kind) | Ltrywith(_, _, _, kind)
-  | Lifthenelse(_, _, _, kind) | Lregion (_, kind) ->
-    kind
-  | Lstaticraise (_, _) ->
-    layout_bottom
-  | Lsequence(_, body) | Levent(body, _) -> compute_expr_layout kinds body
-  | Lwhile _ | Lfor _ | Lassign _ -> layout_unit
-  | Lifused _ ->
-      assert false
-  | Lexclave e -> compute_expr_layout kinds e
-
+let compute_expr_layout free_vars_kind lam =
+  let rec compute_expr_layout kinds = function
+    | Lvar id | Lmutvar id -> begin
+        try Ident.Map.find id kinds
+        with Not_found ->
+        match free_vars_kind id with
+        | Some kind -> kind
+        | None ->
+            Misc.fatal_errorf "Unbound layout for variable %a" Ident.print id
+      end
+    | Lconst cst -> structured_constant_layout cst
+    | Lfunction _ -> layout_function
+    | Lapply { ap_result_layout; _ } -> ap_result_layout
+    | Lsend (_, _, _, _, _, _, _, layout) -> layout
+    | Llet(_, kind, id, _, body) | Lmutlet(kind, id, _, body) ->
+      compute_expr_layout (Ident.Map.add id kind kinds) body
+    | Lletrec(defs, body) ->
+      let kinds =
+        List.fold_left (fun kinds (id, _) -> Ident.Map.add id layout_letrec kinds)
+          kinds defs
+      in
+      compute_expr_layout kinds body
+    | Lprim(p, _, _) ->
+      primitive_result_layout p
+    | Lswitch(_, _, _, kind) | Lstringswitch(_, _, _, _, kind)
+    | Lstaticcatch(_, _, _, kind) | Ltrywith(_, _, _, kind)
+    | Lifthenelse(_, _, _, kind) | Lregion (_, kind) ->
+      kind
+    | Lstaticraise (_, _) ->
+      layout_bottom
+    | Lsequence(_, body) | Levent(body, _) -> compute_expr_layout kinds body
+    | Lwhile _ | Lfor _ | Lassign _ -> layout_unit
+    | Lifused _ ->
+        assert false
+    | Lexclave e -> compute_expr_layout kinds e
+  in
+  compute_expr_layout Ident.Map.empty lam

ocaml/lambda/lambda.mli

@@ -669,4 +669,4 @@ val structured_constant_layout : structured_constant -> layout
 
 val primitive_result_layout : primitive -> layout
 
-val compute_expr_layout : layout Ident.Map.t -> lambda -> layout
+val compute_expr_layout : (Ident.t -> layout option) -> lambda -> layout

ocaml/middle_end/closure/closure.ml

@@ -987,6 +987,10 @@ let close_approx_var { fenv; cenv } id =
 let close_var env id =
   let (ulam, _app) = close_approx_var env id in ulam
 
+let compute_expr_layout kinds lambda =
+  let find_kind id = Ident.Map.find_opt id kinds in
+  compute_expr_layout find_kind lambda
+
 let rec close ({ backend; fenv; cenv ; mutable_vars; kinds; catch_env } as env) lam =
   let module B = (val backend : Backend_intf.S) in
   match lam with
@@ -1147,7 +1151,7 @@ let rec close ({ backend; fenv; cenv ; mutable_vars; kinds; catch_env } as env)
                                 _approx_res)), uargs)
         when nargs > List.length params_layout ->
           let nparams = List.length params_layout in
-          let args_kinds = List.map (Lambda.compute_expr_layout kinds) args in
+          let args_kinds = List.map (compute_expr_layout kinds) args in
           let args = List.map (fun arg -> V.create_local "arg", arg) uargs in
           (* CR mshinwell: Edit when Lapply has kinds *)
           let kinds =
@@ -1194,14 +1198,14 @@ let rec close ({ backend; fenv; cenv ; mutable_vars; kinds; catch_env } as env)
           warning_if_forced_inlined ~loc ~attribute "Unknown function";
           fail_if_probe ~probe "Unknown function";
           (Ugeneric_apply(ufunct, uargs,
-                          List.map (Lambda.compute_expr_layout kinds) args,
+                          List.map (compute_expr_layout kinds) args,
                           ap_result_layout, (pos, mode), dbg), Value_unknown)
       end
   | Lsend(kind, met, obj, args, pos, mode, loc, result_layout) ->
       let (umet, _) = close env met in
       let (uobj, _) = close env obj in
       let dbg = Debuginfo.from_location loc in
-      let args_layout = List.map (Lambda.compute_expr_layout kinds) args in
+      let args_layout = List.map (compute_expr_layout kinds) args in
       (Usend(kind, umet, uobj, close_list env args, args_layout, result_layout, (pos,mode), dbg),
        Value_unknown)
   | Llet(str, kind, id, lam, body) ->