flambda-backend: Local immutable arrays (#1420)

Co-authored-by: Mark Shinwell <[email protected]>
ocaml-flambda · Jun 14, 2023 · be57ed6 · be57ed6
1 parent 3a5d06a
commit be57ed6
Show file tree

Hide file tree

Showing 37 changed files with 2,507 additions and 346 deletions.
diff --git a/asmcomp/cmm_helpers.ml b/asmcomp/cmm_helpers.ml
@@ -790,23 +790,29 @@ let int_array_ref arr ofs dbg =
 let unboxed_float_array_ref arr ofs dbg =
   Cop(Cload (Double, Mutable),
     [array_indexing log2_size_float arr ofs dbg], dbg)
-let float_array_ref arr ofs dbg =
-  box_float dbg Lambda.alloc_heap (unboxed_float_array_ref arr ofs dbg)
+let float_array_ref mode arr ofs dbg =
+  box_float dbg mode (unboxed_float_array_ref arr ofs dbg)
 
-let addr_array_set arr ofs newval dbg =
+let addr_array_set_heap arr ofs newval dbg =
   Cop(Cextcall("caml_modify", typ_void, [], false),
       [array_indexing log2_size_addr arr ofs dbg; newval], dbg)
+
+let addr_array_set_local arr ofs newval dbg =
+  Cop(Cextcall("caml_modify_local", typ_void, [], false),
+      [arr; untag_int ofs dbg; newval], dbg)
+
+let addr_array_set (mode : Lambda.modify_mode) arr ofs newval dbg =
+  match mode with
+  | Modify_heap -> addr_array_set_heap arr ofs newval dbg
+  | Modify_maybe_stack -> addr_array_set_local arr ofs newval dbg
+(* int and float arrays can be written to uniformly regardless of their mode *)
 let int_array_set arr ofs newval dbg =
   Cop(Cstore (Word_int, Assignment),
     [array_indexing log2_size_addr arr ofs dbg; newval], dbg)
 let float_array_set arr ofs newval dbg =
   Cop(Cstore (Double, Assignment),
     [array_indexing log2_size_float arr ofs dbg; newval], dbg)
 
-let addr_array_set_local arr ofs newval dbg =
-  Cop(Cextcall("caml_modify_local", typ_void, [], false),
-      [arr; untag_int ofs dbg; newval], dbg)
-
 let addr_array_initialize arr ofs newval dbg =
   Cop(Cextcall("caml_initialize", typ_void, [], false),
       [array_indexing log2_size_addr arr ofs dbg; newval], dbg)
@@ -2748,28 +2754,28 @@ let bigstring_load size unsafe mode arg1 arg2 dbg =
           idx
           (unaligned_load size ba_data idx dbg)))))
 
-let arrayref_unsafe kind arg1 arg2 dbg =
-  match (kind : Lambda.array_kind) with
-  | Pgenarray ->
+let arrayref_unsafe rkind arg1 arg2 dbg =
+  match (rkind : Lambda.array_ref_kind) with
+  | Pgenarray_ref mode ->
       bind "index" arg2 (fun idx ->
         bind "arr" arg1 (fun arr ->
           Cifthenelse(is_addr_array_ptr arr dbg,
                       dbg,
                       addr_array_ref arr idx dbg,
                       dbg,
-                      float_array_ref arr idx dbg,
+                      float_array_ref mode arr idx dbg,
                       dbg, Any)))
-  | Paddrarray ->
+  | Paddrarray_ref ->
       addr_array_ref arg1 arg2 dbg
-  | Pintarray ->
+  | Pintarray_ref ->
       (* CR mshinwell: for int/addr_array_ref move "dbg" to first arg *)
       int_array_ref arg1 arg2 dbg
-  | Pfloatarray ->
-      float_array_ref arg1 arg2 dbg
+  | Pfloatarray_ref mode ->
+      float_array_ref mode arg1 arg2 dbg
 
-let arrayref_safe kind arg1 arg2 dbg =
-  match (kind : Lambda.array_kind) with
-  | Pgenarray ->
+let arrayref_safe rkind arg1 arg2 dbg =
+  match (rkind : Lambda.array_ref_kind) with
+  | Pgenarray_ref mode ->
       bind "index" arg2 (fun idx ->
       bind "arr" arg1 (fun arr ->
       bind "header" (get_header_without_profinfo arr dbg) (fun hdr ->
@@ -2780,7 +2786,7 @@ let arrayref_safe kind arg1 arg2 dbg =
                         dbg,
                         addr_array_ref arr idx dbg,
                         dbg,
-                        float_array_ref arr idx dbg,
+                        float_array_ref mode arr idx dbg,
                         dbg, Any))
         else
           Cifthenelse(is_addr_array_hdr hdr dbg,
@@ -2791,42 +2797,42 @@ let arrayref_safe kind arg1 arg2 dbg =
             dbg,
             Csequence(
               make_checkbound dbg [float_array_length_shifted hdr dbg; idx],
-              float_array_ref arr idx dbg),
+              float_array_ref mode arr idx dbg),
             dbg, Any))))
-      | Paddrarray ->
-          bind "index" arg2 (fun idx ->
-          bind "arr" arg1 (fun arr ->
-            Csequence(
-              make_checkbound dbg [
-                addr_array_length_shifted
-                  (get_header_without_profinfo arr dbg) dbg; idx],
-              addr_array_ref arr idx dbg)))
-      | Pintarray ->
-          bind "index" arg2 (fun idx ->
-          bind "arr" arg1 (fun arr ->
-            Csequence(
-              make_checkbound dbg [
-                addr_array_length_shifted
-                  (get_header_without_profinfo arr dbg) dbg; idx],
-              int_array_ref arr idx dbg)))
-      | Pfloatarray ->
-          box_float dbg Lambda.alloc_heap (
-            bind "index" arg2 (fun idx ->
-            bind "arr" arg1 (fun arr ->
-              Csequence(
-                make_checkbound dbg [
-                  float_array_length_shifted
-                    (get_header_without_profinfo arr dbg) dbg;
-                  idx],
-                unboxed_float_array_ref arr idx dbg))))
+  | Paddrarray_ref ->
+      bind "index" arg2 (fun idx ->
+      bind "arr" arg1 (fun arr ->
+        Csequence(
+          make_checkbound dbg [
+            addr_array_length_shifted
+              (get_header_without_profinfo arr dbg) dbg; idx],
+          addr_array_ref arr idx dbg)))
+  | Pintarray_ref ->
+      bind "index" arg2 (fun idx ->
+      bind "arr" arg1 (fun arr ->
+        Csequence(
+          make_checkbound dbg [
+            addr_array_length_shifted
+              (get_header_without_profinfo arr dbg) dbg; idx],
+          int_array_ref arr idx dbg)))
+  | Pfloatarray_ref mode ->
+      box_float dbg mode (
+        bind "index" arg2 (fun idx ->
+        bind "arr" arg1 (fun arr ->
+          Csequence(
+            make_checkbound dbg [
+              float_array_length_shifted
+                (get_header_without_profinfo arr dbg) dbg;
+              idx],
+            unboxed_float_array_ref arr idx dbg))))
 
 type ternary_primitive =
   expression -> expression -> expression -> Debuginfo.t -> expression
 
 let setfield_computed ptr init arg1 arg2 arg3 dbg =
   match assignment_kind ptr init with
   | Caml_modify ->
-      return_unit dbg (addr_array_set arg1 arg2 arg3 dbg)
+      return_unit dbg (addr_array_set_heap arg1 arg2 arg3 dbg)
   | Caml_modify_local ->
       return_unit dbg (addr_array_set_local arg1 arg2 arg3 dbg)
   | Caml_initialize ->
@@ -2850,30 +2856,30 @@ let bytesset_safe arg1 arg2 arg3 dbg =
               [add_int str idx dbg; newval],
               dbg))))))
 
-let arrayset_unsafe kind arg1 arg2 arg3 dbg =
-  return_unit dbg (match (kind: Lambda.array_kind) with
-  | Pgenarray ->
+let arrayset_unsafe skind arg1 arg2 arg3 dbg =
+  return_unit dbg (match (skind: Lambda.array_set_kind) with
+  | Pgenarray_set mode ->
       bind "newval" arg3 (fun newval ->
         bind "index" arg2 (fun index ->
           bind "arr" arg1 (fun arr ->
             Cifthenelse(is_addr_array_ptr arr dbg,
                         dbg,
-                        addr_array_set arr index newval dbg,
+                        addr_array_set mode arr index newval dbg,
                         dbg,
                         float_array_set arr index (unbox_float dbg newval)
                           dbg,
                         dbg, Any))))
-  | Paddrarray ->
-      addr_array_set arg1 arg2 arg3 dbg
-  | Pintarray ->
+  | Paddrarray_set mode ->
+      addr_array_set mode arg1 arg2 arg3 dbg
+  | Pintarray_set ->
       int_array_set arg1 arg2 arg3 dbg
-  | Pfloatarray ->
+  | Pfloatarray_set ->
       float_array_set arg1 arg2 arg3 dbg
   )
 
-let arrayset_safe kind arg1 arg2 arg3 dbg =
-  return_unit dbg (match (kind: Lambda.array_kind) with
-  | Pgenarray ->
+let arrayset_safe skind arg1 arg2 arg3 dbg =
+  return_unit dbg (match (skind: Lambda.array_set_kind) with
+  | Pgenarray_set mode ->
       bind "newval" arg3 (fun newval ->
       bind "index" arg2 (fun idx ->
       bind "arr" arg1 (fun arr ->
@@ -2883,7 +2889,7 @@ let arrayset_safe kind arg1 arg2 arg3 dbg =
             make_checkbound dbg [addr_array_length_shifted hdr dbg; idx],
             Cifthenelse(is_addr_array_hdr hdr dbg,
                         dbg,
-                        addr_array_set arr idx newval dbg,
+                        addr_array_set mode arr idx newval dbg,
                         dbg,
                         float_array_set arr idx
                           (unbox_float dbg newval)
@@ -2895,14 +2901,14 @@ let arrayset_safe kind arg1 arg2 arg3 dbg =
             dbg,
             Csequence(
               make_checkbound dbg [addr_array_length_shifted hdr dbg; idx],
-              addr_array_set arr idx newval dbg),
+              addr_array_set mode arr idx newval dbg),
             dbg,
             Csequence(
               make_checkbound dbg [float_array_length_shifted hdr dbg; idx],
               float_array_set arr idx
                 (unbox_float dbg newval) dbg),
             dbg, Any)))))
-  | Paddrarray ->
+  | Paddrarray_set mode ->
       bind "newval" arg3 (fun newval ->
       bind "index" arg2 (fun idx ->
       bind "arr" arg1 (fun arr ->
@@ -2911,8 +2917,8 @@ let arrayset_safe kind arg1 arg2 arg3 dbg =
             addr_array_length_shifted
               (get_header_without_profinfo arr dbg) dbg;
             idx],
-          addr_array_set arr idx newval dbg))))
-  | Pintarray ->
+          addr_array_set mode arr idx newval dbg))))
+  | Pintarray_set ->
       bind "newval" arg3 (fun newval ->
       bind "index" arg2 (fun idx ->
       bind "arr" arg1 (fun arr ->
@@ -2922,7 +2928,7 @@ let arrayset_safe kind arg1 arg2 arg3 dbg =
               (get_header_without_profinfo arr dbg) dbg;
             idx],
           int_array_set arr idx newval dbg))))
-  | Pfloatarray ->
+  | Pfloatarray_set ->
       bind_load "newval" arg3 (fun newval ->
       bind "index" arg2 (fun idx ->
       bind "arr" arg1 (fun arr ->

diff --git a/asmcomp/cmm_helpers.mli b/asmcomp/cmm_helpers.mli
@@ -265,11 +265,15 @@ val addr_array_ref : expression -> expression -> Debuginfo.t -> expression
 val int_array_ref : expression -> expression -> Debuginfo.t -> expression
 val unboxed_float_array_ref :
   expression -> expression -> Debuginfo.t -> expression
-val float_array_ref : expression -> expression -> Debuginfo.t -> expression
-val addr_array_set :
+val float_array_ref :
+  Lambda.alloc_mode -> expression -> expression -> Debuginfo.t -> expression
+val addr_array_set_heap :
   expression -> expression -> expression -> Debuginfo.t -> expression
 val addr_array_set_local :
   expression -> expression -> expression -> Debuginfo.t -> expression
+val addr_array_set :
+  Lambda.modify_mode -> expression -> expression -> expression -> Debuginfo.t ->
+  expression
 val int_array_set :
   expression -> expression -> expression -> Debuginfo.t -> expression
 val float_array_set :
@@ -555,9 +559,9 @@ val bigstring_load :
 
 (** Arrays *)
 
-(** Array access. Args: array, index *)
-val arrayref_unsafe : Lambda.array_kind -> binary_primitive
-val arrayref_safe : Lambda.array_kind -> binary_primitive
+(** Array access.  Args: array, index *)
+val arrayref_unsafe : Lambda.array_ref_kind -> binary_primitive
+val arrayref_safe : Lambda.array_ref_kind -> binary_primitive
 
 type ternary_primitive =
   expression -> expression -> expression -> Debuginfo.t -> expression
@@ -578,8 +582,8 @@ val bytesset_safe : ternary_primitive
     _unboxed_ float. Otherwise, it is expected to be a regular caml value,
     including in the case where the array contains floats.
     Args: array, index, value *)
-val arrayset_unsafe : Lambda.array_kind -> ternary_primitive
-val arrayset_safe : Lambda.array_kind -> ternary_primitive
+val arrayset_unsafe : Lambda.array_set_kind -> ternary_primitive
+val arrayset_safe : Lambda.array_set_kind -> ternary_primitive
 
 (** Set a chunk of data in the given bytes or bigstring structure.
     See also [string_load] and [bigstring_load].

diff --git a/asmcomp/cmmgen.ml b/asmcomp/cmmgen.ml
@@ -1067,10 +1067,10 @@ and transl_prim_2 env p arg1 arg2 dbg =
       bigstring_load size unsafe mode (transl env arg1) (transl env arg2) dbg
 
   (* Array operations *)
-  | Parrayrefu kind ->
-      arrayref_unsafe kind (transl env arg1) (transl env arg2) dbg
-  | Parrayrefs kind ->
-      arrayref_safe kind (transl env arg1) (transl env arg2) dbg
+  | Parrayrefu rkind ->
+      arrayref_unsafe rkind (transl env arg1) (transl env arg2) dbg
+  | Parrayrefs rkind ->
+      arrayref_safe rkind (transl env arg1) (transl env arg2) dbg
 
   (* Boxed integers *)
   | Paddbint (bi, mode) ->
@@ -1155,20 +1155,20 @@ and transl_prim_3 env p arg1 arg2 arg3 dbg =
         (transl env arg1) (transl env arg2) (transl env arg3) dbg
 
   (* Array operations *)
-  | Parraysetu kind ->
+  | Parraysetu skind ->
       let newval =
-        match kind with
-        | Pfloatarray -> transl_unbox_float dbg env arg3
+        match skind with
+        | Pfloatarray_set -> transl_unbox_float dbg env arg3
         | _ -> transl env arg3
       in
-      arrayset_unsafe kind (transl env arg1) (transl env arg2) newval dbg
-  | Parraysets kind ->
+      arrayset_unsafe skind (transl env arg1) (transl env arg2) newval dbg
+  | Parraysets skind ->
       let newval =
-        match kind with
-        | Pfloatarray -> transl_unbox_float dbg env arg3
+        match skind with
+        | Pfloatarray_set -> transl_unbox_float dbg env arg3
         | _ -> transl env arg3
       in
-      arrayset_safe kind (transl env arg1) (transl env arg2) newval dbg
+      arrayset_safe skind (transl env arg1) (transl env arg2) newval dbg
 
   | Pbytes_set(size, unsafe) ->
       bytes_set size unsafe (transl env arg1) (transl env arg2)

diff --git a/boot/ocamlc b/boot/ocamlc
diff --git a/boot/ocamllex b/boot/ocamllex
diff --git a/bytecomp/bytegen.ml b/bytecomp/bytegen.ml
@@ -454,18 +454,23 @@ let comp_primitive p args =
   | Pbytes_load_32(_) -> Kccall("caml_bytes_get32", 2)
   | Pbytes_load_64(_) -> Kccall("caml_bytes_get64", 2)
   | Parraylength _ -> Kvectlength
-  | Parrayrefs Pgenarray -> Kccall("caml_array_get", 2)
-  | Parrayrefs Pfloatarray -> Kccall("caml_floatarray_get", 2)
-  | Parrayrefs _ -> Kccall("caml_array_get_addr", 2)
-  | Parraysets Pgenarray -> Kccall("caml_array_set", 3)
-  | Parraysets Pfloatarray -> Kccall("caml_floatarray_set", 3)
-  | Parraysets _ -> Kccall("caml_array_set_addr", 3)
-  | Parrayrefu Pgenarray -> Kccall("caml_array_unsafe_get", 2)
-  | Parrayrefu Pfloatarray -> Kccall("caml_floatarray_unsafe_get", 2)
-  | Parrayrefu _ -> Kgetvectitem
-  | Parraysetu Pgenarray -> Kccall("caml_array_unsafe_set", 3)
-  | Parraysetu Pfloatarray -> Kccall("caml_floatarray_unsafe_set", 3)
-  | Parraysetu _ -> Ksetvectitem
+  (* In bytecode, nothing is ever actually stack-allocated, so we ignore the
+     array modes (allocation for [Parrayref{s,u}], modification for
+     [Parrayset{s,u}]). *)
+  | Parrayrefs (Pgenarray_ref _) -> Kccall("caml_array_get", 2)
+  | Parrayrefs (Pfloatarray_ref _) -> Kccall("caml_floatarray_get", 2)
+  | Parrayrefs (Paddrarray_ref | Pintarray_ref) ->
+      Kccall("caml_array_get_addr", 2)
+  | Parraysets (Pgenarray_set _) -> Kccall("caml_array_set", 3)
+  | Parraysets Pfloatarray_set -> Kccall("caml_floatarray_set", 3)
+  | Parraysets (Paddrarray_set _ | Pintarray_set) ->
+      Kccall("caml_array_set_addr", 3)
+  | Parrayrefu (Pgenarray_ref _) -> Kccall("caml_array_unsafe_get", 2)
+  | Parrayrefu (Pfloatarray_ref _) -> Kccall("caml_floatarray_unsafe_get", 2)
+  | Parrayrefu (Paddrarray_ref | Pintarray_ref) -> Kgetvectitem
+  | Parraysetu (Pgenarray_set _) -> Kccall("caml_array_unsafe_set", 3)
+  | Parraysetu Pfloatarray_set -> Kccall("caml_floatarray_unsafe_set", 3)
+  | Parraysetu (Paddrarray_set _ | Pintarray_set) -> Ksetvectitem
   | Pctconst c ->
      let const_name = match c with
        | Big_endian -> "big_endian"