Implement let mutable

chamelon/compat.ox.ml

@@ -307,8 +307,7 @@ let mkTpat_alias ~id:(mode, ty) (p, ident, name) =
 type tpat_array_identifier = mutability * Jkind.sort
 
 let mkTpat_array
-    ?id:(mut, arg_sort =
-        (Mutable Alloc.Comonadic.Const.legacy, Jkind.Sort.value)) l =
+    ?id:(mut, arg_sort = (Mutable Value.Comonadic.legacy, Jkind.Sort.value)) l =
   Tpat_array (mut, arg_sort, l)
 
 type tpat_tuple_identifier = string option list

file_formats/cmt_format.ml

@@ -238,15 +238,17 @@ let iter_on_occurrences
             modifs
       | Texp_extension_constructor (lid, path) ->
           f ~namespace:Extension_constructor exp_env path lid
-      | Texp_constant _ | Texp_let _ | Texp_function _ | Texp_apply _
-      | Texp_match _ | Texp_try _ | Texp_tuple _ | Texp_unboxed_tuple _
-      | Texp_variant _ | Texp_array _
+      | Texp_constant _ | Texp_let _ | Texp_letmutable _ | Texp_function _
+      | Texp_apply _ | Texp_match _ | Texp_try _ | Texp_tuple _
+      | Texp_unboxed_tuple _ | Texp_variant _ | Texp_array _
       | Texp_ifthenelse _ | Texp_sequence _ | Texp_while _ | Texp_for _
       | Texp_send _
       | Texp_letmodule _ | Texp_letexception _ | Texp_assert _ | Texp_lazy _
       | Texp_object _ | Texp_pack _ | Texp_letop _ | Texp_unreachable
       | Texp_list_comprehension _ | Texp_array_comprehension _ | Texp_probe _
       | Texp_probe_is_enabled _ | Texp_exclave _
+      (* CR-someday let_mutable: maybe iterate on mutvar? *)
+      | Texp_mutvar _ | Texp_setmutvar _
       | Texp_open _ | Texp_src_pos | Texp_overwrite _ | Texp_hole _ -> ());
       default_iterator.expr sub e);
 

jane/doc/extensions/_11-miscellaneous-extensions/let-mutable.md

@@ -0,0 +1,56 @@
+# The `let mutable` extension
+
+The `let mutable` extension provides a mechanism for creating mutable variables.
+This codifies a pre-existing optimization, where the compiler attempts to
+eliminate allocating a box for a `ref` when it can see the `ref` is only used
+locally in a given scope, instead simply storing the value in a register.
+
+```ocaml
+let triangle n =
+  let mutable total = 0 in
+  for i = 1 to n do
+    total <- total + i
+  done;
+  total
+```
+
+Mutable variables must not escape their scope. For example, a closure can't
+capture a mutable variable.
+
+Local data can be stored in a mutable variable. For example:
+
+```ocaml
+let rec sum (to_sum @ local) =
+  match to_sum with
+  | [] -> 0
+  | hd :: tl -> hd + sum tl
+
+let triangle_list n =
+  let mutable to_sum = [] in
+  for i = 1 to n do exclave_
+    to_sum <- stack_ (i :: to_sum)
+  done;
+  sum to_sum [@nontail]
+```
+
+
+## Restrictions
+
+Mutable `let` declarations may not be recursive, and they may not be used at the
+structure level or in class definitions. The pattern of a mutable `let`
+statement must be a single variable, possibly with a type annotation, e.g. `let
+mutable x, y = ..` and `let mutable add x y = ..` are not allowed. Mutable `let`
+statements must also not use `and`s.
+
+Because closures may not capture mutable variables, some uses that are
+apparently safe from a scope perspective are not possible. For example, the
+following program is rejected:
+
+```ocaml
+let sum xs =
+  let mutable total = 0 in
+  List.iter xs ~f:(fun x -> total <- total + x);
+  total
+```
+
+Mutable variables may not contain unboxed products.

lambda/matching.ml

@@ -4271,7 +4271,7 @@ let assign_pat ~scopes body_layout opt nraise catch_ids loc pat pat_sort lam =
   in
   List.fold_left push_sublet exit rev_sublets
 
-let for_let ~scopes ~arg_sort ~return_layout loc param pat body =
+let for_let ~scopes ~arg_sort ~return_layout loc param mutable_flag pat body =
   match pat.pat_desc with
   | Tpat_any ->
       (* This eliminates a useless variable (and stack slot in bytecode)
@@ -4287,7 +4287,10 @@ let for_let ~scopes ~arg_sort ~return_layout loc param pat body =
          non-polymorphic Ppat_constraint case in type_pat_aux.
       *)
       let k = Typeopt.layout pat.pat_env pat.pat_loc arg_sort pat.pat_type in
-      Llet (Strict, k, id, duid, param, body)
+      begin match mutable_flag with
+      | Asttypes.Mutable -> Lmutlet (k, id, duid, param, body)
+      | Asttypes.Immutable -> Llet (Strict, k, id, duid, param, body)
+      end
   | _ ->
       let opt = ref false in
       let nraise = next_raise_count () in
@@ -4511,7 +4514,7 @@ let for_optional_arg_default
             Loc_unknown))
   in
   for_let ~scopes ~arg_sort:default_arg_sort ~return_layout
-    loc supplied_or_default pat body
+    loc supplied_or_default Immutable pat body
 
 (* Error report *)
 (* CR layouts v5: This file didn't use to have the report_error infrastructure -

lambda/matching.mli

@@ -32,7 +32,7 @@ val for_trywith:
         lambda
 val for_let:
         scopes:scopes -> arg_sort:Jkind.Sort.Const.t -> return_layout:layout ->
-        Location.t -> lambda -> pattern -> lambda ->
+        Location.t -> lambda -> Asttypes.mutable_flag -> pattern -> lambda ->
         lambda
 val for_multiple_match:
         scopes:scopes -> return_layout:layout -> Location.t ->

lambda/transl_array_comprehension.ml

@@ -515,7 +515,7 @@ let iterator ~transl_exp ~scopes ~loc :
                        iter_arr_mut ),
                    [iter_arr.var; Lvar iter_ix],
                    loc ))
-              pattern body
+              Immutable pattern body
         }
     in
     mk_iterator, Array { iter_arr; iter_len }

lambda/transl_list_comprehension.ml

@@ -206,7 +206,8 @@ let iterator ~transl_exp ~scopes = function
       add_bindings =
         (* CR layouts: to change when we allow non-values in sequences *)
         Matching.for_let ~scopes ~arg_sort:Jkind.Sort.Const.for_list_element
-          ~return_layout:layout_any_value pattern.pat_loc (Lvar element) pattern
+          ~return_layout:layout_any_value pattern.pat_loc (Lvar element)
+          Immutable pattern
     }
 
 (** Translates a list comprehension binding

lambda/translcore.ml

@@ -36,6 +36,7 @@ type error =
   | Void_sort of type_expr
   | Unboxed_vector_in_array_comprehension
   | Unboxed_product_in_array_comprehension
+  | Unboxed_product_in_let_mutable
 
 exception Error of Location.t * error
 
@@ -393,6 +394,10 @@ and transl_exp0 ~in_new_scope ~scopes sort e =
       let return_layout = layout_exp sort body in
       transl_let ~scopes ~return_layout rec_flag pat_expr_list
         (event_before ~scopes body (transl_exp ~scopes sort body))
+  | Texp_letmutable(pat_expr, body) ->
+      let return_layout = layout_exp sort body in
+      transl_letmutable ~scopes ~return_layout pat_expr
+        (event_before ~scopes body (transl_exp ~scopes sort body))
   | Texp_function { params; body; ret_sort; ret_mode; alloc_mode;
                     zero_alloc } ->
       let ret_sort = Jkind.Sort.default_for_transl_and_get ret_sort in
@@ -951,11 +956,15 @@ and transl_exp0 ~in_new_scope ~scopes sort e =
       let self = transl_value_path loc e.exp_env path_self in
       let var = transl_value_path loc e.exp_env path in
       Lprim(Pfield_computed Reads_vary, [self; var], loc)
+  | Texp_mutvar id -> Lmutvar id.txt
   | Texp_setinstvar(path_self, path, _, expr) ->
       let loc = of_location ~scopes e.exp_loc in
       let self = transl_value_path loc e.exp_env path_self in
       let var = transl_value_path loc e.exp_env path in
       transl_setinstvar ~scopes loc self var expr
+  | Texp_setmutvar(id, expr_sort, expr) ->
+      Lassign(id.txt, transl_exp ~scopes
+        (Jkind.Sort.default_for_transl_and_get expr_sort) expr)
   | Texp_override(path_self, modifs) ->
       let loc = of_location ~scopes e.exp_loc in
       let self = transl_value_path loc e.exp_env path_self in
@@ -1878,7 +1887,7 @@ and transl_let ~scopes ~return_layout ?(add_regions=false) ?(in_structure=false)
           let mk_body = transl rem in
           fun body ->
             Matching.for_let ~scopes ~arg_sort:sort ~return_layout pat.pat_loc
-              lam pat (mk_body body)
+              lam Immutable pat (mk_body body)
       in
       transl pat_expr_list
   | Recursive ->
@@ -1902,6 +1911,15 @@ and transl_let ~scopes ~return_layout ?(add_regions=false) ?(in_structure=false)
       let lam_bds = List.map2 transl_case pat_expr_list idlist in
       fun body -> Value_rec_compiler.compile_letrec lam_bds body
 
+and transl_letmutable ~scopes ~return_layout
+      {vb_pat=pat; vb_expr=expr; vb_attributes=attr; vb_loc; vb_sort} body =
+  let arg_sort = Jkind_types.Sort.default_to_value_and_get vb_sort in
+  let lam =
+    transl_bound_exp ~scopes ~in_structure:false pat arg_sort expr vb_loc attr
+  in
+  Matching.for_let ~scopes ~return_layout ~arg_sort pat.pat_loc lam Mutable
+    pat body
+
 and transl_setinstvar ~scopes loc self var expr =
   let ptr_or_imm, _ = maybe_pointer expr in
   Lprim(Psetfield_computed (ptr_or_imm, Assignment modify_heap),
@@ -2495,6 +2513,9 @@ let report_error ppf = function
       fprintf ppf
         "Array comprehensions are not yet supported for arrays of unboxed \
          products."
+  | Unboxed_product_in_let_mutable ->
+      fprintf ppf
+        "Mutable lets are not yet supported with unboxed products."
 
 let () =
   Location.register_error_of_exn

lambda/translcore.mli

@@ -54,6 +54,7 @@ type error =
   | Void_sort of Types.type_expr
   | Unboxed_vector_in_array_comprehension
   | Unboxed_product_in_array_comprehension
+  | Unboxed_product_in_let_mutable
 
 exception Error of Location.t * error
 

middle_end/flambda2/from_lambda/lambda_to_flambda.ml

@@ -456,11 +456,13 @@ let rec cps acc env ccenv (lam : L.lambda) (k : cps_continuation)
       let fields = List.map (fun id -> IR.Var id) fields in
       apply_cps_cont_simple k acc env ccenv fields before_unarization)
   | Lmutvar id ->
-    (* CR mshinwell: note: mutable variables of non-singleton layouts are not
-       supported *)
-    let return_id, kind = Env.get_mutable_variable_with_kind env id in
-    apply_cps_cont k acc env ccenv return_id
-      (Flambda_arity.Component_for_creation.Singleton kind)
+    let new_ids_with_kinds, before_unarization =
+      Env.get_mutable_variable_with_kinds env id
+    in
+    let fields =
+      List.map (fun id -> IR.Var id) (List.map fst new_ids_with_kinds)
+    in
+    apply_cps_cont_simple k acc env ccenv fields before_unarization
   | Lconst const ->
     apply_cps_cont_simple k acc env ccenv [IR.Const const]
       (Singleton
@@ -497,24 +499,32 @@ let rec cps acc env ccenv (lam : L.lambda) (k : cps_continuation)
     CC.close_let_rec acc ccenv ~function_declarations:[func] ~body
       ~current_region:
         (Env.current_region env |> Option.map Env.Region_stack_element.region)
-  | Lmutlet (value_kind, id, _duid, defining_expr, body) ->
+  | Lmutlet (layout, id, _duid, defining_expr, body) ->
     (* CR sspies: dropping [debug_uid]; address in subsequent PR. *)
     (* CR mshinwell: user-visibleness needs thinking about here *)
     let temp_id = Ident.create_local "let_mutable" in
     let_cont_nonrecursive_with_extra_params acc env ccenv ~is_exn_handler:false
-      ~params:[temp_id, IR.Not_user_visible, value_kind]
+      ~params:[temp_id, IR.Not_user_visible, layout]
       ~body:(fun acc env ccenv after_defining_expr ->
         cps_tail acc env ccenv defining_expr after_defining_expr k_exn)
       ~handler:(fun acc env ccenv ->
-        let kind =
-          Flambda_kind.With_subkind.from_lambda_values_and_unboxed_numbers_only
-            value_kind
+        let before_unarization =
+          Flambda_arity.Component_for_creation.from_lambda layout
+        in
+        let env, new_ids_with_kinds =
+          Env.register_mutable_variable env id ~before_unarization
         in
-        let env, new_id = Env.register_mutable_variable env id kind in
         let body acc ccenv = cps acc env ccenv body k k_exn in
-        CC.close_let acc ccenv
-          [new_id, kind]
-          User_visible (Simple (Var temp_id)) ~body)
+        let temp_id_unarized : Ident.t list =
+          match Env.get_unboxed_product_fields env temp_id with
+          | None -> [temp_id]
+          | Some (_, temp_id_unarized) -> temp_id_unarized
+        in
+        List.fold_left2
+          (fun body new_id_with_kind temp_id acc ccenv ->
+            CC.close_let acc ccenv [new_id_with_kind] User_visible
+              (Simple (Var temp_id)) ~body)
+          body new_ids_with_kinds temp_id_unarized acc ccenv)
   | Llet ((Strict | Alias | StrictOpt), _, fun_id, duid, Lfunction func, body)
     ->
     (* This case is here to get function names right. *)
@@ -623,21 +633,22 @@ let rec cps acc env ccenv (lam : L.lambda) (k : cps_continuation)
       Misc.fatal_errorf "Lassign on non-mutable variable %a" Ident.print
         being_assigned;
     cps_non_tail_simple acc env ccenv new_value
-      (fun acc env ccenv new_value _arity ->
-        let new_value = must_be_singleton_simple new_value in
-        let env, new_id = Env.update_mutable_variable env being_assigned in
+      (fun acc env ccenv new_values _arity ->
+        let env = Env.update_mutable_variable env being_assigned in
         let body acc ccenv =
           let body acc ccenv = cps acc env ccenv body k k_exn in
           CC.close_let acc ccenv
             [id, Flambda_kind.With_subkind.tagged_immediate]
             Not_user_visible (Simple (Const L.const_unit)) ~body
         in
-        let value_kind =
-          snd (Env.get_mutable_variable_with_kind env being_assigned)
+        let new_ids_with_kinds, _before_unarization =
+          Env.get_mutable_variable_with_kinds env being_assigned
         in
-        CC.close_let acc ccenv
-          [new_id, value_kind]
-          User_visible (Simple new_value) ~body)
+        List.fold_left2
+          (fun body new_id_with_kind new_value acc ccenv ->
+            CC.close_let acc ccenv [new_id_with_kind] User_visible
+              (Simple new_value) ~body)
+          body new_ids_with_kinds new_values acc ccenv)
       k_exn
   | Llet
       ((Strict | Alias | StrictOpt), _layout, id, _duid, defining_expr, Lvar id')
@@ -966,19 +977,20 @@ let rec cps acc env ccenv (lam : L.lambda) (k : cps_continuation)
       Misc.fatal_errorf "Lassign on non-mutable variable %a" Ident.print
         being_assigned;
     cps_non_tail_simple acc env ccenv new_value
-      (fun acc env ccenv new_value _arity ->
-        let new_value = must_be_singleton_simple new_value in
-        let env, new_id = Env.update_mutable_variable env being_assigned in
+      (fun acc env ccenv new_values _arity ->
+        let env = Env.update_mutable_variable env being_assigned in
         let body acc ccenv =
           apply_cps_cont_simple k acc env ccenv [Const L.const_unit]
             (Singleton Flambda_kind.With_subkind.tagged_immediate)
         in
-        let _, value_kind =
-          Env.get_mutable_variable_with_kind env being_assigned
+        let new_ids_with_kinds, _before_unarization =
+          Env.get_mutable_variable_with_kinds env being_assigned
         in
-        CC.close_let acc ccenv
-          [new_id, value_kind]
-          User_visible (Simple new_value) ~body)
+        List.fold_left2
+          (fun body new_id_with_kind new_value acc ccenv ->
+            CC.close_let acc ccenv [new_id_with_kind] User_visible
+              (Simple new_value) ~body)
+          body new_ids_with_kinds new_values acc ccenv)
       k_exn
   | Levent (body, _event) -> cps acc env ccenv body k k_exn
   | Lifused _ ->