memory-profiler 1.3.0 → 1.5.0

data/ext/memory/profiler/capture.c

@@ -4,10 +4,10 @@
 #include "capture.h"
 #include "allocations.h"
 #include "events.h"
+#include "table.h"
 
-#include "ruby.h"
-#include "ruby/debug.h"
-#include "ruby/st.h"
+#include <ruby/debug.h>
+#include <ruby/st.h>
 #include <stdatomic.h>
 #include <stdio.h>
 
@@ -20,6 +20,7 @@ static VALUE Memory_Profiler_Capture = Qnil;
 // Event symbols:
 static VALUE sym_newobj, sym_freeobj;
 
+
 // Main capture state (per-instance).
 struct Memory_Profiler_Capture {
 	// Master switch - is tracking active? (set by start/stop).
@@ -29,15 +30,19 @@ struct Memory_Profiler_Capture {
 	int paused;
 
 	// Tracked classes: class => VALUE (wrapped Memory_Profiler_Capture_Allocations).
-	st_table *tracked_classes;
+	st_table *tracked;
+	
+	// Custom object table: object (address) => state hash
+	// Uses system malloc (GC-safe), updates addresses during compaction
+	struct Memory_Profiler_Object_Table *states;
 
 	// Total number of allocations and frees seen since tracking started.
 	size_t new_count;
 	size_t free_count;
 };
 
-// GC mark callback for tracked_classes table.
-static int Memory_Profiler_Capture_tracked_classes_mark(st_data_t key, st_data_t value, st_data_t arg) {
+// GC mark callback for tracked table.
+static int Memory_Profiler_Capture_tracked_mark(st_data_t key, st_data_t value, st_data_t arg) {
 	// Mark class as un-movable:
 	// - We don't want to re-index the table if the class moves.
 	// - We don't want objects in `freeobj` to have invalid class pointers (maybe helps).
@@ -54,16 +59,22 @@ static int Memory_Profiler_Capture_tracked_classes_mark(st_data_t key, st_data_t
 static void Memory_Profiler_Capture_mark(void *ptr) {
 	struct Memory_Profiler_Capture *capture = ptr;
 	
-	if (capture->tracked_classes) {
-		st_foreach(capture->tracked_classes, Memory_Profiler_Capture_tracked_classes_mark, 0);
+	if (capture->tracked) {
+		st_foreach(capture->tracked, Memory_Profiler_Capture_tracked_mark, 0);
 	}
+	
+	Memory_Profiler_Object_Table_mark(capture->states);
 }
 
 static void Memory_Profiler_Capture_free(void *ptr) {
 	struct Memory_Profiler_Capture *capture = ptr;
 		
-	if (capture->tracked_classes) {
-		st_free_table(capture->tracked_classes);
+	if (capture->tracked) {
+		st_free_table(capture->tracked);
+	}
+	
+	if (capture->states) {
+		Memory_Profiler_Object_Table_free(capture->states);
 	}
 	
 	xfree(capture);
@@ -73,20 +84,20 @@ static size_t Memory_Profiler_Capture_memsize(const void *ptr) {
 	const struct Memory_Profiler_Capture *capture = ptr;
 	size_t size = sizeof(struct Memory_Profiler_Capture);
 	
-	if (capture->tracked_classes) {
-		size += capture->tracked_classes->num_entries * (sizeof(st_data_t) + sizeof(struct Memory_Profiler_Capture_Allocations));
+	if (capture->tracked) {
+		size += capture->tracked->num_entries * (sizeof(st_data_t) + sizeof(struct Memory_Profiler_Capture_Allocations));
 	}
 	
 	return size;
 }
 
 // Foreach callback for st_foreach_with_replace (iteration logic).
-static int Memory_Profiler_Capture_tracked_classes_foreach(st_data_t key, st_data_t value, st_data_t argp, int error) {
+static int Memory_Profiler_Capture_tracked_foreach(st_data_t key, st_data_t value, st_data_t argp, int error) {
 	return ST_REPLACE;
 }
 
 // Replace callback for st_foreach_with_replace (update logic).
-static int Memory_Profiler_Capture_tracked_classes_update(st_data_t *key, st_data_t *value, st_data_t argp, int existing) {
+static int Memory_Profiler_Capture_tracked_update(st_data_t *key, st_data_t *value, st_data_t argp, int existing) {
 	// Update wrapped Allocations VALUE if it moved:
 	VALUE old_allocations = (VALUE)*value;
 	VALUE new_allocations = rb_gc_location(old_allocations);
@@ -100,12 +111,17 @@ static int Memory_Profiler_Capture_tracked_classes_update(st_data_t *key, st_dat
 static void Memory_Profiler_Capture_compact(void *ptr) {
 	struct Memory_Profiler_Capture *capture = ptr;
 	
-	// Update tracked_classes keys and allocations values in-place:
-	if (capture->tracked_classes && capture->tracked_classes->num_entries > 0) {
-		if (st_foreach_with_replace(capture->tracked_classes, Memory_Profiler_Capture_tracked_classes_foreach, Memory_Profiler_Capture_tracked_classes_update, 0)) {
-			rb_raise(rb_eRuntimeError, "tracked_classes modified during GC compaction");
+	// Update tracked keys and allocations values in-place:
+	if (capture->tracked && capture->tracked->num_entries > 0) {
+		if (st_foreach_with_replace(capture->tracked, Memory_Profiler_Capture_tracked_foreach, Memory_Profiler_Capture_tracked_update, 0)) {
+			rb_raise(rb_eRuntimeError, "tracked modified during GC compaction");
 		}
 	}
+	
+	// Update custom object table (system malloc, safe during GC)
+	if (capture->states) {
+		Memory_Profiler_Object_Table_compact(capture->states);
+	}
 }
 
 static const rb_data_type_t Memory_Profiler_Capture_type = {
@@ -138,6 +154,9 @@ const char *event_flag_name(rb_event_flag_t event_flag) {
 }
 
 // Process a NEWOBJ event. All allocation tracking logic is here.
+// object_id parameter is the Integer object_id, NOT the raw object.
+// Process a NEWOBJ event. All allocation tracking logic is here.
+// object parameter is the actual object being allocated.
 static void Memory_Profiler_Capture_process_newobj(VALUE self, VALUE klass, VALUE object) {
 	struct Memory_Profiler_Capture *capture;
 	TypedData_Get_Struct(self, struct Memory_Profiler_Capture, &Memory_Profiler_Capture_type, capture);
@@ -153,7 +172,7 @@ static void Memory_Profiler_Capture_process_newobj(VALUE self, VALUE klass, VALU
 	VALUE allocations;
 	struct Memory_Profiler_Capture_Allocations *record;
 	
-	if (st_lookup(capture->tracked_classes, (st_data_t)klass, &allocations_data)) {
+	if (st_lookup(capture->tracked, (st_data_t)klass, &allocations_data)) {
 		// Existing record
 		allocations = (VALUE)allocations_data;
 		record = Memory_Profiler_Allocations_get(allocations);
@@ -164,65 +183,66 @@ static void Memory_Profiler_Capture_process_newobj(VALUE self, VALUE klass, VALU
 		record->callback = Qnil;
 		record->new_count = 1;
 		record->free_count = 0;
-		record->states = st_init_numtable();
 		
 		allocations = Memory_Profiler_Allocations_wrap(record);
-		st_insert(capture->tracked_classes, (st_data_t)klass, (st_data_t)allocations);
+		st_insert(capture->tracked, (st_data_t)klass, (st_data_t)allocations);
 		RB_OBJ_WRITTEN(self, Qnil, klass);
 		RB_OBJ_WRITTEN(self, Qnil, allocations);
 	}
 	
-	// Only store state if there's a callback
+	VALUE data = Qnil;
 	if (!NIL_P(record->callback)) {
-		VALUE state = rb_funcall(record->callback, rb_intern("call"), 3, klass, sym_newobj, Qnil);
-		
-		// Store state using object as key (works because object is alive):
-		st_insert(record->states, (st_data_t)object, (st_data_t)state);
-	} else {
-		// Store state as nil:
-		st_insert(record->states, (st_data_t)object, (st_data_t)Qnil);
+		data = rb_funcall(record->callback, rb_intern("call"), 3, klass, sym_newobj, Qnil);
 	}
-
+	
+	struct Memory_Profiler_Object_Table_Entry *entry = Memory_Profiler_Object_Table_insert(capture->states, object);
+	RB_OBJ_WRITTEN(self, Qnil, object);
+	RB_OBJ_WRITE(self, &entry->klass, klass);
+	RB_OBJ_WRITE(self, &entry->data, data);
+	RB_OBJ_WRITE(self, &entry->allocations, allocations);
+	
+	if (DEBUG) fprintf(stderr, "[NEWOBJ] Object inserted into table: %p\n", (void*)object);
+	
 	// Resume the capture:
 	capture->paused -= 1;
 }
 
 // Process a FREEOBJ event. All deallocation tracking logic is here.
-static void Memory_Profiler_Capture_process_freeobj(VALUE capture_value, VALUE klass, VALUE object) {
+// freeobj_data parameter is [state_hash, object] array from event handler.
+static void Memory_Profiler_Capture_process_freeobj(VALUE capture_value, VALUE unused_klass, VALUE object) {
 	struct Memory_Profiler_Capture *capture;
 	TypedData_Get_Struct(capture_value, struct Memory_Profiler_Capture, &Memory_Profiler_Capture_type, capture);
 	
 	// Pause the capture to prevent infinite loop:
 	capture->paused += 1;
 	
-	// Look up allocations record for this class
-	st_data_t allocations_data;
-	if (!st_lookup(capture->tracked_classes, (st_data_t)klass, &allocations_data)) {
-		// The class is not tracked, so we are not tracking this object.
+	struct Memory_Profiler_Object_Table_Entry *entry = Memory_Profiler_Object_Table_lookup(capture->states, object);
+	
+	if (!entry) {
+		if (DEBUG) fprintf(stderr, "[FREEOBJ] Object not found in table: %p\n", (void*)object);
 		goto done;
+	} else {
+		if (DEBUG) fprintf(stderr, "[FREEOBJ] Object found in table: %p\n", (void*)object);
 	}
 	
-	VALUE allocations = (VALUE)allocations_data;
+	VALUE klass = entry->klass;
+	VALUE data = entry->data;
+	VALUE allocations = entry->allocations;
+	
+	// Delete by entry pointer (faster - no second lookup!)
+	Memory_Profiler_Object_Table_delete_entry(capture->states, entry);
+	
 	struct Memory_Profiler_Capture_Allocations *record = Memory_Profiler_Allocations_get(allocations);
 	
-	if (!record->states) {
-		// There is no state table for this class, so we are not tracking it.
-		goto done;
-	}
+	// Increment global free count
+	capture->free_count++;
 	
-	st_data_t state_data;
-	if (st_delete(record->states, (st_data_t *)&object, &state_data)) {
-		VALUE state = (VALUE)state_data;
-		
-		// Increment global free count
-		capture->free_count++;
-		
-		// Increment per-class free count
-		record->free_count++;
-		
-		if (!NIL_P(record->callback) && !NIL_P(state)) {
-			rb_funcall(record->callback, rb_intern("call"), 3, klass, sym_freeobj, state);
-		}
+	// Increment per-class free count
+	record->free_count++;
+	
+	// Call callback if present
+	if (!NIL_P(record->callback) && !NIL_P(data)) {
+		rb_funcall(record->callback, rb_intern("call"), 3, klass, sym_freeobj, data);
 	}
 
 done:
@@ -283,11 +303,11 @@ int Memory_Profiler_Capture_trackable_p(VALUE object) {
 
 // Event hook callback with RAW_ARG
 // Signature: (VALUE data, rb_trace_arg_t *trace_arg)
-static void Memory_Profiler_Capture_event_callback(VALUE data, void *ptr) {
+static void Memory_Profiler_Capture_event_callback(VALUE self, void *ptr) {
 	rb_trace_arg_t *trace_arg = (rb_trace_arg_t *)ptr;
 	
 	struct Memory_Profiler_Capture *capture;
-	TypedData_Get_Struct(data, struct Memory_Profiler_Capture, &Memory_Profiler_Capture_type, capture);
+	TypedData_Get_Struct(self, struct Memory_Profiler_Capture, &Memory_Profiler_Capture_type, capture);
 	
 	VALUE object = rb_tracearg_object(trace_arg);
 	
@@ -295,38 +315,23 @@ static void Memory_Profiler_Capture_event_callback(VALUE data, void *ptr) {
 	if (!Memory_Profiler_Capture_trackable_p(object)) return;
 	
 	rb_event_flag_t event_flag = rb_tracearg_event_flag(trace_arg);
-	VALUE klass = rb_class_of(object);
-	if (!klass) return;
 	
-	if (DEBUG) {
-		const char *klass_name = "(ignored)";
-		if (event_flag == RUBY_INTERNAL_EVENT_NEWOBJ) {
-			klass_name = rb_class2name(klass);
-		}
-		fprintf(stderr, "Memory_Profiler_Capture_event_callback: object=%p, event_flag=%s, klass=%s\n", (void*)object, event_flag_name(event_flag), klass_name);
-	}
-
 	if (event_flag == RUBY_INTERNAL_EVENT_NEWOBJ) {
 		// Skip NEWOBJ if disabled (during callback) to prevent infinite recursion
 		if (capture->paused) return;
 		
-		// Check if class is already freed (T_NONE). This can happen when:
-		// It's safe to unconditionally call here:
-		object = rb_obj_id(object);
-
-		Memory_Profiler_Events_enqueue(MEMORY_PROFILER_EVENT_TYPE_NEWOBJ, data, klass, object);
+		VALUE klass = rb_obj_class(object);
+		
+		// Skip if klass is not a Class
+		if (rb_type(klass) != RUBY_T_CLASS) return;
+		
+		// Enqueue actual object (not object_id) - queue retains it until processed
+		// Ruby 3.5 compatible: no need for FL_SEEN_OBJ_ID or rb_obj_id
+		if (DEBUG) fprintf(stderr, "[NEWOBJ] Enqueuing event for object: %p\n", (void*)object);
+		Memory_Profiler_Events_enqueue(MEMORY_PROFILER_EVENT_TYPE_NEWOBJ, self, klass, object);
 	} else if (event_flag == RUBY_INTERNAL_EVENT_FREEOBJ) {
-		// We only care about objects that have been seen before (i.e. have an object ID):
-		if (RB_FL_TEST(object, FL_SEEN_OBJ_ID)) {
-			// For freeobj, we only care about klasses that we are tracking.
-			// This prevents us from enqueuing klass objects that might be freed.
-			if (!st_lookup(capture->tracked_classes, (st_data_t)klass, NULL)) return;
-			
-			// It's only safe to call here if the object already has an object ID.
-			object = rb_obj_id(object);
-			
-			Memory_Profiler_Events_enqueue(MEMORY_PROFILER_EVENT_TYPE_FREEOBJ, data, klass, object);
-		}
+		if (DEBUG) fprintf(stderr, "[FREEOBJ] Enqueuing event for object: %p\n", (void*)object);
+		Memory_Profiler_Events_enqueue(MEMORY_PROFILER_EVENT_TYPE_FREEOBJ, self, Qnil, object);
 	}
 }
 
@@ -339,10 +344,17 @@ static VALUE Memory_Profiler_Capture_alloc(VALUE klass) {
 		rb_raise(rb_eRuntimeError, "Failed to allocate Memory::Profiler::Capture");
 	}
 	
-	capture->tracked_classes = st_init_numtable();
+	capture->tracked = st_init_numtable();
 	
-	if (!capture->tracked_classes) {
-		rb_raise(rb_eRuntimeError, "Failed to initialize hash table");
+	if (!capture->tracked) {
+		rb_raise(rb_eRuntimeError, "Failed to initialize tracked hash table");
+	}
+	
+	// Initialize custom object table (uses system malloc, GC-safe)
+	capture->states = Memory_Profiler_Object_Table_new(1024);
+	if (!capture->states) {
+		st_free_table(capture->tracked);
+		rb_raise(rb_eRuntimeError, "Failed to initialize object table");
 	}
 	
 	// Initialize allocation tracking counters
@@ -424,7 +436,7 @@ static VALUE Memory_Profiler_Capture_track(int argc, VALUE *argv, VALUE self) {
 	st_data_t allocations_data;
 	VALUE allocations;
 	
-	if (st_lookup(capture->tracked_classes, (st_data_t)klass, &allocations_data)) {
+	if (st_lookup(capture->tracked, (st_data_t)klass, &allocations_data)) {
 		allocations = (VALUE)allocations_data;
 		struct Memory_Profiler_Capture_Allocations *record = Memory_Profiler_Allocations_get(allocations);
 		RB_OBJ_WRITE(self, &record->callback, callback);
@@ -433,12 +445,12 @@ static VALUE Memory_Profiler_Capture_track(int argc, VALUE *argv, VALUE self) {
 		RB_OBJ_WRITE(self, &record->callback, callback);
 		record->new_count = 0;
 		record->free_count = 0;
-		record->states = st_init_numtable();
+		// NOTE: States table removed - now at Capture level
 		
 		// Wrap the record in a VALUE
 		allocations = Memory_Profiler_Allocations_wrap(record);
 		
-		st_insert(capture->tracked_classes, (st_data_t)klass, (st_data_t)allocations);
+		st_insert(capture->tracked, (st_data_t)klass, (st_data_t)allocations);
 		RB_OBJ_WRITTEN(self, Qnil, klass);
 		RB_OBJ_WRITTEN(self, Qnil, allocations);
 	}
@@ -452,7 +464,7 @@ static VALUE Memory_Profiler_Capture_untrack(VALUE self, VALUE klass) {
 	TypedData_Get_Struct(self, struct Memory_Profiler_Capture, &Memory_Profiler_Capture_type, capture);
 	
 	st_data_t allocations_data;
-	if (st_delete(capture->tracked_classes, (st_data_t *)&klass, &allocations_data)) {
+	if (st_delete(capture->tracked, (st_data_t *)&klass, &allocations_data)) {
 		// The wrapped Allocations VALUE will be GC'd naturally
 		// No manual cleanup needed
 	}
@@ -465,7 +477,7 @@ static VALUE Memory_Profiler_Capture_tracking_p(VALUE self, VALUE klass) {
 	struct Memory_Profiler_Capture *capture;
 	TypedData_Get_Struct(self, struct Memory_Profiler_Capture, &Memory_Profiler_Capture_type, capture);
 	
-	return st_lookup(capture->tracked_classes, (st_data_t)klass, NULL) ? Qtrue : Qfalse;
+	return st_lookup(capture->tracked, (st_data_t)klass, NULL) ? Qtrue : Qfalse;
 }
 
 // Get count of live objects for a specific class (O(1) lookup!)
@@ -474,7 +486,7 @@ static VALUE Memory_Profiler_Capture_retained_count_of(VALUE self, VALUE klass)
 	TypedData_Get_Struct(self, struct Memory_Profiler_Capture, &Memory_Profiler_Capture_type, capture);
 	
 	st_data_t allocations_data;
-	if (st_lookup(capture->tracked_classes, (st_data_t)klass, &allocations_data)) {
+	if (st_lookup(capture->tracked, (st_data_t)klass, &allocations_data)) {
 		VALUE allocations = (VALUE)allocations_data;
 		struct Memory_Profiler_Capture_Allocations *record = Memory_Profiler_Allocations_get(allocations);
 		if (record->free_count <= record->new_count) {
@@ -486,7 +498,7 @@ static VALUE Memory_Profiler_Capture_retained_count_of(VALUE self, VALUE klass)
 }
 
 // Iterator to reset each class record
-static int Memory_Profiler_Capture_tracked_classes_clear(st_data_t key, st_data_t value, st_data_t arg) {
+static int Memory_Profiler_Capture_tracked_clear(st_data_t key, st_data_t value, st_data_t arg) {
 	VALUE allocations = (VALUE)value;
 	
 	Memory_Profiler_Allocations_clear(allocations);
@@ -507,8 +519,14 @@ static VALUE Memory_Profiler_Capture_clear(VALUE self) {
 		rb_raise(rb_eRuntimeError, "Cannot clear while capture is running - call stop() first!");
 	}
 	
-	// Reset all counts to 0 (don't free, just reset) - pass self for write barriers:
-	st_foreach(capture->tracked_classes, Memory_Profiler_Capture_tracked_classes_clear, 0);
+	// Reset all counts to 0 (don't free, just reset):
+	st_foreach(capture->tracked, Memory_Profiler_Capture_tracked_clear, 0);
+	
+	// Clear custom object table by recreating it
+	if (capture->states) {
+		Memory_Profiler_Object_Table_free(capture->states);
+		capture->states = Memory_Profiler_Object_Table_new(1024);
+	}
 	
 	// Reset allocation tracking counters
 	capture->new_count = 0;
@@ -535,18 +553,118 @@ static VALUE Memory_Profiler_Capture_each(VALUE self) {
 	
 	RETURN_ENUMERATOR(self, 0, 0);
 	
-	st_foreach(capture->tracked_classes, Memory_Profiler_Capture_each_allocation, 0);
+	st_foreach(capture->tracked, Memory_Profiler_Capture_each_allocation, 0);
 	
 	return self;
 }
 
+// Struct for filtering states during each_object iteration
+struct Memory_Profiler_Each_Object_Arguments {
+	VALUE self;
+	
+	// The allocations wrapper to filter by (Qnil = no filter).
+	VALUE allocations;
+};
+
+// Cleanup function to ensure table is made weak again
+static VALUE Memory_Profiler_Capture_each_object_ensure(VALUE arg) {
+	struct Memory_Profiler_Each_Object_Arguments *arguments = (struct Memory_Profiler_Each_Object_Arguments *)arg;
+	struct Memory_Profiler_Capture *capture;
+	TypedData_Get_Struct(arguments->self, struct Memory_Profiler_Capture, &Memory_Profiler_Capture_type, capture);
+	
+	// Make table weak again
+	Memory_Profiler_Object_Table_decrement_strong(capture->states);
+	
+	return Qnil;
+}
+
+// Main iteration function
+static VALUE Memory_Profiler_Capture_each_object_body(VALUE arg) {
+	struct Memory_Profiler_Each_Object_Arguments *arguments = (struct Memory_Profiler_Each_Object_Arguments *)arg;
+	struct Memory_Profiler_Capture *capture;
+	TypedData_Get_Struct(arguments->self, struct Memory_Profiler_Capture, &Memory_Profiler_Capture_type, capture);
+	
+	// Iterate custom object table entries
+	if (capture->states) {
+		if (DEBUG) fprintf(stderr, "[ITER] Iterating table, capacity=%zu, count=%zu\n", capture->states->capacity, capture->states->count);
+		
+		for (size_t i = 0; i < capture->states->capacity; i++) {
+			struct Memory_Profiler_Object_Table_Entry *entry = &capture->states->entries[i];
+			
+			// Skip empty or deleted slots (0 = not set)
+			if (entry->object == 0) {
+				continue;
+			}
+			
+			// Filter by allocations if specified
+			if (!NIL_P(arguments->allocations)) {
+				if (entry->allocations != arguments->allocations) continue;
+			}
+			
+			rb_yield_values(2, entry->object, entry->allocations);
+		}
+	}
+	
+	return arguments->self;
+}
+
+// Iterate over tracked object IDs, optionally filtered by class
+// Called as: 
+//   capture.each_object(String) { |object_id, state| ... }  # Specific class
+//   capture.each_object { |object_id, state| ... }          // All objects
+// 
+// Yields object_id as Integer. Caller can:
+//   - Format as hex: "0x%x" % object_id
+//   - Convert to object with ObjectSpace._id2ref (may raise RangeError if recycled)
+// Future-proof for Ruby 3.5 where _id2ref is deprecated
+static VALUE Memory_Profiler_Capture_each_object(int argc, VALUE *argv, VALUE self) {
+	struct Memory_Profiler_Capture *capture;
+	TypedData_Get_Struct(self, struct Memory_Profiler_Capture, &Memory_Profiler_Capture_type, capture);
+	
+	VALUE klass;
+	rb_scan_args(argc, argv, "01", &klass);
+	
+	RETURN_ENUMERATOR(self, argc, argv);
+	
+	// Make table strong so objects won't be collected during iteration
+	Memory_Profiler_Object_Table_increment_strong(capture->states);
+	
+	// Process all pending events and run GC to clean up stale objects:
+	Memory_Profiler_Events_process_all();
+	
+	// If class provided, look up its allocations wrapper
+	VALUE allocations = Qnil;
+	if (!NIL_P(klass)) {
+		st_data_t allocations_data;
+		if (st_lookup(capture->tracked, (st_data_t)klass, &allocations_data)) {
+			allocations = (VALUE)allocations_data;
+		} else {
+			// Class not tracked - nothing to iterate
+			Memory_Profiler_Object_Table_decrement_strong(capture->states);
+			return self;
+		}
+	}
+	
+	// Setup arguments for iteration
+	struct Memory_Profiler_Each_Object_Arguments arguments = {
+		.self = self,
+		.allocations = allocations
+	};
+	
+	// Use rb_ensure to guarantee cleanup even if exception is raised
+	return rb_ensure(
+		Memory_Profiler_Capture_each_object_body, (VALUE)&arguments,
+		Memory_Profiler_Capture_each_object_ensure, (VALUE)&arguments
+	);
+}
+
 // Get allocations for a specific class
 static VALUE Memory_Profiler_Capture_aref(VALUE self, VALUE klass) {
 	struct Memory_Profiler_Capture *capture;
 	TypedData_Get_Struct(self, struct Memory_Profiler_Capture, &Memory_Profiler_Capture_type, capture);
 	
 	st_data_t allocations_data;
-	if (st_lookup(capture->tracked_classes, (st_data_t)klass, &allocations_data)) {
+	if (st_lookup(capture->tracked, (st_data_t)klass, &allocations_data)) {
 		return (VALUE)allocations_data;
 	}
 	
@@ -559,20 +677,6 @@ struct Memory_Profiler_Allocations_Statistics {
 	VALUE per_class_counts;
 };
 
-// Iterator callback to count states per class
-static int Memory_Profiler_Capture_count_states(st_data_t key, st_data_t value, st_data_t argument) {
-	struct Memory_Profiler_Allocations_Statistics *statistics = (struct Memory_Profiler_Allocations_Statistics *)argument;
-	VALUE klass = (VALUE)key;
-	VALUE allocations = (VALUE)value;
-	struct Memory_Profiler_Capture_Allocations *record = Memory_Profiler_Allocations_get(allocations);
-	
-	size_t states_count = record->states ? record->states->num_entries : 0;
-	statistics->total_tracked_objects += states_count;
-	
-	rb_hash_aset(statistics->per_class_counts, klass, SIZET2NUM(states_count));
-	return ST_CONTINUE;
-}
-
 // Get internal statistics for debugging
 // Returns hash with internal state sizes
 static VALUE Memory_Profiler_Capture_statistics(VALUE self) {
@@ -582,18 +686,11 @@ static VALUE Memory_Profiler_Capture_statistics(VALUE self) {
 	VALUE statistics = rb_hash_new();
 	
 	// Tracked classes count
-	rb_hash_aset(statistics, ID2SYM(rb_intern("tracked_classes_count")), SIZET2NUM(capture->tracked_classes->num_entries));
-		
-	// Count states entries for each tracked class
-	struct Memory_Profiler_Allocations_Statistics allocations_statistics = {
-		.total_tracked_objects = 0,
-		.per_class_counts = rb_hash_new()
-	};
-	
-	st_foreach(capture->tracked_classes, Memory_Profiler_Capture_count_states, (st_data_t)&allocations_statistics);
+	rb_hash_aset(statistics, ID2SYM(rb_intern("tracked_count")), SIZET2NUM(capture->tracked->num_entries));
 	
-	rb_hash_aset(statistics, ID2SYM(rb_intern("total_tracked_objects")), SIZET2NUM(allocations_statistics.total_tracked_objects));
-	rb_hash_aset(statistics, ID2SYM(rb_intern("tracked_objects_per_class")), allocations_statistics.per_class_counts);
+	// Custom object table size
+	size_t states_size = capture->states ? Memory_Profiler_Object_Table_size(capture->states) : 0;
+	rb_hash_aset(statistics, ID2SYM(rb_intern("object_table_size")), SIZET2NUM(states_size));
 	
 	return statistics;
 }
@@ -643,6 +740,7 @@ void Init_Memory_Profiler_Capture(VALUE Memory_Profiler)
 	rb_define_method(Memory_Profiler_Capture, "tracking?", Memory_Profiler_Capture_tracking_p, 1);
 	rb_define_method(Memory_Profiler_Capture, "retained_count_of", Memory_Profiler_Capture_retained_count_of, 1);
 	rb_define_method(Memory_Profiler_Capture, "each", Memory_Profiler_Capture_each, 0);
+	rb_define_method(Memory_Profiler_Capture, "each_object", Memory_Profiler_Capture_each_object, -1);  // -1 = variable args
 	rb_define_method(Memory_Profiler_Capture, "[]", Memory_Profiler_Capture_aref, 1);
 	rb_define_method(Memory_Profiler_Capture, "clear", Memory_Profiler_Capture_clear, 0);
 	rb_define_method(Memory_Profiler_Capture, "statistics", Memory_Profiler_Capture_statistics, 0);

data/ext/memory/profiler/events.c

@@ -1,11 +1,10 @@
 // Released under the MIT License.
 // Copyright, 2025, by Samuel Williams.
 
-#include "ruby.h"
-#include "ruby/debug.h"
-
 #include "events.h"
 #include "capture.h"
+
+#include <ruby/debug.h>
 #include <stdio.h>
 
 enum {
@@ -105,7 +104,10 @@ static void Memory_Profiler_Events_mark_queue(struct Memory_Profiler_Queue *queu
 		
 		rb_gc_mark_movable(event->capture);
 		rb_gc_mark_movable(event->klass);
-		rb_gc_mark_movable(event->object);
+		
+		if (event->type == MEMORY_PROFILER_EVENT_TYPE_NEWOBJ) {
+			rb_gc_mark_movable(event->object);
+		}
 	}
 }
 
@@ -130,7 +132,10 @@ static void Memory_Profiler_Events_compact_queue(struct Memory_Profiler_Queue *q
 		
 		event->capture = rb_gc_location(event->capture);
 		event->klass = rb_gc_location(event->klass);
-		event->object = rb_gc_location(event->object);
+
+		if (event->type == MEMORY_PROFILER_EVENT_TYPE_NEWOBJ) {
+			event->object = rb_gc_location(event->object);
+		}
 	}
 }
 
@@ -179,7 +184,7 @@ int Memory_Profiler_Events_enqueue(
 	VALUE object
 ) {
 	struct Memory_Profiler_Events *events = Memory_Profiler_Events_instance();
-
+	
 	// Always enqueue to the available queue - it won't be touched during processing:
 	struct Memory_Profiler_Event *event = Memory_Profiler_Queue_push(events->available);
 	if (event) {

data/ext/memory/profiler/events.h

@@ -20,10 +20,10 @@ struct Memory_Profiler_Event {
 	// Which Capture instance this event belongs to:
 	VALUE capture;
 	
-	// The class of the object:
+	// The class of the allocated object (Qnil for FREEOBJ):
 	VALUE klass;
 
-	// The object itself (for NEWOBJ and FREEOBJ):
+	// The object pointer being alllocated or freed.
 	VALUE object;
 };
 
@@ -32,7 +32,11 @@ struct Memory_Profiler_Events;
 struct Memory_Profiler_Events* Memory_Profiler_Events_instance(void);
 
 // Enqueue an event to the global queue.
+// object parameter semantics:
+//   - NEWOBJ: the actual object being allocated (queue retains it)
+//   - FREEOBJ: Array with state data for postponed processing
 // Returns non-zero on success, zero on failure.
+// Ruby 3.5 compatible: no FL_SEEN_OBJ_ID or object_id needed
 int Memory_Profiler_Events_enqueue(
 	enum Memory_Profiler_Event_Type type,
 	VALUE capture,

data/ext/memory/profiler/profiler.c

@@ -3,15 +3,26 @@
 
 #include "capture.h"
 
+// Return the memory address of an object as a hex string
+// This matches the format used by ObjectSpace.dump_all
+static VALUE Memory_Profiler_address_of(VALUE module, VALUE object) {
+	char buffer[32];
+	snprintf(buffer, sizeof(buffer), "0x%lx", (unsigned long)object);
+	return rb_str_new_cstr(buffer);
+}
+
 void Init_Memory_Profiler(void)
 {
 #ifdef HAVE_RB_EXT_RACTOR_SAFE
-    rb_ext_ractor_safe(true);
+	rb_ext_ractor_safe(true);
 #endif
-    
-    VALUE Memory = rb_const_get(rb_cObject, rb_intern("Memory"));
-    VALUE Memory_Profiler = rb_define_module_under(Memory, "Profiler");
-    
-    Init_Memory_Profiler_Capture(Memory_Profiler);
+	
+	VALUE Memory = rb_const_get(rb_cObject, rb_intern("Memory"));
+	VALUE Memory_Profiler = rb_define_module_under(Memory, "Profiler");
+	
+	// Add Memory::Profiler.address_of(object) module function:
+	rb_define_module_function(Memory_Profiler, "address_of", Memory_Profiler_address_of, 1);
+	
+	Init_Memory_Profiler_Capture(Memory_Profiler);
 }