Add [Obj.uniqely_reachable_words]

This function takes in a list of objects and for each one computes the number
of words of memory that can be reachabed from that object, but no others in
the list.

ocaml/runtime/extern.c

@@ -274,6 +274,7 @@ Caml_inline int extern_lookup_position(value obj,
       return 0;
     }
     if (pos_table.entries[h].obj == obj) {
+      *h_out = h;
       *pos_out = pos_table.entries[h].pos;
       return 1;
     }
@@ -295,6 +296,27 @@ static void extern_record_location(value obj, uintnat h)
   if (obj_counter >= pos_table.threshold) extern_resize_position_table();
 }
 
+/* Record the output position for the given object [obj]. */
+/* The [h] parameter is the index in the hash table where the object
+   must be inserted.  It was determined during lookup. */
+
+static void extern_record_location_with_data(value obj, uintnat h, uintnat data)
+{
+  if (extern_flags & NO_SHARING) return;
+  bitvect_set(pos_table.present, h);
+  pos_table.entries[h].obj = obj;
+  pos_table.entries[h].pos = data;
+  obj_counter++;
+  if (obj_counter >= pos_table.threshold) extern_resize_position_table();
+}
+
+static void extern_update_location_with_data(value obj, uintnat h, uintnat data)
+{
+  if (extern_flags & NO_SHARING) return;
+  pos_table.entries[h].obj = obj;
+  pos_table.entries[h].pos = data;
+}
+
 /* To buffer the output */
 
 static char * extern_userprovided_output;
@@ -1136,18 +1158,54 @@ CAMLexport void caml_serialize_block_float_8(void * data, intnat len)
 #endif
 }
 
-CAMLprim value caml_obj_reachable_words(value v)
-{
-  intnat size;
-  struct extern_item * sp;
-  uintnat h = 0;
-  uintnat pos;
+/* Performs traversal through the OCaml object reachability graph to deterime
+   how much memory an object has access to.
+
+   Assumes that the position_table has already been initialized using
+   [caml_obj_reachable_words_init].  We can run this function multiple times
+   without clearing the position table to share data between runs starting
+   from different roots. The mode must be an integer between 1 and 3 inclusive
+   and the identifier must be a positive integer.
+
+   For each value node visited, we record its traversal status in the [pos] field
+   of its entry in [position_table.entries]. The statuses are:
+   * >0: node has been visited exactly once, from root with identifier equal to
+         this integer.
+   * -1: node has been visited at least two times. It will be ignored in future
+         traversals. Also used by mode 3 to show that this node has already been
+	 counted.
+   * -2: node is a root that hasn't been processed yet.
+   * -3: node is a root that has been visited by mode 2.
+
+   The expected usage is:
+   - Iterate through list of roots, running in mode 1 for each. This initializes
+     information about which nodes are roots. Returns 0.
+   - Iterate through list of roots, running in mode 2 for each. This performs
+     the main traversal, for each descendant recording that it is reachable from
+     the current root. Returns the total size of elements marked, that is ones
+     that are reachable from the current root and can be reached from at most
+     one root from the ones that already ran.
+   - Iterate through list of roots, running in mode 3 for each. This uses the
+     information computed in the previous step to return the sum of sizes of
+     elements that are reachable from this root but no other root. */
+CAMLprim value caml_obj_reachable_words_once(value root, value mode_v, value identifier_v) {
+  intnat mode = Int_val(mode_v), identifier = Int_val(identifier_v);
 
-  obj_counter = 0;
-  extern_flags = 0;
-  extern_init_position_table();
+  struct extern_item * sp;
+  intnat size;
+  uintnat mark, new_mark;
+  value v = root;
+  uintnat h;
+  int previously_marked, should_traverse;
   sp = extern_stack;
   size = 0;
+
+  if (mode == 1) {
+    h = Hash(v);
+    extern_record_location_with_data(v, h, -2);
+    return Long_val(0);
+  }
+
   while (1) {
     if (Is_long(v)) {
       /* Tagged integers contribute 0 to the size, nothing to do */
@@ -1157,45 +1215,85 @@ CAMLprim value caml_obj_reachable_words(value v)
          between major heap blocks and out-of-heap blocks,
          and the test above is always false,
          so we end up counting out-of-heap blocks too. */
-    } else if (extern_lookup_position(v, &pos, &h)) {
-      /* Already seen and counted, nothing to do */
     } else {
-      header_t hd = Hd_val(v);
-      tag_t tag = Tag_hd(hd);
-      mlsize_t sz = Wosize_hd(hd);
-      /* Infix pointer: go back to containing closure */
-      if (tag == Infix_tag) {
-        v = v - Infix_offset_hd(hd);
-        continue;
+      if ((previously_marked = extern_lookup_position(v, &mark, &h))) {
+        if (mark == -2 && v == root) {
+          /* mode == 2 */
+          should_traverse = 1;
+	  new_mark = -3;
+	} else if (mark == -3 && v == root) {
+	  /* mode == 3 || mode == 2 */
+          should_traverse = (mode == 3);
+	  new_mark = -1;
+	} else if (mark == -1 || mark == -2 || mark == -3) {
+          /* ignore (-1) or root but not ours (-2/-3) */
+	  should_traverse = 0;
+	} else {
+          /* mark is identifier (ours or not) */
+          should_traverse = !(mark == identifier && mode == 2);
+          new_mark = -1;
+	}
+      } else {
+        /* mode == 2 && not root */
+        should_traverse = 1;
+	new_mark = identifier;
       }
-      /* Remember that we've visited this block */
-      extern_record_location(v, h);
-      /* The block contributes to the total size */
-      size += 1 + sz;           /* header word included */
-      if (tag < No_scan_tag) {
-        /* i is the position of the first field to traverse recursively */
-        uintnat i =
-          tag == Closure_tag ? Start_env_closinfo(Closinfo_val(v)) : 0;
-        if (i < sz) {
-          if (i < sz - 1) {
-            /* Remember that we need to count fields i + 1 ... sz - 1 */
-            sp++;
-            if (sp >= extern_stack_limit) sp = extern_resize_stack(sp);
-            sp->v = &Field(v, i + 1);
-            sp->count = sz - i - 1;
-          }
-          /* Continue with field i */
-          v = Field(v, i);
+
+      if (should_traverse) {
+        header_t hd = Hd_val(v);
+        tag_t tag = Tag_hd(hd);
+        mlsize_t sz = Wosize_hd(hd);
+        /* Infix pointer: go back to containing closure */
+        if (tag == Infix_tag) {
+          v = v - Infix_offset_hd(hd);
           continue;
         }
+        if (!previously_marked) {
+          extern_record_location_with_data(v, h, new_mark);
+        } else {
+          extern_update_location_with_data(v, h, new_mark);
+        }
+        /* The block contributes to the total size */
+        size += 1 + sz;           /* header word included */
+        if (tag < No_scan_tag) {
+          /* i is the position of the first field to traverse recursively */
+          uintnat i =
+            tag == Closure_tag ? Start_env_closinfo(Closinfo_val(v)) : 0;
+          if (i < sz) {
+            if (i < sz - 1) {
+              /* Remember that we need to count fields i + 1 ... sz - 1 */
+              sp++;
+              if (sp >= extern_stack_limit) sp = extern_resize_stack(sp);
+              sp->v = &Field(v, i + 1);
+              sp->count = sz - i - 1;
+            }
+            /* Continue with field i */
+            v = Field(v, i);
+            continue;
+          }
+        }
       }
     }
     /* Pop one more item to traverse, if any */
     if (sp == extern_stack) break;
     v = *((sp->v)++);
     if (--(sp->count) == 0) sp--;
   }
+
+  return Val_long(size);
+}
+
+CAMLprim value caml_obj_reachable_words_init(value v)
+{
+  obj_counter = 0;
+  extern_flags = 0;
+  extern_init_position_table();
+  return Val_unit;
+}
+
+CAMLprim value caml_obj_reachable_words_cleanup(value v)
+{
   extern_free_stack();
   extern_free_position_table();
-  return Val_long(size);
+  return Val_unit;
 }

ocaml/stdlib/obj.ml

@@ -36,7 +36,27 @@ external tag : t -> int = "caml_obj_tag" [@@noalloc]
    restriction is likely not respected by callees of this module. *)
 external size : t -> int = "%obj_size"
 let [@inline always] size t = size (Sys.opaque_identity t)
-external reachable_words : t -> int = "caml_obj_reachable_words"
+external reachable_words_init : unit -> unit = "caml_obj_reachable_words_init"
+external reachable_words_once : t -> mode:int -> identifier:int -> int = "caml_obj_reachable_words_once"
+external reachable_words_cleanup : unit -> unit = "caml_obj_reachable_words_cleanup"
+let reachable_words (t : t) : int =
+  reachable_words_init ();
+  let ret = reachable_words_once t ~mode:1 ~identifier:1 in
+  reachable_words_cleanup ();
+  ret
+let uniquely_reachable_words (ts : t list) : int list =
+  reachable_words_init ();
+  let rec traverse ~mode i = function
+    | [] -> ()
+    | hd :: tl -> (let _ = reachable_words_once hd ~mode ~identifier:i in traverse ~mode (i+1) tl) in
+  let rec accumulate ~mode i = function
+    | [] -> []
+    | hd :: tl -> (let v = reachable_words_once hd ~mode ~identifier:i in v :: accumulate ~mode (i+1) tl) in
+  let () = traverse ~mode:1 1 ts in
+  let () = traverse ~mode:2 1 ts in
+  let ret = accumulate ~mode:3 1 ts in
+  reachable_words_cleanup ();
+  ret
 external field : t -> int -> t = "%obj_field"
 let [@inline always] field t index = field (Sys.opaque_identity t) index
 external set_field : t -> int -> t -> unit = "%obj_set_field"

ocaml/stdlib/obj.mli

@@ -32,7 +32,7 @@ val is_block : t -> bool
 external is_int : t -> bool = "%obj_is_int"
 external tag : t -> int = "caml_obj_tag" [@@noalloc]
 val size : t -> int
-external reachable_words : t -> int = "caml_obj_reachable_words"
+val reachable_words : t -> int
   (**
      Computes the total size (in words, including the headers) of all
      heap blocks accessible from the argument.  Statically
@@ -42,6 +42,14 @@ external reachable_words : t -> int = "caml_obj_reachable_words"
      @since 4.04
   *)
 
+val uniquely_reachable_words : t list -> int list
+(** For each element of the list, computes the total size (as defined
+    above by [reachable_words]) of all heap blocks accessible from the
+    argument but excluding all blocks accessible from previous arguments.
+
+    @since 4.15
+  *)
+
 val field : t -> int -> t
 
 (** When using flambda:

ocaml/testsuite/tests/lib-obj/uniquely_reachable_words.ml

@@ -0,0 +1,79 @@
+(* TEST
+*)
+
+let native =
+  match Sys.backend_type with
+  | Sys.Native -> true
+  | Sys.Bytecode -> false
+  | Sys.Other s -> print_endline s; assert false
+
+let sizes xs = Obj.uniquely_reachable_words (List.map Obj.repr xs)
+
+let expect_sizes xs exp_bytecode exp_native =
+  let actual = sizes xs in
+  let expected = if native then exp_native else exp_bytecode in
+  List.combine actual expected
+  |> List.iteri (fun i (a, e) ->
+    if a <> e then
+      Printf.printf "index = %i; actual = %i; expected = %i\n" i a e)
+
+type node = { id: int; used_memory: int list; mutable children: node list }
+let make id ch = { id; used_memory = List.init (Int.shift_left 1 id) (fun _ -> id); children = ch }
+
+let f () =
+  let n10 = make 10 [] in
+  let n9 = make 9 [] in
+  let n8 = make 8 [] in
+  let n7 = make 7 [n8] in
+  let n6 = make 6 [n7; n8] in
+  n7.children <- n6 :: n7.children;
+  let n5 = make 5 [n9] in
+  let n4 = make 4 [n6] in
+  let n3 = make 3 [] in
+  let n2 = make 2 [n3; n4] in
+  let n1 = make 1 [n5; n10] in
+  let n0 = make 0 [n3; n5] in
+  let n14 = make 14 [] in
+  let n13 = make 13 [n14] in
+  let n12 = make 12 [n14] in
+  let n11 = make 11 [n12; n13] in
+  (*  /-> 10
+   * 1 --> 5 --> 9
+   *   /
+   * 0 --> 3       ->8<
+   *   /          /    \
+   * 2 --> 4 --> 6 <--> 7
+   *
+   *   /-> 12 >--\
+   * 11 -> 13 >- 14
+   *)
+  expect_sizes [n0; n1; n2] [13; 3092; 1445] [13; 3092; 1442]; (* Proper roots *)
+  expect_sizes [n0; n2; n1] [13; 1445; 3092] [13; 1442; 3092]; (* check permutation doesn't matter *)
+  expect_sizes [n1; n0; n2] [3092; 13; 1445] [3092; 13; 1442];
+  expect_sizes [n1; n2; n0] [3092; 1445; 13] [3092; 1442; 13];
+  expect_sizes [n2; n0; n1] [1445; 13; 3092] [1442; 13; 3092];
+  expect_sizes [n2; n1; n0] [1445; 3092; 13] [1442; 3092; 13];
+  expect_sizes [n1; n2] [4735; 1473] [4735; 1470];
+  expect_sizes [n0; n2] [1656; 1445] [1656; 1442];
+  expect_sizes [n0; n1] [41; 3092] [41; 3092];
+
+  expect_sizes [n6; n7] [202; 394] [199; 391]; (* Cycles between roots *)
+  expect_sizes [n6; n7; n2] [202; 394; 105] [199; 391; 105];
+  expect_sizes [n6; n7; n8] [202; 394; 772] [199; 391; 772];
+
+  expect_sizes [n5; n9] [103; 1540] [103; 1540]; (* Root is parent of another root *)
+  expect_sizes [n5; n9; n3] [103; 1540; 28] [103; 1540; 28];
+  expect_sizes [n5; n9; n3; n0] [103; 1540; 28; 13] [103; 1540; 28; 13];
+  expect_sizes [n1; n10; n5] [16; 3076; 1643] [16; 3076; 1643];
+
+  expect_sizes [n12; n13] [12295; 24583] [12292; 24580]; (* Multiple owners *)
+  expect_sizes [n12; n13; n14] [12295; 24583; 49156] [12292; 24580; 49156];
+  expect_sizes [n11; n12] [30737; 12295] [30734; 12292];
+  expect_sizes [n12; n11] [12295; 30737] [12292; 30734];
+  expect_sizes [n11] [92188] [92185];
+
+  expect_sizes [n8; n9; n10] [772; 1540; 3076] [772; 1540; 3076]; (* Leaves *)
+
+  print_endline "OK"
+
+let () = f ()

ocaml/testsuite/tests/lib-obj/uniquely_reachable_words.reference

@@ -0,0 +1 @@
+OK