From c20066b24cf1c1c235d1f9402b7986b5085cad53 Mon Sep 17 00:00:00 2001
From: Maxime Chevalier-Boisvert <maxime.chevalierboisvert@shopify.com>
Date: Sun, 20 Sep 2020 14:23:14 -0400
Subject: [PATCH] Added method to align code block write position

---
 ujit_asm.c     | 26 +++++++++++++++++++++++---
 ujit_asm.h     |  3 +++
 ujit_compile.c |  3 +++
 3 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/ujit_asm.c b/ujit_asm.c
index 8a9cc7f287..6e12abe513 100644
--- a/ujit_asm.c
+++ b/ujit_asm.c
@@ -138,9 +138,23 @@ void cb_init(codeblock_t* cb, size_t mem_size)
     cb->num_refs = 0;
 }
 
-/**
-Set the current write position
-*/
+// Align the current write position to a multiple of bytes
+void cb_align_pos(codeblock_t* cb, size_t multiple)
+{
+    // Compute the pointer modulo the given alignment boundary
+    uint8_t* ptr = &cb->mem_block[cb->write_pos];
+    size_t rem = ((size_t)ptr) % multiple;
+
+    // If the pointer is already aligned, stop
+    if (rem != 0)
+        return;
+
+    // Pad the pointer by the necessary amount to align it
+    size_t pad = multiple - rem;
+    cb->write_pos += pad;
+}
+
+// Set the current write position
 void cb_set_pos(codeblock_t* cb, size_t pos)
 {
     assert (pos < cb->mem_size);
@@ -1469,6 +1483,12 @@ void sub(codeblock_t* cb, x86opnd_t opnd0, x86opnd_t opnd1)
     );
 }
 
+/// Undefined opcode
+void ud2(codeblock_t* cb)
+{
+    cb_write_bytes(cb, 2, 0x0F, 0x0B);
+}
+
 /// xor - Exclusive bitwise OR
 void xor(codeblock_t* cb, x86opnd_t opnd0, x86opnd_t opnd1)
 {
diff --git a/ujit_asm.h b/ujit_asm.h
index 2ed00b1723..92f90d90c3 100644
--- a/ujit_asm.h
+++ b/ujit_asm.h
@@ -182,7 +182,9 @@ x86opnd_t imm_opnd(int64_t val);
 // Constant pointer operand
 x86opnd_t const_ptr_opnd(void* ptr);
 
+// Code block methods
 void cb_init(codeblock_t* cb, size_t mem_size);
+void cb_align_pos(codeblock_t* cb, size_t multiple);
 void cb_set_pos(codeblock_t* cb, size_t pos);
 uint8_t* cb_get_ptr(codeblock_t* cb, size_t index);
 void cb_write_byte(codeblock_t* cb, uint8_t byte);
@@ -283,6 +285,7 @@ void sar(codeblock_t* cb, x86opnd_t opnd0, x86opnd_t opnd1);
 void shl(codeblock_t* cb, x86opnd_t opnd0, x86opnd_t opnd1);
 void shr(codeblock_t* cb, x86opnd_t opnd0, x86opnd_t opnd1);
 void sub(codeblock_t* cb, x86opnd_t opnd0, x86opnd_t opnd1);
+void ud2(codeblock_t* cb);
 void xor(codeblock_t* cb, x86opnd_t opnd0, x86opnd_t opnd1);
 
 #endif
diff --git a/ujit_compile.c b/ujit_compile.c
index 254a986d85..367fbf8ba9 100644
--- a/ujit_compile.c
+++ b/ujit_compile.c
@@ -141,6 +141,9 @@ ujit_compile_insn(rb_iseq_t *iseq, unsigned int insn_idx, unsigned int* next_uji
         rb_bug("out of executable memory");
     }
 
+    // Align the current write positon to cache line boundaries
+    cb_align_pos(cb, 64);
+
     // Get a pointer to the current write position in the code block
     uint8_t *code_ptr = &cb->mem_block[cb->write_pos];
     //printf("write pos: %ld\n", cb->write_pos);