From ecf2613f4b460297d6badc3f0cc4b7d31360bddf Mon Sep 17 00:00:00 2001
From: Gabriel Scherer <gabriel.scherer@gmail.com>
Date: Fri, 25 Aug 2023 01:02:02 +0200
Subject: [PATCH] bigarray: fix #12491 (memory accounting of shared bigarray
 data)

---
 runtime/bigarray.c | 43 ++++++++++++++++++++++++++++---------------
 1 file changed, 28 insertions(+), 15 deletions(-)

diff --git a/runtime/bigarray.c b/runtime/bigarray.c
index 08bb5ae8f0ca..dcf12c83680f 100644
--- a/runtime/bigarray.c
+++ b/runtime/bigarray.c
@@ -229,48 +229,60 @@ static int ba_compute_size(int flags, int num_dims, intnat * dim, uintnat *num_e
                          size));
 }
 
-static caml_ba_alloc_gen(int flags, int num_dims, void * data, intnat * dim, int owns_data)
-{
-    
-}
+/* [caml_ba_alloc_gen] will allocate a new bigarray object in the heap.
+
+   If [data] is NULL, the memory for the contents is also allocated.
+   (with [malloc]) by [caml_ba_alloc_gen].
+
+   If [data] is NULL or [owns_data] is non-zero, we consider that [data] is now
+   owned by the OCaml heap, and adjust GC pacing accordingly --
+   the external memory considered is the size of the bigarray.
 
-/* [caml_ba_alloc] will allocate a new bigarray object in the heap.
-   If [data] is NULL, the memory for the contents is also allocated
-   (with [malloc]) by [caml_ba_alloc].
    [data] cannot point into the OCaml heap.
    [dim] may point into an object in the OCaml heap.
 */
-CAMLexport value
-caml_ba_alloc(int flags, int num_dims, void * data, intnat * dim)
+static value caml_ba_alloc_gen(int flags, int num_dims, void * data, intnat * dim, int owns_data)
 {
   uintnat asize, num_elts, size;
-  int i, is_managed;
   value res;
   struct caml_ba_array * b;
   intnat dimcopy[CAML_BA_MAX_NUM_DIMS];
 
   CAMLassert(num_dims >= 0 && num_dims <= CAML_BA_MAX_NUM_DIMS);
   CAMLassert((flags & CAML_BA_KIND_MASK) < CAML_BA_FIRST_UNIMPLEMENTED_KIND);
-  for (i = 0; i < num_dims; i++) dimcopy[i] = dim[i];
+  for (int i = 0; i < num_dims; i++) dimcopy[i] = dim[i];
   if (!ba_compute_size(flags, num_dims, dimcopy, &num_elts, &size))
     caml_raise_out_of_memory();
   if (data == NULL) {
     data = malloc(size);
     if (data == NULL && size != 0) caml_raise_out_of_memory();
+    owns_data = 1;
     flags |= CAML_BA_MANAGED;
   }
   asize = SIZEOF_BA_ARRAY + num_dims * sizeof(intnat);
-  is_managed = ((flags & CAML_BA_MANAGED_MASK) == CAML_BA_MANAGED);
-  res = caml_alloc_custom_mem(&caml_ba_ops, asize, is_managed ? size : 0);
+  res = caml_alloc_custom_mem(&caml_ba_ops, asize, owns_data ? size : 0);
   b = Caml_ba_array_val(res);
   b->data = data;
   b->num_dims = num_dims;
   b->flags = flags;
   b->proxy = NULL;
-  for (i = 0; i < num_dims; i++) b->dim[i] = dimcopy[i];
+  for (int i = 0; i < num_dims; i++) b->dim[i] = dimcopy[i];
   return res;
 }
 
+/* [caml_ba_alloc] will allocate a new bigarray object in the heap.
+   If [data] is NULL, the memory for the contents is also allocated.
+   (with [malloc]) by [caml_ba_alloc].
+   [data] cannot point into the OCaml heap.
+   [dim] may point into an object in the OCaml heap.
+*/
+CAMLexport value
+caml_ba_alloc(int flags, int num_dims, void * data, intnat * dim)
+{
+  int is_managed = ((flags & CAML_BA_MANAGED_MASK) == CAML_BA_MANAGED);
+  return caml_ba_alloc_gen(flags, num_dims, data, dim, is_managed);
+}
+
 /* Same as caml_ba_alloc, but dimensions are passed as a list of
    arguments */
 
@@ -1106,7 +1118,8 @@ static value caml_ba_inherit(value vb, int num_dims, intnat * dim)
   CAMLparam1 (vb);
   CAMLlocal1 (res);
 #define b (Caml_ba_array_val(vb))
-  res = caml_ba_alloc(b->flags, num_dims, b->data, dim);
+  int owns_data = 0;
+  res = caml_ba_alloc_gen(b->flags, num_dims, b->data, dim, owns_data);
    /* Copy the finalization function from the original array (PR#8568) */
   Custom_ops_val(res) = Custom_ops_val(vb);
   /* Create or update proxy in case of managed bigarray */