First pass at code page GC object.

4 files changed, 69 insertions, 27 deletions

diff --git a/yjit_asm.c b/yjit_asm.c
index 1b44587e3c..31e1528ff0 100644
--- a/yjit_asm.c
+++ b/yjit_asm.c
@@ -224,18 +224,11 @@ uint8_t* alloc_exec_mem(uint32_t mem_size)
 // How many code pages to allocate at once
 #define PAGES_PER_ALLOC 512
 
-typedef struct free_list_node
-{
-    uint8_t* page_ptr;
-
-    struct free_list_node *next;
-
-} freelist_t;
-
-freelist_t *freelist = NULL;
+// Head of the list of free code pages
+code_page_t *freelist = NULL;
 
 // Allocate a single code page from a pool of free pages
-uint8_t* alloc_code_page()
+code_page_t* alloc_code_page()
 {
     fprintf(stderr, "allocating code page\n");
 
@@ -246,29 +239,25 @@ uint8_t* alloc_code_page()
 
         // Do this in reverse order so we allocate our pages in order
         for (int i = PAGES_PER_ALLOC - 1; i >= 0; --i) {
-            freelist_t* node = malloc(sizeof(freelist_t));
-            node->page_ptr = code_chunk + i * CODE_PAGE_SIZE;
-            node->next = freelist;
-            freelist = node;
+            code_page_t* code_page = malloc(sizeof(code_page_t));
+            code_page->mem_block = code_chunk + i * CODE_PAGE_SIZE;
+            code_page->page_size = CODE_PAGE_SIZE;
+            code_page->_next = freelist;
+            freelist = code_page;
         }
     }
 
-    freelist_t* free_node = freelist;
-    uint8_t* page_ptr = freelist->page_ptr;
-
-    freelist = freelist->next;
-    free(free_node);
+    code_page_t* free_page = freelist;
+    freelist = freelist->_next;
 
-    return page_ptr;
+    return free_page;
 }
 
 // Put a code page back into the allocation pool
-void free_code_page(uint8_t* page_ptr)
+void free_code_page(code_page_t* code_page)
 {
-    freelist_t* node = malloc(sizeof(freelist_t));
-    node->page_ptr = page_ptr;
-    node->next = freelist;
-    freelist = node;
+    code_page->_next = freelist;
+    freelist = code_page;
 }
 
 // Initialize a code block object
diff --git a/yjit_asm.h b/yjit_asm.h
index c94b54240c..93762987bc 100644
--- a/yjit_asm.h
+++ b/yjit_asm.h
@@ -131,6 +131,20 @@ typedef struct X86Opnd
 
 } x86opnd_t;
 
+// Struct representing a code page
+typedef struct code_page_struct
+{
+    // Chunk of executable memory
+    uint8_t* mem_block;
+
+    // Size of the executable memory chunk
+    uint32_t page_size;
+
+    // Next node in the free list (private)
+    struct code_page_struct* _next;
+
+} code_page_t;
+
 // Dummy none/null operand
 static const x86opnd_t NO_OPND = { OPND_NONE, 0, .as.imm = 0 };
 
@@ -242,8 +256,8 @@ x86opnd_t const_ptr_opnd(const void *ptr);
 
 // Machine code allocation
 uint8_t* alloc_exec_mem(uint32_t mem_size);
-uint8_t* alloc_code_page();
-void free_code_page(uint8_t* page_ptr);
+code_page_t* alloc_code_page();
+void free_code_page(code_page_t* code_page);
 
 // Code block methods
 void cb_init(codeblock_t* cb, uint8_t* mem_block, uint32_t mem_size);
diff --git a/yjit_core.h b/yjit_core.h
index 45b7c18557..cff46648cc 100644
--- a/yjit_core.h
+++ b/yjit_core.h
@@ -242,8 +242,12 @@ typedef struct yjit_block_version
     VALUE receiver_klass;
     VALUE callee_cme;
 
+    // Code page this block lives on
+    VALUE code_page;
+
     // Index one past the last instruction in the iseq
     uint32_t end_idx;
+
 } block_t;
 
 // Context object methods
diff --git a/yjit_iface.c b/yjit_iface.c
index 1cb455e0c3..a7d18e0687 100644
--- a/yjit_iface.c
+++ b/yjit_iface.c
@@ -855,6 +855,9 @@ rb_yjit_iseq_mark(const struct rb_iseq_constant_body *body)
                 memcpy(&object, value_address, SIZEOF_VALUE);
                 rb_gc_mark_movable(object);
             }
+
+            // Mark the machine code page this block lives on
+            rb_gc_mark_movable(block->code_page);
         }
     }
 }
@@ -895,6 +898,9 @@ rb_yjit_iseq_update_references(const struct rb_iseq_constant_body *body)
                     memcpy(value_address, &possibly_moved, SIZEOF_VALUE);
                 }
             }
+
+            // Update the machine code page this block lives on
+            block->code_page = rb_gc_location(block->code_page);
         }
     }
 }
@@ -917,6 +923,35 @@ rb_yjit_iseq_free(const struct rb_iseq_constant_body *body)
     rb_darray_free(body->yjit_blocks);
 }
 
+static void
+yjit_code_page_free(void *code_page)
+{
+    free_code_page((code_page_t*)code_page);
+}
+
+// Custom type for interacting with the GC
+static const rb_data_type_t yjit_code_page_type = {
+    "yjit_code_page",
+    {NULL, yjit_code_page_free, NULL, NULL},
+    0, 0, RUBY_TYPED_FREE_IMMEDIATELY
+};
+
+// Allocate a code page and wrap it into a Ruby object owned by the GC
+VALUE rb_yjit_code_page_alloc()
+{
+    code_page_t* code_page = alloc_code_page();
+    VALUE cp_obj = TypedData_Wrap_Struct(0, &yjit_code_page_type, code_page);
+    return cp_obj;
+}
+
+// Unwrap the Ruby object representing a code page
+code_page_t *rb_yjit_code_page_unwrap(VALUE cp_obj)
+{
+    code_page_t * code_page;
+    TypedData_Get_Struct(cp_obj, code_page_t, &yjit_code_page_type, code_page);
+    return code_page;
+}
+
 bool
 rb_yjit_enabled_p(void)
 {