feat(tracer): handle more Ruby types

Author: zah

gems/native-tracer/ext/native_tracer/src/lib.rs

@@ -6,6 +6,7 @@ use std::{
     os::raw::{c_char, c_void},
     path::Path,
     ptr,
+    collections::HashMap,
 };
 
 use rb_sys::{
@@ -21,9 +22,12 @@ use rb_sys::{
     rb_raise, rb_eIOError,
     rb_sym2id, rb_id2name, rb_id2sym, rb_obj_classname, rb_num2long
 };
-use rb_sys::{RARRAY_LEN, RARRAY_CONST_PTR, RSTRING_LEN, RSTRING_PTR, RB_INTEGER_TYPE_P, RB_TYPE_P, RB_SYMBOL_P, NIL_P};
+use rb_sys::{RARRAY_LEN, RARRAY_CONST_PTR, RSTRING_LEN, RSTRING_PTR, RB_INTEGER_TYPE_P, RB_TYPE_P, RB_SYMBOL_P, RB_FLOAT_TYPE_P, NIL_P};
 use rb_sys::{Qtrue, Qfalse, Qnil};
-use runtime_tracing::{Tracer, Line, ValueRecord, TypeKind, EventLogKind, TraceLowLevelEvent, CallRecord, FullValueRecord, ReturnRecord, RecordEvent};
+use runtime_tracing::{
+    Tracer, Line, ValueRecord, TypeKind, TypeSpecificInfo, TypeRecord, FieldTypeRecord, TypeId,
+    EventLogKind, TraceLowLevelEvent, CallRecord, FullValueRecord, ReturnRecord, RecordEvent,
+};
 
 #[repr(C)]
 struct RTypedData {
@@ -58,6 +62,15 @@ extern "C" {
         data_type: *const rb_data_type_t,
     ) -> VALUE;
     fn rb_check_typeddata(obj: VALUE, data_type: *const rb_data_type_t) -> *mut c_void;
+    fn rb_num2dbl(val: VALUE) -> f64;
+    fn rb_obj_is_kind_of(obj: VALUE, class: VALUE) -> VALUE;
+    fn rb_path2class(path: *const c_char) -> VALUE;
+    fn rb_const_defined(klass: VALUE, name: ID) -> VALUE;
+    fn rb_const_get(klass: VALUE, name: ID) -> VALUE;
+    static rb_cTime: VALUE;
+    static rb_cRegexp: VALUE;
+    static rb_cStruct: VALUE;
+    static rb_cRange: VALUE;
 }
 
 struct Recorder {
@@ -69,6 +82,66 @@ struct Recorder {
     inst_meth_id: ID,
     parameters_id: ID,
     class_id: ID,
+    struct_types: HashMap<String, runtime_tracing::TypeId>,
+}
+
+fn value_type_id(val: &ValueRecord) -> runtime_tracing::TypeId {
+    use ValueRecord::*;
+    match val {
+        Int { type_id, .. }
+        | Float { type_id, .. }
+        | Bool { type_id, .. }
+        | String { type_id, .. }
+        | Sequence { type_id, .. }
+        | Tuple { type_id, .. }
+        | Struct { type_id, .. }
+        | Variant { type_id, .. }
+        | Reference { type_id, .. }
+        | Raw { type_id, .. }
+        | Error { type_id, .. }
+        | None { type_id } => *type_id,
+        Cell { .. } => runtime_tracing::NONE_TYPE_ID,
+    }
+}
+
+unsafe fn struct_value(
+    recorder: &mut Recorder,
+    class_name: &str,
+    field_names: &[&str],
+    field_values: &[VALUE],
+    depth: usize,
+    to_s_id: ID,
+) -> ValueRecord {
+    let mut vals = Vec::new();
+    for v in field_values {
+        vals.push(to_value(recorder, *v, depth - 1, to_s_id));
+    }
+
+    let type_id = if let Some(id) = recorder.struct_types.get(class_name) {
+        *id
+    } else {
+        let field_types: Vec<FieldTypeRecord> = field_names
+            .iter()
+            .zip(&vals)
+            .map(|(n, v)| FieldTypeRecord {
+                name: (*n).to_string(),
+                type_id: value_type_id(v),
+            })
+            .collect();
+        let typ = TypeRecord {
+            kind: TypeKind::Struct,
+            lang_type: class_name.to_string(),
+            specific_info: TypeSpecificInfo::Struct { fields: field_types },
+        };
+        let id = recorder.tracer.ensure_raw_type_id(typ);
+        recorder.struct_types.insert(class_name.to_string(), id);
+        id
+    };
+
+    ValueRecord::Struct {
+        field_values: vals,
+        type_id,
+    }
 }
 
 unsafe extern "C" fn recorder_free(ptr: *mut c_void) {
@@ -122,6 +195,7 @@ unsafe extern "C" fn ruby_recorder_alloc(klass: VALUE) -> VALUE {
         inst_meth_id,
         parameters_id,
         class_id,
+        struct_types: HashMap::new(),
     });
     let ty = std::ptr::addr_of!(RECORDER_TYPE) as *const rb_data_type_t;
     rb_data_typed_object_wrap(klass, Box::into_raw(recorder) as *mut c_void, ty)
@@ -186,35 +260,40 @@ unsafe fn value_to_string(val: VALUE, to_s_id: ID) -> Option<String> {
     Some(String::from_utf8_lossy(slice).to_string())
 }
 
-unsafe fn to_value(tracer: &mut Tracer, val: VALUE, depth: usize, to_s_id: ID) -> ValueRecord {
+unsafe fn to_value(recorder: &mut Recorder, val: VALUE, depth: usize, to_s_id: ID) -> ValueRecord {
     if depth == 0 {
-        let type_id = tracer.ensure_type_id(TypeKind::Error, "No type");
+        let type_id = recorder.tracer.ensure_type_id(TypeKind::Error, "No type");
         return ValueRecord::None { type_id };
     }
     if NIL_P(val) {
-        let type_id = tracer.ensure_type_id(TypeKind::Error, "No type");
+        let type_id = recorder.tracer.ensure_type_id(TypeKind::Error, "No type");
         return ValueRecord::None { type_id };
     }
     if val == (Qtrue as VALUE) || val == (Qfalse as VALUE) {
-        let type_id = tracer.ensure_type_id(TypeKind::Bool, "Bool");
+        let type_id = recorder.tracer.ensure_type_id(TypeKind::Bool, "Bool");
         return ValueRecord::Bool { b: val == (Qtrue as VALUE), type_id };
     }
     if RB_INTEGER_TYPE_P(val) {
         let i = rb_num2long(val) as i64;
-        let type_id = tracer.ensure_type_id(TypeKind::Int, "Integer");
+        let type_id = recorder.tracer.ensure_type_id(TypeKind::Int, "Integer");
         return ValueRecord::Int { i, type_id };
     }
+    if RB_FLOAT_TYPE_P(val) {
+        let f = rb_num2dbl(val);
+        let type_id = recorder.tracer.ensure_type_id(TypeKind::Float, "Float");
+        return ValueRecord::Float { f, type_id };
+    }
     if RB_SYMBOL_P(val) {
         let id = rb_sym2id(val);
         let name = CStr::from_ptr(rb_id2name(id)).to_str().unwrap_or("");
-        let type_id = tracer.ensure_type_id(TypeKind::String, "Symbol");
+        let type_id = recorder.tracer.ensure_type_id(TypeKind::String, "Symbol");
         return ValueRecord::String { text: name.to_string(), type_id };
     }
     if RB_TYPE_P(val, rb_sys::ruby_value_type::RUBY_T_STRING) {
         let ptr = RSTRING_PTR(val);
         let len = RSTRING_LEN(val) as usize;
         let slice = std::slice::from_raw_parts(ptr as *const u8, len);
-        let type_id = tracer.ensure_type_id(TypeKind::String, "String");
+        let type_id = recorder.tracer.ensure_type_id(TypeKind::String, "String");
         return ValueRecord::String { text: String::from_utf8_lossy(slice).to_string(), type_id };
     }
     if RB_TYPE_P(val, rb_sys::ruby_value_type::RUBY_T_ARRAY) {
@@ -223,14 +302,106 @@ unsafe fn to_value(tracer: &mut Tracer, val: VALUE, depth: usize, to_s_id: ID) -
         let ptr = RARRAY_CONST_PTR(val);
         for i in 0..len {
             let elem = *ptr.add(i);
-            elements.push(to_value(tracer, elem, depth - 1, to_s_id));
+            elements.push(to_value(recorder, elem, depth - 1, to_s_id));
+        }
+        let type_id = recorder.tracer.ensure_type_id(TypeKind::Seq, "Array");
+        return ValueRecord::Sequence { elements, is_slice: false, type_id };
+    }
+    if RB_TYPE_P(val, rb_sys::ruby_value_type::RUBY_T_HASH) {
+        let pairs = rb_funcall(val, rb_intern(b"to_a\0".as_ptr() as *const c_char), 0);
+        let len = RARRAY_LEN(pairs) as usize;
+        let ptr = RARRAY_CONST_PTR(pairs);
+        let mut elements = Vec::new();
+        for i in 0..len {
+            let pair = *ptr.add(i);
+            if RARRAY_LEN(pair) < 2 { continue; }
+            let pair_ptr = RARRAY_CONST_PTR(pair);
+            let key = *pair_ptr.add(0);
+            let val_elem = *pair_ptr.add(1);
+            elements.push(struct_value(recorder, "Pair", &["k", "v"], &[key, val_elem], depth, to_s_id));
         }
-        let type_id = tracer.ensure_type_id(TypeKind::Seq, "Array");
+        let type_id = recorder.tracer.ensure_type_id(TypeKind::Seq, "Hash");
         return ValueRecord::Sequence { elements, is_slice: false, type_id };
     }
+    if rb_obj_is_kind_of(val, rb_cRange) != 0 {
+        let begin_val = rb_funcall(val, rb_intern(b"begin\0".as_ptr() as *const c_char), 0);
+        let end_val = rb_funcall(val, rb_intern(b"end\0".as_ptr() as *const c_char), 0);
+        return struct_value(recorder, "Range", &["begin", "end"], &[begin_val, end_val], depth, to_s_id);
+    }
+    let set_id = rb_intern(b"Set\0".as_ptr() as *const c_char);
+    if rb_const_defined(rb_cObject, set_id) != 0 {
+        let set_cls = rb_const_get(rb_cObject, set_id);
+        if rb_obj_is_kind_of(val, set_cls) != 0 {
+            let arr = rb_funcall(val, rb_intern(b"to_a\0".as_ptr() as *const c_char), 0);
+            let len = RARRAY_LEN(arr) as usize;
+            let ptr = RARRAY_CONST_PTR(arr);
+            let mut elements = Vec::new();
+            for i in 0..len {
+                let elem = *ptr.add(i);
+                elements.push(to_value(recorder, elem, depth - 1, to_s_id));
+            }
+            let type_id = recorder.tracer.ensure_type_id(TypeKind::Seq, "Set");
+            return ValueRecord::Sequence { elements, is_slice: false, type_id };
+        }
+    }
+    if rb_obj_is_kind_of(val, rb_cTime) != 0 {
+        let sec = rb_funcall(val, rb_intern(b"to_i\0".as_ptr() as *const c_char), 0);
+        let nsec = rb_funcall(val, rb_intern(b"nsec\0".as_ptr() as *const c_char), 0);
+        return struct_value(recorder, "Time", &["sec", "nsec"], &[sec, nsec], depth, to_s_id);
+    }
+    if rb_obj_is_kind_of(val, rb_cRegexp) != 0 {
+        let src = rb_funcall(val, rb_intern(b"source\0".as_ptr() as *const c_char), 0);
+        let opts = rb_funcall(val, rb_intern(b"options\0".as_ptr() as *const c_char), 0);
+        return struct_value(recorder, "Regexp", &["source", "options"], &[src, opts], depth, to_s_id);
+    }
+    if rb_obj_is_kind_of(val, rb_cStruct) != 0 {
+        let class_name = cstr_to_string(rb_obj_classname(val)).unwrap_or_else(|| "Struct".to_string());
+        let members = rb_funcall(val, rb_intern(b"members\0".as_ptr() as *const c_char), 0);
+        let values = rb_funcall(val, rb_intern(b"values\0".as_ptr() as *const c_char), 0);
+        let len = RARRAY_LEN(values) as usize;
+        let mem_ptr = RARRAY_CONST_PTR(members);
+        let val_ptr = RARRAY_CONST_PTR(values);
+        let mut names: Vec<&str> = Vec::new();
+        let mut vals: Vec<VALUE> = Vec::new();
+        for i in 0..len {
+            let sym = *mem_ptr.add(i);
+            let id = rb_sym2id(sym);
+            let cstr = rb_id2name(id);
+            let name = CStr::from_ptr(cstr).to_str().unwrap_or("?");
+            names.push(name);
+            vals.push(*val_ptr.add(i));
+        }
+        return struct_value(recorder, &class_name, &names, &vals, depth, to_s_id);
+    }
+    let open_struct_id = rb_intern(b"OpenStruct\0".as_ptr() as *const c_char);
+    if rb_const_defined(rb_cObject, open_struct_id) != 0 {
+        let open_struct = rb_const_get(rb_cObject, open_struct_id);
+        if rb_obj_is_kind_of(val, open_struct) != 0 {
+            let h = rb_funcall(val, rb_intern(b"to_h\0".as_ptr() as *const c_char), 0);
+            return to_value(recorder, h, depth - 1, to_s_id);
+        }
+    }
     let class_name = cstr_to_string(rb_obj_classname(val)).unwrap_or_else(|| "Object".to_string());
+    // generic object
+    let ivars = rb_funcall(val, rb_intern(b"instance_variables\0".as_ptr() as *const c_char), 0);
+    let len = RARRAY_LEN(ivars) as usize;
+    let ptr = RARRAY_CONST_PTR(ivars);
+    let mut names: Vec<&str> = Vec::new();
+    let mut vals: Vec<VALUE> = Vec::new();
+    for i in 0..len {
+        let sym = *ptr.add(i);
+        let id = rb_sym2id(sym);
+        let cstr = rb_id2name(id);
+        let name = CStr::from_ptr(cstr).to_str().unwrap_or("?");
+        names.push(name);
+        let value = rb_funcall(val, rb_intern(b"instance_variable_get\0".as_ptr() as *const c_char), 1, sym);
+        vals.push(value);
+    }
+    if !names.is_empty() {
+        return struct_value(recorder, &class_name, &names, &vals, depth, to_s_id);
+    }
     let text = value_to_string(val, to_s_id).unwrap_or_default();
-    let type_id = tracer.ensure_type_id(TypeKind::Raw, &class_name);
+    let type_id = recorder.tracer.ensure_type_id(TypeKind::Raw, &class_name);
     ValueRecord::Raw { r: text, type_id }
 }
 
@@ -244,7 +415,7 @@ unsafe fn record_variables(recorder: &mut Recorder, binding: VALUE) -> Vec<FullV
         let id = rb_sym2id(sym);
         let name = CStr::from_ptr(rb_id2name(id)).to_str().unwrap_or("");
         let value = rb_funcall(binding, recorder.local_get_id, 1, sym);
-        let val_rec = to_value(&mut recorder.tracer, value, 10, recorder.to_s_id);
+        let val_rec = to_value(recorder, value, 10, recorder.to_s_id);
         recorder.tracer.register_variable_with_full_value(name, val_rec.clone());
         let var_id = recorder.tracer.ensure_variable_id(name);
         result.push(FullValueRecord { variable_id: var_id, value: val_rec });
@@ -276,7 +447,7 @@ unsafe fn record_parameters(recorder: &mut Recorder, binding: VALUE, defined_cla
         }
         let name = CStr::from_ptr(name_c).to_str().unwrap_or("");
         let value = rb_funcall(binding, recorder.local_get_id, 1, name_sym);
-        let val_rec = to_value(&mut recorder.tracer, value, 10, recorder.to_s_id);
+        let val_rec = to_value(recorder, value, 10, recorder.to_s_id);
         if register {
             recorder.tracer.register_variable_with_full_value(name, val_rec.clone());
             let var_id = recorder.tracer.ensure_variable_id(name);
@@ -378,7 +549,7 @@ unsafe extern "C" fn event_hook_raw(data: VALUE, arg: *mut rb_trace_arg_t) {
             // register parameter types first
             let _ = record_parameters(recorder, binding, defined_class, mid, false);
         }
-        let self_rec = to_value(&mut recorder.tracer, self_val, 10, recorder.to_s_id);
+        let self_rec = to_value(recorder, self_val, 10, recorder.to_s_id);
         recorder
             .tracer
             .register_variable_with_full_value("self", self_rec.clone());
@@ -405,7 +576,7 @@ unsafe extern "C" fn event_hook_raw(data: VALUE, arg: *mut rb_trace_arg_t) {
     } else if (ev & RUBY_EVENT_RETURN) != 0 {
         recorder.tracer.register_step(Path::new(&path), Line(line));
         let ret = rb_tracearg_return_value(arg);
-        let val_rec = to_value(&mut recorder.tracer, ret, 10, recorder.to_s_id);
+        let val_rec = to_value(recorder, ret, 10, recorder.to_s_id);
         recorder.tracer.register_variable_with_full_value("<return_value>", val_rec.clone());
         recorder.tracer.events.push(TraceLowLevelEvent::Return(ReturnRecord { return_value: val_rec }));
     } else if (ev & RUBY_EVENT_RAISE) != 0 {