flambda-backend: Ltail for lambda and use in dissect_letrec (#1313)

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

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

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) ->