From ca47899ccf9547223f4c64fc4b0837796bee09af Mon Sep 17 00:00:00 2001
From: Maxime Chevalier-Boisvert <maxime.chevalierboisvert@shopify.com>
Date: Mon, 14 Sep 2020 13:31:34 -0400
Subject: [PATCH] Small refactorings in ujit_compile.c

---
 ujit_asm.c     | 16 ----------------
 ujit_asm.h     |  4 ----
 ujit_compile.c | 44 ++++++++++++++++++++++++++++++++++++--------
 3 files changed, 36 insertions(+), 28 deletions(-)

diff --git a/ujit_asm.c b/ujit_asm.c
index 6b6c586045..dcf1be07f1 100644
--- a/ujit_asm.c
+++ b/ujit_asm.c
@@ -8,9 +8,6 @@
 
 #include "ujit_asm.h"
 
-// TODO: give ujit_examples.h some more meaningful file name
-#include "ujit_examples.h"
-
 // Dummy none/null operand
 const x86opnd_t NO_OPND = { OPND_NONE, 0, .imm = 0 };
 
@@ -213,19 +210,6 @@ void cb_write_int(codeblock_t* cb, uint64_t val, size_t num_bits)
     }
 }
 
-// Ruby instruction prologue and epilogue functions
-void cb_write_prologue(codeblock_t* cb)
-{
-    for (size_t i = 0; i < sizeof(ujit_pre_call_bytes); ++i)
-        cb_write_byte(cb, ujit_pre_call_bytes[i]);
-}
-
-void cb_write_epilogue(codeblock_t* cb)
-{
-    for (size_t i = 0; i < sizeof(ujit_post_call_bytes); ++i)
-        cb_write_byte(cb, ujit_post_call_bytes[i]);
-}
-
 // Allocate a new label with a given name
 size_t cb_new_label(codeblock_t* cb, const char* name)
 {
diff --git a/ujit_asm.h b/ujit_asm.h
index d0de30628e..a977dd9478 100644
--- a/ujit_asm.h
+++ b/ujit_asm.h
@@ -188,10 +188,6 @@ void cb_write_label(codeblock_t* cb, size_t label_idx);
 void cb_label_ref(codeblock_t* cb, size_t label_idx);
 void cb_link_labels(codeblock_t* cb);
 
-// Ruby instruction prologue and epilogue functions
-void cb_write_prologue(codeblock_t* cb);
-void cb_write_epilogue(codeblock_t* cb);
-
 // Encode individual instructions into a code block
 void add(codeblock_t* cb, x86opnd_t opnd0, x86opnd_t opnd1);
 void call_label(codeblock_t* cb, size_t label_idx);
diff --git a/ujit_compile.c b/ujit_compile.c
index 82ea7291f8..3f3f8fd2c6 100644
--- a/ujit_compile.c
+++ b/ujit_compile.c
@@ -8,11 +8,34 @@
 #include "ujit_compile.h"
 #include "ujit_asm.h"
 
+// TODO: give ujit_examples.h some more meaningful file name
+#include "ujit_examples.h"
+
 static codeblock_t block;
 static codeblock_t* cb = NULL;
 
+// Hash table of encoded instructions
 extern st_table *rb_encoded_insn_data;
 
+static void ujit_init();
+
+// Ruby instruction entry
+static void
+ujit_instr_entry(codeblock_t* cb)
+{
+    for (size_t i = 0; i < sizeof(ujit_pre_call_bytes); ++i)
+        cb_write_byte(cb, ujit_pre_call_bytes[i]);
+}
+
+// Ruby instruction exit
+static void
+ujit_instr_exit(codeblock_t* cb)
+{
+    for (size_t i = 0; i < sizeof(ujit_post_call_bytes); ++i)
+        cb_write_byte(cb, ujit_post_call_bytes[i]);
+}
+
+// Keep track of mapping from instructions to generated code
 // See comment for rb_encoded_insn_data in iseq.c
 static void
 addr2insn_bookkeeping(void *code_ptr, int insn)
@@ -33,12 +56,10 @@ addr2insn_bookkeeping(void *code_ptr, int insn)
 uint8_t *
 ujit_compile_insn(rb_iseq_t *iseq, size_t insn_idx)
 {
-    // Allocate the code block if not previously allocated
+    // If not previously done, initialize ujit
     if (!cb)
     {
-        // 4MB ought to be enough for anybody
-        cb = &block;
-        cb_init(cb, 4000000);
+        ujit_init();
     }
 
     int insn = (int)iseq->body->iseq_encoded[insn_idx];
@@ -57,14 +78,14 @@ ujit_compile_insn(rb_iseq_t *iseq, size_t insn_idx)
     if (insn == BIN(nop))
     {
         // Write the pre call bytes
-        cb_write_prologue(cb);
+        ujit_instr_entry(cb);
 
         add(cb, RSI, imm_opnd(8));                  // increment PC
         mov(cb, mem_opnd(64, RDI, 0), RSI);         // write new PC to EC object, not necessary for nop bytecode?
         mov(cb, RAX, RSI);                          // return new PC
 
         // Write the post call bytes
-        cb_write_epilogue(cb);
+        ujit_instr_exit(cb);
 
         addr2insn_bookkeeping(code_ptr, insn);
 
@@ -74,7 +95,7 @@ ujit_compile_insn(rb_iseq_t *iseq, size_t insn_idx)
     if (insn == BIN(pop))
     {
         // Write the pre call bytes
-        cb_write_prologue(cb);
+        ujit_instr_entry(cb);
 
         sub(cb, mem_opnd(64, RDI, 8), imm_opnd(8)); // decrement SP
         add(cb, RSI, imm_opnd(8));                  // increment PC
@@ -82,7 +103,7 @@ ujit_compile_insn(rb_iseq_t *iseq, size_t insn_idx)
         mov(cb, RAX, RSI);                          // return new PC
 
         // Write the post call bytes
-        cb_write_epilogue(cb);
+        ujit_instr_exit(cb);
 
         addr2insn_bookkeeping(code_ptr, insn);
 
@@ -91,3 +112,10 @@ ujit_compile_insn(rb_iseq_t *iseq, size_t insn_idx)
 
     return 0;
 }
+
+static void ujit_init()
+{
+    // 4MB ought to be enough for anybody
+    cb = &block;
+    cb_init(cb, 4000000);
+}