Call counts in Spacetime (#1180)

2 files changed, 58 insertions, 14 deletions

diff --git a/asmrun/spacetime.c b/asmrun/spacetime.c
index aaf7f6af8e..fa1e87d121 100644
--- a/asmrun/spacetime.c
+++ b/asmrun/spacetime.c
@@ -118,14 +118,26 @@ static void reinitialise_free_node_block(void)
 extern value val_process_id;
 #endif
 
-static uint32_t version_number = 0;
+enum {
+  FEATURE_CALL_COUNTS = 1,
+} features;
+
+static uint16_t version_number = 0;
 static uint32_t magic_number_base = 0xace00ace;
 
 static void caml_spacetime_write_magic_number_internal(struct channel* chan)
 {
-  value magic_number =
+  value magic_number;
+  uint16_t features = 0;
+
+#ifdef WITH_SPACETIME_CALL_COUNTS
+  features |= FEATURE_CALL_COUNTS;
+#endif
+
+  magic_number =
     Val_long(((uint64_t) magic_number_base)
-             | (((uint64_t) version_number) << 32));
+             | (((uint64_t) version_number) << 32)
+             | (((uint64_t) features) << 48));
 
   Lock(chan);
   caml_output_val(chan, magic_number, Val_long(0));
@@ -568,7 +580,8 @@ static c_node* allocate_c_node(void)
 
   node->gc_header =
     Make_header(sizeof(c_node)/sizeof(uintnat) - 1, C_node_tag, Caml_black);
-  node->data.callee_node = Val_unit;
+  node->data.call.callee_node = Val_unit;
+  node->data.call.call_count = Val_long(0);
   node->next = Val_unit;
 
   return node;
@@ -582,7 +595,7 @@ static c_node* allocate_c_node(void)
    call (e.g. [List.map] when not inlined). */
 static void* last_indirect_node_hole_ptr_callee;
 static value* last_indirect_node_hole_ptr_node_hole;
-static value* last_indirect_node_hole_ptr_result;
+static call_point* last_indirect_node_hole_ptr_result;
 
 CAMLprim value* caml_spacetime_indirect_node_hole_ptr
       (void* callee, value* node_hole, value caller_node)
@@ -596,7 +609,11 @@ CAMLprim value* caml_spacetime_indirect_node_hole_ptr
 
   if (callee == last_indirect_node_hole_ptr_callee
       && node_hole == last_indirect_node_hole_ptr_node_hole) {
-    return last_indirect_node_hole_ptr_result;
+#ifdef WITH_SPACETIME_CALL_COUNTS
+    last_indirect_node_hole_ptr_result->call_count =
+      Val_long (Long_val (last_indirect_node_hole_ptr_result->call_count) + 1);
+#endif
+    return &(last_indirect_node_hole_ptr_result->callee_node);
   }
 
   last_indirect_node_hole_ptr_callee = callee;
@@ -613,8 +630,12 @@ CAMLprim value* caml_spacetime_indirect_node_hole_ptr
     CAMLassert(caml_spacetime_classify_c_node(c_node) == CALL);
 
     if (c_node->pc == encoded_callee) {
-      last_indirect_node_hole_ptr_result = &(c_node->data.callee_node);
-      return last_indirect_node_hole_ptr_result;
+#ifdef WITH_SPACETIME_CALL_COUNTS
+      c_node->data.call.call_count =
+        Val_long (Long_val(c_node->data.call.call_count) + 1);
+#endif
+      last_indirect_node_hole_ptr_result = &(c_node->data.call);
+      return &(last_indirect_node_hole_ptr_result->callee_node);
     }
     else {
       node_hole = &c_node->next;
@@ -629,7 +650,7 @@ CAMLprim value* caml_spacetime_indirect_node_hole_ptr
        Perform the initialization equivalent to that emitted by
        [Spacetime.code_for_function_prologue] for direct tail call
        sites. */
-    c_node->data.callee_node = Encode_tail_caller_node(caller_node);
+    c_node->data.call.callee_node = Encode_tail_caller_node(caller_node);
   }
 
   *node_hole = caml_spacetime_stored_pointer_of_c_node(c_node);
@@ -637,9 +658,13 @@ CAMLprim value* caml_spacetime_indirect_node_hole_ptr
   CAMLassert(((uintnat) *node_hole) % sizeof(value) == 0);
   CAMLassert(*node_hole != Val_unit);
 
-  last_indirect_node_hole_ptr_result = &(c_node->data.callee_node);
+#ifdef WITH_SPACETIME_CALL_COUNTS
+  c_node->data.call.call_count =
+    Val_long (Long_val(c_node->data.call.call_count) + 1);
+#endif
+  last_indirect_node_hole_ptr_result = &(c_node->data.call);
 
-  return last_indirect_node_hole_ptr_result;
+  return &(last_indirect_node_hole_ptr_result->callee_node);
 }
 
 /* Some notes on why caml_call_gc doesn't need a distinguished node.
@@ -846,7 +871,7 @@ static NOINLINE void* find_trie_node_from_libunwind(int for_allocation,
 
     CAMLassert(caml_spacetime_classify_c_node(node) == expected_type);
     CAMLassert(pc_inside_c_node_matches(node, pc));
-    node_hole = &node->data.callee_node;
+    node_hole = &node->data.call.callee_node;
   }
 
   if (must_initialise_node_for_allocation) {
diff --git a/asmrun/spacetime_offline.c b/asmrun/spacetime_offline.c
index 6f79e49fa3..d8e815cefa 100644
--- a/asmrun/spacetime_offline.c
+++ b/asmrun/spacetime_offline.c
@@ -144,6 +144,12 @@ CAMLprim value caml_spacetime_ocaml_direct_call_point_callee_node
   return Direct_callee_node(node, Long_val(offset));
 }
 
+CAMLprim value caml_spacetime_ocaml_direct_call_point_call_count
+(value node, value offset)
+{
+  return Direct_call_count(node, Long_val(offset));
+}
+
 CAMLprim value caml_spacetime_ocaml_indirect_call_point_callees
       (value node, value offset)
 {
@@ -197,10 +203,23 @@ CAMLprim value caml_spacetime_c_node_callee_node(value node)
   /* This might be an uninitialised tail call point: for example if an OCaml
      callee was indirectly called but the callee wasn't instrumented (e.g. a
      leaf function that doesn't allocate). */
-  if (Is_tail_caller_node_encoded(c_node->data.callee_node)) {
+  if (Is_tail_caller_node_encoded(c_node->data.call.callee_node)) {
     return Val_unit;
   }
-  return c_node->data.callee_node;
+  return c_node->data.call.callee_node;
+}
+
+CAMLprim value caml_spacetime_c_node_call_count(value node)
+{
+  c_node* c_node;
+  CAMLassert(node != (value) NULL);
+  CAMLassert(Is_c_node(node));
+  c_node = caml_spacetime_offline_c_node_of_stored_pointer_not_null(node);
+  CAMLassert(caml_spacetime_offline_classify_c_node(c_node) == CALL);
+  if (Is_tail_caller_node_encoded(c_node->data.call.callee_node)) {
+    return Val_long(0);
+  }
+  return c_node->data.call.call_count;
 }
 
 CAMLprim value caml_spacetime_c_node_profinfo(value node)