Try to break the code page refactoring into smaller steps

yjit_asm.c

@@ -219,20 +219,12 @@ uint8_t* alloc_exec_mem(uint32_t mem_size)
 #endif
 }
 
-// Size of code pages to allocate
-#define CODE_PAGE_SIZE 16 * 1024
-
-// How many code pages to allocate at once
-#define PAGES_PER_ALLOC 512
-
 // Head of the list of free code pages
 code_page_t *freelist = NULL;
 
 // Allocate a single code page from a pool of free pages
 code_page_t* alloc_code_page()
 {
-    fprintf(stderr, "allocating code page\n");
-
     // If the free list is empty
     if (!freelist) {
         // Allocate many pages at once
@@ -242,6 +234,7 @@ code_page_t* alloc_code_page()
         for (int i = PAGES_PER_ALLOC - 1; i >= 0; --i) {
             code_page_t* code_page = malloc(sizeof(code_page_t));
             code_page->mem_block = code_chunk + i * CODE_PAGE_SIZE;
+            assert ((intptr_t)code_page->mem_block % CODE_PAGE_SIZE == 0);
             code_page->page_size = CODE_PAGE_SIZE;
             code_page->_next = freelist;
             freelist = code_page;
@@ -264,6 +257,7 @@ void free_code_page(code_page_t* code_page)
 // Initialize a code block object
 void cb_init(codeblock_t* cb, uint8_t* mem_block, uint32_t mem_size)
 {
+    assert (mem_block);
     cb->mem_block = mem_block;
     cb->mem_size = mem_size;
     cb->write_pos = 0;
@@ -290,6 +284,14 @@ void cb_set_pos(codeblock_t* cb, uint32_t pos)
     cb->write_pos = pos;
 }
 
+// Set the current write position from a pointer
+void cb_set_write_ptr(codeblock_t* cb, uint8_t* code_ptr)
+{
+    intptr_t pos = code_ptr - cb->mem_block;
+    assert (pos < cb->mem_size);
+    cb->write_pos = (uint32_t)pos;
+}
+
 // Get a direct pointer into the executable memory block
 uint8_t* cb_get_ptr(codeblock_t* cb, uint32_t index)
 {
@@ -297,6 +299,12 @@ uint8_t* cb_get_ptr(codeblock_t* cb, uint32_t index)
     return &cb->mem_block[index];
 }
 
+// Get a direct pointer to the current write position
+uint8_t* cb_get_write_ptr(codeblock_t* cb)
+{
+    return cb_get_ptr(cb, cb->write_pos);
+}
+
 // Write a byte at the current position
 void cb_write_byte(codeblock_t* cb, uint8_t byte)
 {

yjit_asm.h

@@ -5,6 +5,12 @@
 #include <stddef.h>
 #include <stdbool.h>
 
+// Size of code pages to allocate
+#define CODE_PAGE_SIZE 16 * 1024
+
+// How many code pages to allocate at once
+#define PAGES_PER_ALLOC 512
+
 // Maximum number of labels to link
 #define MAX_LABELS 32
 
@@ -263,7 +269,9 @@ void free_code_page(code_page_t* code_page);
 void cb_init(codeblock_t* cb, uint8_t* mem_block, uint32_t mem_size);
 void cb_align_pos(codeblock_t* cb, uint32_t multiple);
 void cb_set_pos(codeblock_t* cb, uint32_t pos);
+void cb_set_write_ptr(codeblock_t* cb, uint8_t* code_ptr);
 uint8_t* cb_get_ptr(codeblock_t* cb, uint32_t index);
+uint8_t* cb_get_write_ptr(codeblock_t* cb);
 void cb_write_byte(codeblock_t* cb, uint8_t byte);
 void cb_write_bytes(codeblock_t* cb, uint32_t num_bytes, ...);
 void cb_write_int(codeblock_t* cb, uint64_t val, uint32_t num_bits);

yjit_iface.c

@@ -971,6 +971,10 @@ VALUE rb_yjit_code_page_alloc(void)
 {
     code_page_t* code_page = alloc_code_page();
     VALUE cp_obj = TypedData_Wrap_Struct(0, &yjit_code_page_type, code_page);
+
+    // Write a pointer to the wrapper object at the beginning of the code page
+    *((VALUE*)code_page->mem_block) = cp_obj;
+
     return cp_obj;
 }
 
@@ -982,6 +986,44 @@ code_page_t *rb_yjit_code_page_unwrap(VALUE cp_obj)
     return code_page;
 }
 
+// Get the code page wrapper object for a code pointer
+VALUE rb_yjit_code_page_from_ptr(uint8_t* code_ptr)
+{
+    VALUE* page_start = (VALUE*)((intptr_t)code_ptr & ~(CODE_PAGE_SIZE - 1));
+    VALUE wrapper = *page_start;
+    return wrapper;
+}
+
+// Get the inline code block corresponding to a code pointer
+void rb_yjit_get_cb(codeblock_t* cb, uint8_t* code_ptr)
+{
+    VALUE page_wrapper = rb_yjit_code_page_from_ptr(code_ptr);
+    code_page_t *code_page = rb_yjit_code_page_unwrap(page_wrapper);
+
+    // A pointer to the page wrapper object is written at the start of the code page
+    uint8_t* mem_block = code_page->mem_block + sizeof(VALUE);
+    uint32_t mem_size = (code_page->page_size/2) - sizeof(VALUE);
+    RUBY_ASSERT(mem_block);
+
+    // Map the code block to this memory region
+    cb_init(cb, mem_block, mem_size);
+}
+
+// Get the outlined code block corresponding to a code pointer
+void rb_yjit_get_ocb(codeblock_t* cb, uint8_t* code_ptr)
+{
+    VALUE page_wrapper = rb_yjit_code_page_from_ptr(code_ptr);
+    code_page_t *code_page = rb_yjit_code_page_unwrap(page_wrapper);
+
+    // A pointer to the page wrapper object is written at the start of the code page
+    uint8_t* mem_block = code_page->mem_block + (code_page->page_size/2);
+    uint32_t mem_size = code_page->page_size/2;
+    RUBY_ASSERT(mem_block);
+
+    // Map the code block to this memory region
+    cb_init(cb, mem_block, mem_size);
+}
+
 bool
 rb_yjit_enabled_p(void)
 {

yjit_iface.h

@@ -126,4 +126,9 @@ const VALUE *rb_yjit_count_side_exit_op(const VALUE *exit_pc);
 void yjit_unlink_method_lookup_dependency(block_t *block);
 void yjit_block_assumptions_free(block_t *block);
 
+VALUE rb_yjit_code_page_alloc(void);
+code_page_t *rb_yjit_code_page_unwrap(VALUE cp_obj);
+void rb_yjit_get_cb(codeblock_t* cb, uint8_t* code_ptr);
+void rb_yjit_get_ocb(codeblock_t* cb, uint8_t* code_ptr);
+
 #endif // #ifndef YJIT_IFACE_H