Flambda support (#49)

* Flambda support for local allocations and regions
* Optimise away redundant Regions in Flambda

lambda/lambda.ml

@@ -1000,3 +1000,21 @@ let max_arity () =
 
 let reset () =
   raise_count := 0
+
+let join_mode a b =
+  match a, b with
+  | Alloc_local, _ | _, Alloc_local -> Alloc_local
+  | Alloc_heap, Alloc_heap -> Alloc_heap
+
+let sub_mode a b =
+  match a, b with
+  | Alloc_heap, _ -> true
+  | _, Alloc_local -> true
+  | Alloc_local, Alloc_heap -> false
+
+let eq_mode a b =
+  match a, b with
+  | Alloc_heap, Alloc_heap -> true
+  | Alloc_local, Alloc_local -> true
+  | Alloc_heap, Alloc_local -> false
+  | Alloc_local, Alloc_heap -> false

lambda/lambda.mli

@@ -451,6 +451,10 @@ val max_arity : unit -> int
       This is unlimited ([max_int]) for bytecode, but limited
       (currently to 126) for native code. *)
 
+val join_mode : alloc_mode -> alloc_mode -> alloc_mode
+val sub_mode : alloc_mode -> alloc_mode -> bool
+val eq_mode : alloc_mode -> alloc_mode -> bool
+
 (***********************)
 (* For static failures *)
 (***********************)

lambda/translcore.ml

@@ -93,17 +93,6 @@ let transl_value_mode mode =
   let alloc_mode = Types.Value_mode.regional_to_global_alloc mode in
   transl_alloc_mode alloc_mode
 
-let join_mode a b =
-  match a, b with
-  | Alloc_local, _ | _, Alloc_local -> Alloc_local
-  | Alloc_heap, Alloc_heap -> Alloc_heap
-
-let sub_mode a b =
-  match a, b with
-  | Alloc_heap, _ -> true
-  | _, Alloc_local -> true
-  | Alloc_local, Alloc_heap -> false
-
 let transl_apply_position position =
   match position with
   | Nontail -> Apply_nontail
@@ -941,10 +930,11 @@ and transl_tupled_function
       ~scopes ~arity ~mode loc return
       repr partial (param:Ident.t) cases =
   match cases with
-  | {c_lhs={pat_desc = Tpat_tuple pl}} :: _
+  | {c_lhs={pat_desc = Tpat_tuple pl; pat_mode }} :: _
     when !Clflags.native_code
       && arity = 1
       && mode = Alloc_heap
+      && transl_value_mode pat_mode = Alloc_heap
       && List.length pl <= (Lambda.max_arity ()) ->
       begin try
         let size = List.length pl in

middle_end/flambda/augment_specialised_args.ml

@@ -468,6 +468,7 @@ module Make (T : S) = struct
             spec_args_bound_in_the_wrapper;
           kind = Direct (Closure_id.wrap new_fun_var);
           dbg = Debuginfo.none;
+          position = Apply_nontail;
           inline = Default_inline;
           specialise = Default_specialise;
         }
@@ -527,9 +528,15 @@ module Make (T : S) = struct
         for_one_function.existing_specialised_args
         Variable.Map.empty
     in
+    let alloc_mode =
+      (* Wrapper closes over no more values than the original function,
+         so can share the same alloc mode *)
+      function_decl.alloc_mode
+    in
     let new_function_decl =
       Flambda.create_function_declaration
         ~params:wrapper_params
+        ~alloc_mode
         ~body:wrapper_body
         ~stub:true
         ~dbg:Debuginfo.none
@@ -607,8 +614,12 @@ module Make (T : S) = struct
           Variable.Set.elements (Variable.Map.keys
             for_one_function.new_inner_to_new_outer_vars)
         in
+        let last_mode =
+          List.fold_left (fun _mode p -> Parameter.alloc_mode p)
+            function_decl.alloc_mode function_decl.params
+        in
         let new_params =
-          List.map Parameter.wrap new_params
+          List.map (fun p -> Parameter.wrap p last_mode) new_params
         in
         function_decl.params @ new_params
       in
@@ -618,6 +629,7 @@ module Make (T : S) = struct
       let rewritten_function_decl =
         Flambda.create_function_declaration
           ~params:all_params
+          ~alloc_mode:function_decl.alloc_mode
           ~body:function_decl.body
           ~stub:function_decl.stub
           ~dbg:function_decl.dbg

middle_end/flambda/build_export_info.ml

@@ -258,6 +258,10 @@ let rec approx_of_expr (env : Env.t) (flam : Flambda.t) : Export_info.approx =
         Closure_id.Map.find closure_id results
       | _ -> Value_unknown
     end
+  | Region body ->
+    approx_of_expr env body
+  | Tail body ->
+    approx_of_expr env body
   | Assign _ -> Value_id (Env.new_unit_descr env)
   | For _ -> Value_id (Env.new_unit_descr env)
   | While _ -> Value_id (Env.new_unit_descr env)

middle_end/flambda/closure_conversion.ml

@@ -87,6 +87,7 @@ let tupled_function_call_stub original_params unboxed_version ~closure_bound_var
            redundancy here (func is also unboxed_version) *)
         kind = Direct (Closure_id.wrap unboxed_version);
         dbg = Debuginfo.none;
+        position = Apply_nontail;
         inline = Default_inline;
         specialise = Default_specialise;
       })
@@ -99,8 +100,10 @@ let tupled_function_call_stub original_params unboxed_version ~closure_bound_var
         pos + 1, Flambda.create_let param lam body)
       (0, call) params
   in
-  let tuple_param = Parameter.wrap tuple_param_var in
-  Flambda.create_function_declaration ~params:[tuple_param]
+  (* Tupled functions are always Alloc_heap. See translcore.ml *)
+  let alloc_mode = Lambda.Alloc_heap in
+  let tuple_param = Parameter.wrap tuple_param_var alloc_mode in
+  Flambda.create_function_declaration ~params:[tuple_param] ~alloc_mode
     ~body ~stub:true ~dbg:Debuginfo.none ~inline:Default_inline
     ~specialise:Default_specialise ~is_a_functor:false
     ~closure_origin:(Closure_origin.create (Closure_id.wrap closure_bound_var))
@@ -207,15 +210,15 @@ let rec close t env (lam : Lambda.lambda) : Flambda.t =
            initial_value = var;
            body;
            contents_kind = block_kind })
-  | Lfunction { kind; params; body; attr; loc; (* FIXME mode *) } ->
+  | Lfunction { kind; params; body; attr; loc; mode } ->
     let name = Names.anon_fn_with_loc loc in
     let closure_bound_var = Variable.create name in
     (* CR-soon mshinwell: some of this is now very similar to the let rec case
        below *)
     let set_of_closures_var = Variable.create Names.set_of_closures in
     let set_of_closures =
       let decl =
-        Function_decl.create ~let_rec_ident:None ~closure_bound_var ~kind
+        Function_decl.create ~let_rec_ident:None ~closure_bound_var ~kind ~mode
           ~params:(List.map fst params) ~body ~attr ~loc
       in
       close_functions t env (Function_decls.create [decl])
@@ -228,7 +231,7 @@ let rec close t env (lam : Lambda.lambda) : Flambda.t =
     Flambda.create_let set_of_closures_var set_of_closures
       (name_expr (Project_closure (project_closure)) ~name)
   | Lapply { ap_func; ap_args; ap_loc;
-             ap_tailcall = _; ap_inlined; ap_specialised; } ->
+             ap_tailcall = _; ap_inlined; ap_specialised; ap_position } ->
     Lift_code.lifting_helper (close_list t env ap_args)
       ~evaluation_order:`Right_to_left
       ~name:Names.apply_arg
@@ -241,6 +244,7 @@ let rec close t env (lam : Lambda.lambda) : Flambda.t =
               args;
               kind = Indirect;
               dbg = Debuginfo.from_location ap_loc;
+              position = ap_position;
               inline = ap_inlined;
               specialise = ap_specialised;
             })))
@@ -255,13 +259,13 @@ let rec close t env (lam : Lambda.lambda) : Flambda.t =
          will be named after the corresponding identifier in the [let rec]. *)
       List.map (function
           | (let_rec_ident,
-             Lambda.Lfunction { kind; params; body; attr; loc (* FIXME mode *) }) ->
+             Lambda.Lfunction { kind; params; body; attr; loc; mode }) ->
             let closure_bound_var =
               Variable.create_with_same_name_as_ident let_rec_ident
             in
             let function_declaration =
               Function_decl.create ~let_rec_ident:(Some let_rec_ident)
-                ~closure_bound_var ~kind ~params:(List.map fst params) ~body
+                ~closure_bound_var ~kind ~mode ~params:(List.map fst params) ~body
                 ~attr ~loc
             in
             Some function_declaration
@@ -312,7 +316,7 @@ let rec close t env (lam : Lambda.lambda) : Flambda.t =
       in
       Let_rec (defs, close t env body)
     end
-  | Lsend (kind, meth, obj, args, _FIXME, loc) ->
+  | Lsend (kind, meth, obj, args, position, loc) ->
     let meth_var = Variable.create Names.meth in
     let obj_var = Variable.create Names.obj in
     let dbg = Debuginfo.from_location loc in
@@ -322,7 +326,7 @@ let rec close t env (lam : Lambda.lambda) : Flambda.t =
           ~evaluation_order:`Right_to_left
           ~name:Names.send_arg
           ~create_body:(fun args ->
-              Send { kind; meth = meth_var; obj = obj_var; args; dbg; })))
+              Send { kind; meth = meth_var; obj = obj_var; args; dbg; position })))
   | Lprim ((Pdivint Safe | Pmodint Safe
            | Pdivbint { is_safe = Safe } | Pmodbint { is_safe = Safe }) as prim,
            [arg1; arg2], loc)
@@ -568,8 +572,8 @@ 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 _ ->
-    Misc.fatal_error "FIXME: Lregion unimplemented in Flambda"
+  | Lregion body ->
+    Region (close t env body)
 
 (** Perform closure conversion on a set of function declarations, returning a
     set of closures.  (The set will often only contain a single function;
@@ -600,22 +604,34 @@ and close_functions t external_env function_declarations : Flambda.named =
        not marked as stub but certainly should *)
     let stub = Function_decl.stub decl in
     let param_vars = List.map (Env.find_var closure_env) params in
-    let params = List.map Parameter.wrap param_vars in
+    let nheap =
+      match Function_decl.mode decl, Function_decl.kind decl with
+      | _, Curried {nlocal} -> List.length params - nlocal
+      | Alloc_heap, Tupled -> List.length params
+      | Alloc_local, Tupled -> 0
+    in
+    let params = List.mapi (fun i v ->
+      let alloc_mode : Lambda.alloc_mode =
+        if i < nheap then Alloc_heap else Alloc_local in
+      Parameter.wrap v alloc_mode) param_vars
+    in
     let closure_bound_var = Function_decl.closure_bound_var decl in
     let unboxed_version = Variable.rename closure_bound_var in
     let body = close t closure_env body in
     let closure_origin =
       Closure_origin.create (Closure_id.wrap unboxed_version)
     in
     let fun_decl =
-      Flambda.create_function_declaration ~params ~body ~stub ~dbg
+      Flambda.create_function_declaration
+        ~params ~alloc_mode:(Function_decl.mode decl)
+        ~body ~stub ~dbg
         ~inline:(Function_decl.inline decl)
         ~specialise:(Function_decl.specialise decl)
         ~is_a_functor:(Function_decl.is_a_functor decl)
         ~closure_origin
     in
     match Function_decl.kind decl with
-    | Curried _ (* FIXME nlocal *) ->
+    | Curried _ ->
        Variable.Map.add closure_bound_var fun_decl map
     | Tupled ->
       let unboxed_version = Variable.rename closure_bound_var in
@@ -661,12 +677,12 @@ and close_list t sb l = List.map (close t sb) l
 and close_let_bound_expression t ?let_rec_ident let_bound_var env
       (lam : Lambda.lambda) : Flambda.named =
   match lam with
-  | Lfunction { kind; params; body; attr; loc; } ->
+  | Lfunction { kind; params; body; attr; loc; mode } ->
     (* Ensure that [let] and [let rec]-bound functions have appropriate
        names. *)
     let closure_bound_var = Variable.rename let_bound_var in
     let decl =
-      Function_decl.create ~let_rec_ident ~closure_bound_var ~kind
+      Function_decl.create ~let_rec_ident ~closure_bound_var ~kind ~mode
         ~params:(List.map fst params) ~body ~attr ~loc
     in
     let set_of_closures_var = Variable.rename let_bound_var in

middle_end/flambda/closure_conversion_aux.ml

@@ -87,14 +87,15 @@ module Function_decls = struct
       let_rec_ident : Ident.t;
       closure_bound_var : Variable.t;
       kind : Lambda.function_kind;
+      mode : Lambda.alloc_mode;
       params : Ident.t list;
       body : Lambda.lambda;
       free_idents_of_body : Ident.Set.t;
       attr : Lambda.function_attribute;
       loc : Lambda.scoped_location
     }
 
-    let create ~let_rec_ident ~closure_bound_var ~kind ~params ~body
+    let create ~let_rec_ident ~closure_bound_var ~kind ~mode ~params ~body
         ~attr ~loc =
       let let_rec_ident =
         match let_rec_ident with
@@ -104,6 +105,7 @@ module Function_decls = struct
       { let_rec_ident;
         closure_bound_var;
         kind;
+        mode;
         params;
         body;
         free_idents_of_body = Lambda.free_variables body;
@@ -114,6 +116,7 @@ module Function_decls = struct
     let let_rec_ident t = t.let_rec_ident
     let closure_bound_var t = t.closure_bound_var
     let kind t = t.kind
+    let mode t = t.mode
     let params t = t.params
     let body t = t.body
     let free_idents t = t.free_idents_of_body
@@ -160,7 +163,7 @@ module Function_decls = struct
         (all_params function_decls))
       (let_rec_idents function_decls)
 
-  let create function_decls =
+  let create (function_decls : Function_decl.t list) =
     { function_decls;
       all_free_idents = all_free_idents function_decls;
     }

middle_end/flambda/closure_conversion_aux.mli

@@ -56,6 +56,7 @@ module Function_decls : sig
        : let_rec_ident:Ident.t option
       -> closure_bound_var:Variable.t
       -> kind:Lambda.function_kind
+      -> mode:Lambda.alloc_mode
       -> params:Ident.t list
       -> body:Lambda.lambda
       -> attr:Lambda.function_attribute
@@ -65,6 +66,7 @@ module Function_decls : sig
     val let_rec_ident : t -> Ident.t
     val closure_bound_var : t -> Variable.t
     val kind : t -> Lambda.function_kind
+    val mode : t -> Lambda.alloc_mode
     val params : t -> Ident.t list
     val body : t -> Lambda.lambda
     val inline : t -> Lambda.inline_attribute

middle_end/flambda/effect_analysis.ml

@@ -45,6 +45,10 @@ let rec no_effects (flam : Flambda.t) =
     (* If there is a [raise] in [body], the whole [Try_with] may have an
        effect, so there is no need to test the handler. *)
     no_effects body
+  | Region body ->
+    no_effects body
+  | Tail body ->
+    no_effects body
   | While _ | For _ | Apply _ | Send _ | Assign _ | Static_raise _ -> false
   | Proved_unreachable -> true
 

middle_end/flambda/export_info_for_pack.ml

@@ -139,7 +139,8 @@ and import_function_declarations_for_pack_aux units pack
   let funs =
     Variable.Map.map
       (fun (function_decl : Flambda.function_declaration) ->
-        Flambda.create_function_declaration ~params:function_decl.params
+        Flambda.create_function_declaration
+          ~params:function_decl.params ~alloc_mode:function_decl.alloc_mode
           ~body:(import_code_for_pack units pack function_decl.body)
           ~stub:function_decl.stub ~dbg:function_decl.dbg
           ~inline:function_decl.inline

middle_end/flambda/extract_projections.ml

@@ -107,6 +107,7 @@ let rec analyse_expr ~which_variables expr =
       check_free_variable from_value;
       check_free_variable to_value
     | Let _ | Let_rec _ | Static_catch _ | While _ | Try_with _
+    | Region _ | Tail _
     | Proved_unreachable -> ()
   in
   let for_named (named : Flambda.named) =