ext/calleree/calleree.c in calleree-0.1.0 vs ext/calleree/calleree.c in calleree-0.2.0

@@ -1,60 +1,81 @@
 #include "ruby/ruby.h"
 #include "ruby/debug.h"
 #include "ruby/st.h"
 
+#define USE_LOCINDEX 0 // Ruby 3.1 feature trial
+
 VALUE rb_mCalleree;
 
 struct eree_data {
     VALUE tp; // TracePoint.new(:call)
-    VALUE files; // name -> [ {[fname, 1], ... } ]
-    struct st_table* caller_callee;  // [posnum | posnum] -> cnt
+    VALUE files; // name -> [ nil, posnum1, posnum,2 ... ]
+    struct st_table* caller_callee;  // { (posnumX | posnumY) => cnt, ... }
     int last_posnum;
 } eree_data;
 
-static int
+#if !USE_LOCINDEX
+static unsigned int
 posnum(struct eree_data *data, VALUE path, int line)
 {
     VALUE lines, posnumv;
 
     if (NIL_P(lines = rb_hash_aref(data->files, path))) {
         rb_hash_aset(data->files, path, lines = rb_ary_new());
     }
     if (NIL_P(posnumv = rb_ary_entry(lines, line))) {
-        rb_ary_store(lines, line, posnumv = INT2FIX(data->last_posnum++));
+        int posnum = ++data->last_posnum;
+        if (RB_UNLIKELY(posnum == 0)) rb_raise(rb_eRuntimeError, "posnum overflow");
+        rb_ary_store(lines, line, posnumv = INT2FIX(posnum));
     }
-    return FIX2INT(posnumv);
+    return (unsigned int)FIX2INT(posnumv);
 }
+#endif
 
 static int
 increment_i(st_data_t *key, st_data_t *value, st_data_t arg, int existing)
 {
     *value = (existing ? *value : 0) + 1;
     return ST_CONTINUE;
 }
 
 static void
-erree_called(VALUE tpval, void *ptr)
+eree_called(VALUE tpval, void *ptr)
 {
     struct eree_data *data = ptr;
+
+#if !USE_LOCINDEX
     VALUE frames[2];
     int lines[2];
     // int start, int limit, VALUE *buff, int *lines);
-    rb_profile_frames(1, 2, frames, lines);
-    int callee_posnum = posnum(data, rb_profile_frame_path(frames[0]), lines[0]);
-    int caller_posnum = posnum(data, rb_profile_frame_path(frames[1]), lines[1]);
-    // rb_p(rb_profile_frame_path(frames[1]));
-    st_data_t eree_pair = (((VALUE)caller_posnum << 32) | ((VALUE)callee_posnum));
+    int r = rb_profile_frames(1, 2, frames, lines);
 
-    st_update(data->caller_callee, eree_pair, increment_i, 0);
+    if (r >= 2) {
+        unsigned int callee_posnum = posnum(data, rb_profile_frame_absolute_path(frames[0]), lines[0]);
+        unsigned int caller_posnum = posnum(data, rb_profile_frame_absolute_path(frames[1]), lines[1]);
+
+        st_data_t eree_pair = (((VALUE)caller_posnum << 32) | ((VALUE)callee_posnum));
+        st_update(data->caller_callee, eree_pair, increment_i, 0);
+    }
+#else
+    unsigned int locs[2];
+    int r = rb_profile_locindex(0, 2, locs, 0);
+    if (r >= 2) {
+        unsigned int callee_posnum = locs[0];// posnum(data, rb_profile_frame_absolute_path (frames[0]), lines[0]);
+        unsigned int caller_posnum = locs[1];// posnum(data, rb_profile_frame_absolute_path (frames[1]), lines[1]);
+        st_data_t eree_pair = (((VALUE)caller_posnum << 32) | ((VALUE)callee_posnum));
+        st_update(data->caller_callee, eree_pair, increment_i, 0);
+    }
+#endif
 }
 
 static VALUE
 eree_start(VALUE self)
 {
     if (!eree_data.tp) {
-        eree_data.tp = rb_tracepoint_new(0, RUBY_EVENT_CALL, erree_called, &eree_data);
+        eree_data.tp = rb_tracepoint_new(0, RUBY_EVENT_CALL, eree_called, &eree_data);
+        rb_gc_register_mark_object(eree_data.tp);
         eree_data.caller_callee = st_init_numtable();
     }
     rb_tracepoint_enable(eree_data.tp);
     return self;
 }
@@ -70,13 +91,25 @@
 
 static int
 raw_result_i(st_data_t key, st_data_t val, st_data_t data)
 {
     VALUE ary = (VALUE)data;
-    int callee_posnum = (int)(key & 0xffffffff);
-    int caller_posnum = (int)(key >> 32);
-    rb_ary_push(ary, rb_ary_new_from_args(3, INT2FIX(caller_posnum), INT2FIX(callee_posnum), INT2FIX((int)val)));
+    unsigned int callee_posnum = (unsigned int)(key & 0xffffffff);
+    unsigned int caller_posnum = (unsigned int)(key >> 32);
+#if !USE_LOCINDEX
+    rb_ary_push(ary, rb_ary_new_from_args(3, UINT2NUM(caller_posnum), UINT2NUM(callee_posnum), INT2FIX((int)val)));
+#else
+    VALUE er_path, ee_path;
+    int er_line, ee_line;
+    rb_locindex_resolve(callee_posnum, &ee_path, &ee_line);
+    rb_locindex_resolve(caller_posnum, &er_path, &er_line);
+
+    rb_ary_push(ary, rb_ary_new_from_args(3,
+                                          rb_ary_new_from_args(2, er_path, INT2FIX(er_line)),
+                                          rb_ary_new_from_args(2, ee_path, INT2FIX(ee_line)),
+                                          INT2FIX((int)val)));
+#endif
     return ST_CONTINUE;
 }
 
 static VALUE
 eree_raw_result(VALUE self, VALUE clear_p)
@@ -90,23 +123,29 @@
         }
     }
     return ary;
 }
 
+#if !USE_LOCINDEX
 static VALUE
 eree_raw_posmap(VALUE self)
 {
     return eree_data.files;
 }
+#endif
 
 void
 Init_calleree(void)
 {
     eree_data.files = rb_hash_new();
+    rb_funcall(eree_data.files, rb_intern("compare_by_identity"), 0);
     rb_gc_register_mark_object(eree_data.files);
 
     rb_mCalleree = rb_define_module("Calleree");
     rb_define_singleton_method(rb_mCalleree, "_start", eree_start, 0);
     rb_define_singleton_method(rb_mCalleree, "_stop", eree_stop, 0);
     rb_define_singleton_method(rb_mCalleree, "raw_result", eree_raw_result, 1);
+
+#if !USE_LOCINDEX
     rb_define_singleton_method(rb_mCalleree, "raw_posmap", eree_raw_posmap, 0);
+#endif
 }