Refactor memory scanning code

natalie-lang · Mar 3, 2025 · d76894a · d76894a
1 parent d20780e
commit d76894a
Show file tree

Hide file tree

Showing 4 changed files with 44 additions and 55 deletions.
diff --git a/include/natalie/gc/heap.hpp b/include/natalie/gc/heap.hpp
@@ -1,6 +1,7 @@
 #pragma once
 
 #include <array>
+#include <functional>
 #include <mutex>
 #include <setjmp.h>
 #include <stdlib.h>
@@ -81,6 +82,9 @@ class Heap {
     bool collect_all_at_exit() const { return m_collect_all_at_exit; }
     void set_collect_all_at_exit(bool collect) { m_collect_all_at_exit = collect; }
 
+    void scan_memory(Cell::Visitor &visitor, void *start, void *end);
+    void scan_memory(Cell::Visitor &visitor, void *start, void *end, std::function<void(Cell *)> fn);
+
 private:
     friend Allocator;
 

diff --git a/src/fiber_object.cpp b/src/fiber_object.cpp
@@ -1,5 +1,4 @@
 #include "natalie.hpp"
-#include "natalie/thread_object.hpp"
 
 #define MINICORO_IMPL
 #include "minicoro/minicoro.h"
@@ -321,14 +320,13 @@ NO_SANITIZE_ADDRESS void FiberObject::visit_children_from_stack(Visitor &visitor
     if (m_status == Status::Terminated)
         return;
 
-    for (char *ptr = reinterpret_cast<char *>(m_end_of_stack); ptr < m_start_of_stack; ptr += sizeof(intptr_t)) {
-        Cell *potential_cell = *reinterpret_cast<Cell **>(ptr);
-        if (Heap::the().is_a_heap_cell_in_use(potential_cell))
-            visitor.visit(potential_cell);
 #ifdef __SANITIZE_ADDRESS__
+    Heap::the().scan_memory(visitor, m_end_of_stack, m_start_of_stack, [&visitor](Cell *potential_cell) {
         visit_children_from_asan_fake_stack(visitor, potential_cell);
+    });
+#else
+    Heap::the().scan_memory(visitor, m_end_of_stack, m_start_of_stack);
 #endif
-    }
 }
 
 #ifdef __SANITIZE_ADDRESS__
@@ -339,11 +337,7 @@ NO_SANITIZE_ADDRESS void FiberObject::visit_children_from_asan_fake_stack(Visito
 
     if (!real_stack) return;
 
-    for (char *ptr = reinterpret_cast<char *>(begin); ptr < end; ptr += sizeof(intptr_t)) {
-        Cell *potential_cell = *reinterpret_cast<Cell **>(ptr);
-        if (Heap::the().is_a_heap_cell_in_use(potential_cell))
-            visitor.visit(potential_cell);
-    }
+    Heap::the().scan_memory(visitor, begin, end);
 }
 #else
 void FiberObject::visit_children_from_asan_fake_stack(Visitor &visitor, Cell *potential_cell) const { }
@@ -361,7 +355,6 @@ HashObject *FiberObject::ensure_thread_storage() {
         m_thread_storage = new HashObject {};
     return m_thread_storage;
 }
-
 }
 
 extern "C" {

diff --git a/src/gc.cpp b/src/gc.cpp
@@ -37,13 +37,7 @@ NO_SANITIZE_ADDRESS void Heap::visit_roots_from_asan_fake_stack(Cell::Visitor &v
 
     if (!real_stack) return;
 
-    for (char *ptr = reinterpret_cast<char *>(begin_fake_frame); ptr < end_fake_frame; ptr += sizeof(intptr_t)) {
-        Cell *potential_cell = *reinterpret_cast<Cell **>(ptr);
-        if (!potential_cell)
-            continue;
-        if (is_a_heap_cell_in_use(potential_cell))
-            visitor.visit(potential_cell);
-    }
+    scan_memory(visitor, begin_fake_frame, end_fake_frame);
 }
 #else
 void Heap::visit_roots_from_asan_fake_stack(Cell::Visitor &visitor, Cell *potential_cell) { }
@@ -56,28 +50,19 @@ NO_SANITIZE_ADDRESS void Heap::visit_roots(Cell::Visitor &visitor) {
     // step over stack, saving potential pointers
     auto start_of_stack = ThreadObject::current()->start_of_stack();
     assert(start_of_stack > end_of_stack);
-
-    for (char *ptr = reinterpret_cast<char *>(end_of_stack); ptr < start_of_stack; ptr += sizeof(intptr_t)) {
-        Cell *potential_cell = *reinterpret_cast<Cell **>(ptr);
-        if (!potential_cell)
-            continue;
-        if (is_a_heap_cell_in_use(potential_cell))
-            visitor.visit(potential_cell);
 #ifdef __SANITIZE_ADDRESS__
+    scan_memory(visitor, end_of_stack, start_of_stack, [&visitor](Cell *potential_cell) {
         visit_roots_from_asan_fake_stack(visitor, potential_cell);
+    });
+#else
+    scan_memory(visitor, end_of_stack, start_of_stack);
 #endif
-    }
 
     // step over any registers, saving potential pointers
     jmp_buf jump_buf;
     setjmp(jump_buf);
-    for (char *ptr = reinterpret_cast<char *>(jump_buf); ptr < reinterpret_cast<char *>(jump_buf) + sizeof(jump_buf); ptr += sizeof(intptr_t)) {
-        Cell *potential_cell = *reinterpret_cast<Cell **>(ptr);
-        if (!potential_cell)
-            continue;
-        if (is_a_heap_cell_in_use(potential_cell))
-            visitor.visit(potential_cell);
-    }
+    auto start = reinterpret_cast<std::byte *>(jump_buf);
+    scan_memory(visitor, start, start + sizeof(jump_buf));
 }
 
 void Heap::collect() {
@@ -253,6 +238,27 @@ void Heap::dump() const {
     printf("Total allocations: %lld\n", allocation_count);
 }
 
+void Heap::scan_memory(Cell::Visitor &visitor, void *start, void *end) {
+    for (auto *ptr = reinterpret_cast<std::byte *>(start); ptr < end; ptr += sizeof(intptr_t)) {
+        Cell *potential_cell = *reinterpret_cast<Cell **>(ptr);
+        if (!potential_cell)
+            continue;
+        if (is_a_heap_cell_in_use(potential_cell))
+            visitor.visit(potential_cell);
+    }
+}
+
+void Heap::scan_memory(Cell::Visitor &visitor, void *start, void *end, std::function<void(Cell *)> fn) {
+    for (auto *ptr = reinterpret_cast<std::byte *>(start); ptr < end; ptr += sizeof(intptr_t)) {
+        Cell *potential_cell = *reinterpret_cast<Cell **>(ptr);
+        if (!potential_cell)
+            continue;
+        if (is_a_heap_cell_in_use(potential_cell))
+            visitor.visit(potential_cell);
+        fn(potential_cell);
+    }
+}
+
 Cell *HeapBlock::find_next_free_cell() {
     assert(has_free());
     --m_free_count;

diff --git a/src/thread_object.cpp b/src/thread_object.cpp
@@ -842,19 +842,14 @@ NO_SANITIZE_ADDRESS void ThreadObject::visit_children_from_stack(Visitor &visito
     if (m_suspend_status == SuspendStatus::Launching)
         return;
 
-    // Walk the stack looking for variables...
-    for (char *ptr = reinterpret_cast<char *>(m_end_of_stack); ptr < m_start_of_stack; ptr += sizeof(intptr_t)) {
-        Cell *potential_cell = *reinterpret_cast<Cell **>(ptr);
-        if (!potential_cell)
-            continue;
-        if (Heap::the().is_a_heap_cell_in_use(potential_cell))
-            visitor.visit(potential_cell);
 #ifdef __SANITIZE_ADDRESS__
+    Heap::the().scan_memory(visitor, m_end_of_stack, m_start_of_stack, [&visitor](Cell *potential_cell) {
         visit_children_from_asan_fake_stack(visitor, potential_cell);
+    });
+#else
+    Heap::the().scan_memory(visitor, m_end_of_stack, m_start_of_stack);
 #endif
-    }
 
-    // Check in the registers for any variables...
     visit_children_from_context(visitor);
 }
 
@@ -869,24 +864,15 @@ NO_SANITIZE_ADDRESS void ThreadObject::visit_children_from_asan_fake_stack(Visit
 
     if (!real_stack) return;
 
-    for (char *ptr = reinterpret_cast<char *>(begin); ptr < end; ptr += sizeof(intptr_t)) {
-        Cell *potential_cell = *reinterpret_cast<Cell **>(ptr);
-        if (Heap::the().is_a_heap_cell_in_use(potential_cell))
-            visitor.visit(potential_cell);
-    }
+    Heap::the().scan_memory(visitor, begin, end);
 }
 #else
 void ThreadObject::visit_children_from_asan_fake_stack(Visitor &visitor, Cell *potential_cell) const { }
 #endif
 
 NO_SANITIZE_ADDRESS void ThreadObject::visit_children_from_context(Visitor &visitor) const {
-    for (char *i = (char *)m_context; i < (char *)m_context + sizeof(ucontext_t); ++i) {
-        Cell *potential_cell = *reinterpret_cast<Cell **>(i);
-        if (!potential_cell)
-            continue;
-        if (Heap::the().is_a_heap_cell_in_use(potential_cell))
-            visitor.visit(potential_cell);
-    }
+    auto start = reinterpret_cast<std::byte *>(m_context);
+    Heap::the().scan_memory(visitor, start, start + sizeof(ucontext_t));
 }
 
 }