diff --git a/src/codegen.h b/src/codegen.h
index 961bef5..d45ade7 100644
--- a/src/codegen.h
+++ b/src/codegen.h
@@ -65,12 +65,14 @@ static inline codeblock_t *codeblock_tree_find(uint32_t phys)
         while (block)
         {
                 if (phys == block->phys)
-                        return block;
+                        break;
                 else if (phys < block->phys)
                         block = block->left;
                 else
                         block = block->right;
         }
+        
+        return block;
 }
 
 static inline void codeblock_tree_add(codeblock_t *new_block)
diff --git a/src/codegen_x86-64.c b/src/codegen_x86-64.c
index 7181fda..9fc6783 100644
--- a/src/codegen_x86-64.c
+++ b/src/codegen_x86-64.c
@@ -126,6 +126,8 @@ static void delete_block(codeblock_t *block)
                 fatal("Deleting deleted block\n");
         block->pc = 0;
 
+        codeblock_tree_delete(block);
+        
         if (block->prev)
         {
                 block->prev->next = block->next;
@@ -464,6 +466,8 @@ void codegen_block_end()
 //        pclog("End block %i\n", block_num);
 
         recomp_page = -1;
+
+        codeblock_tree_add(block);
 }
 
 void codegen_flush()