From ef08af93769982047e59356d76c7e72a67970335 Mon Sep 17 00:00:00 2001
From: Maxime Chevalier-Boisvert <maxime.chevalierboisvert@shopify.com>
Date: Tue, 12 Jan 2021 17:03:54 -0500
Subject: [PATCH] Progress on code invalidation

---
 ujit_codegen.c | 14 +++++++++-----
 ujit_codegen.h |  5 ++++-
 ujit_core.c    | 29 ++++++++++++++++++++++++++---
 ujit_core.h    |  2 ++
 4 files changed, 41 insertions(+), 9 deletions(-)

diff --git a/ujit_codegen.c b/ujit_codegen.c
index 8ba71bf37f..4833e4ab1f 100644
--- a/ujit_codegen.c
+++ b/ujit_codegen.c
@@ -133,8 +133,7 @@ uint8_t* ujit_compile_entry(const rb_iseq_t *iseq, uint32_t insn_idx)
     ctx_t ctx = { 0 };
 
     // Compile the block starting at this instruction
-    uint32_t num_instrs = 0;
-    ujit_compile_block(iseq, insn_idx, &ctx, &num_instrs);
+    uint32_t num_instrs = ujit_compile_block(iseq, insn_idx, &ctx);
 
     // If no instructions were compiled
     if (num_instrs == 0) {
@@ -154,8 +153,8 @@ uint8_t* ujit_compile_entry(const rb_iseq_t *iseq, uint32_t insn_idx)
 /*
 Compile a sequence of bytecode instructions starting at `insn_idx`.
 */
-void
-ujit_compile_block(const rb_iseq_t *iseq, uint32_t insn_idx, ctx_t* ctx, uint32_t* num_instrs)
+uint32_t
+ujit_compile_block(/*version_t* version,*/ const rb_iseq_t *iseq, uint32_t insn_idx, ctx_t* ctx)
 {
     assert (cb != NULL);
     VALUE *encoded = iseq->body->iseq_encoded;
@@ -175,10 +174,13 @@ ujit_compile_block(const rb_iseq_t *iseq, uint32_t insn_idx, ctx_t* ctx, uint32_
 
     // Initialize JIT state object
     jitstate_t jit = {
+        NULL,
         iseq,
         insn_idx
     };
 
+    uint32_t num_instrs = 0;
+
     // For each instruction to compile
     for (;;) {
         // Set the current instruction
@@ -209,7 +211,7 @@ ujit_compile_block(const rb_iseq_t *iseq, uint32_t insn_idx, ctx_t* ctx, uint32_
     	// Move to the next instruction
         p_last_op = p_desc;
         insn_idx += insn_len(opcode);
-        (*num_instrs)++;
+        num_instrs++;
 
         // If this instruction terminates this block
         if (p_desc->is_branch) {
@@ -234,6 +236,8 @@ ujit_compile_block(const rb_iseq_t *iseq, uint32_t insn_idx, ctx_t* ctx, uint32_
             pc += insn_len(opcode);
         }
     }
+
+    return num_instrs;
 }
 
 static bool
diff --git a/ujit_codegen.h b/ujit_codegen.h
index 413ecb35ac..d01a73bde0 100644
--- a/ujit_codegen.h
+++ b/ujit_codegen.h
@@ -11,6 +11,9 @@ codeblock_t* ocb;
 // Code generation state
 typedef struct JITState
 {
+    // Block version being compiled
+    version_t* version;
+
     // Instruction sequence this is associated with
     const rb_iseq_t *iseq;
 
@@ -42,7 +45,7 @@ typedef struct OpDesc
 
 uint8_t* ujit_compile_entry(const rb_iseq_t *iseq, uint32_t insn_idx);
 
-void ujit_compile_block(const rb_iseq_t *iseq, uint32_t insn_idx, ctx_t* ctx, uint32_t* num_instrs);
+uint32_t ujit_compile_block(const rb_iseq_t *iseq, uint32_t insn_idx, ctx_t* ctx);
 
 void ujit_init_codegen(void);
 
diff --git a/ujit_core.c b/ujit_core.c
index c3166d0f09..c316675185 100644
--- a/ujit_core.c
+++ b/ujit_core.c
@@ -95,9 +95,8 @@ version_t* gen_block_version(blockid_t blockid, const ctx_t* ctx)
 
     // Compile the block version
     ctx_t ctx_copy = *ctx;
-    uint32_t num_instrs = 0;
     p_version->start_pos = cb->write_pos;
-    ujit_compile_block(blockid.iseq, blockid.idx, &ctx_copy, &num_instrs);
+    ujit_compile_block(blockid.iseq, blockid.idx, &ctx_copy);
     p_version->end_pos = cb->write_pos;
 
     // Keep track of the new block version
@@ -159,7 +158,7 @@ uint8_t* branch_stub_hit(uint32_t branch_idx, uint32_t target_idx)
 }
 
 // Get a version or stub corresponding to a branch target
-// TODO: need incoming and target versioning contexts
+// TODO: need incoming and target contexts
 uint8_t* get_branch_target(
     blockid_t target,
     const ctx_t* ctx,
@@ -237,6 +236,30 @@ void gen_branch(
     assert (num_branches < MAX_BRANCHES);
     branch_entries[num_branches] = branch_entry;
     num_branches++;
+}
+
+// Invalidate one specific block version
+void invalidate(version_t* version)
+{
+    // All branches jumping to the block should be atomically patched with jumps going to a stub instead.
+
+    // There can also be other blocks falling through to the invalidated block because they immediately precede it.
+    // - If an incoming fall-through branch is too short to be patched, we may need to invalidate its block
+    // - This may not be an issue in practice, because the block we go to could have space
+    // - We can force any block that may need to be invalidated to have sufficient space to contain a jump to a stub
+
+    // If the block is an entry point, it needs to be unmapped from its iseq
+    // Unmap/remap anything at this iseq/idx
+
+    // Optional: may want to recompile a new deoptimized entry point
+    // Call continuation addresses on the stack can also be atomically replaced by jumps going to the stub.
+
+
+
+
+
+
+
 }
 
 int blockid_cmp(st_data_t arg0, st_data_t arg1)
diff --git a/ujit_core.h b/ujit_core.h
index 4c80bbc5d4..07a18ed727 100644
--- a/ujit_core.h
+++ b/ujit_core.h
@@ -122,6 +122,8 @@ void gen_branch(
     branchgen_fn gen_fn
 );
 
+void invalidate(version_t* version);
+
 void ujit_init_core(void);
 
 #endif // #ifndef UJIT_CORE_H