[bug] Fixed not rewriting injected code by Compiler

* The bug was introduced by the latest changes. The rewriter
    only interated basic blocks, but not possibly injected code
    that converts one register state into another via a small
    code sequence injected at the end of the function.

  * Resolved some static analysis problems. The reported
    use-after-free was a false positive constructing by
    contradicting conditions - so new asserts guard this.

src/asmjit/arm/a64rapass.cpp

@@ -685,13 +685,9 @@ Error ARMRAPass::build_cfg_nodes() noexcept {
 // ========================
 
 ASMJIT_FAVOR_SPEED Error ARMRAPass::rewrite() noexcept {
-  size_t virt_count = cc()._virt_regs.size();
-
-  for (RABlock* block : _pov.iterate_reverse()) {
-    BaseNode* node = block->first();
-    BaseNode* stop = block->last();
-
-    for (;;) {
+  const size_t virt_count = cc()._virt_regs.size();
+  return rewrite_iterate([&](BaseNode* node, BaseNode* stop, RABlock* block) noexcept -> Error {
+    while (node != stop) {
       BaseNode* next = node->next();
 
       if (node->is_inst()) {
@@ -744,7 +740,10 @@ ASMJIT_FAVOR_SPEED Error ARMRAPass::rewrite() noexcept {
 
               BaseNode* prev = node->prev();
               cc().remove_node(node);
-              block->set_last(prev);
+
+              if (block) {
+                block->set_last(prev);
+              }
             }
           }
         }
@@ -819,15 +818,11 @@ ASMJIT_FAVOR_SPEED Error ARMRAPass::rewrite() noexcept {
         }
       }
 
-      if (node == stop) {
-        break;
-      }
-
       node = next;
     }
-  }
 
-  return Error::kOk;
+    return Error::kOk;
+  });
 }
 
 // a64::ARMRAPass - Prolog & Epilog

src/asmjit/core/api-config.h

@@ -16,7 +16,7 @@
 #define ASMJIT_LIBRARY_MAKE_VERSION(major, minor, patch) ((major << 16) | (minor << 8) | (patch))
 
 //! AsmJit library version, see \ref ASMJIT_LIBRARY_MAKE_VERSION for a version format reference.
-#define ASMJIT_LIBRARY_VERSION ASMJIT_LIBRARY_MAKE_VERSION(1, 18, 0)
+#define ASMJIT_LIBRARY_VERSION ASMJIT_LIBRARY_MAKE_VERSION(1, 18, 1)
 
 //! \def ASMJIT_ABI_NAMESPACE
 //!

src/asmjit/core/ralocal.cpp

@@ -801,6 +801,7 @@ Error RALocalAllocator::alloc_instruction(InstNode* node) noexcept {
             RegMask available_regs = _available_regs[group] & ~_cur_assignment.assigned(group);
             if (available_regs) {
               uint32_t tmp_reg_id = pick_best_suitable_register(group, available_regs);
+              ASMJIT_ASSERT(tmp_reg_id != RAAssignment::kPhysNone);
 
               ASMJIT_PROPAGATE(on_move_reg(group, this_work_reg, this_work_id, tmp_reg_id, this_phys_id));
               _clobbered_regs[group] |= Support::bit_mask<RegMask>(tmp_reg_id);
@@ -1098,7 +1099,7 @@ Error RALocalAllocator::alloc_branch(InstNode* node, RABlock* target, RABlock* c
   ASMJIT_PROPAGATE(spill_regs_before_entry(target));
 
   if (target->has_entry_assignment()) {
-    BaseNode* injection_point = _pass.extra_block()->prev();
+    BaseNode* injection_point = _pass._injection_end->prev();
     BaseNode* prev_cursor = _cc.set_cursor(injection_point);
 
     _tmp_assignment.copy_from(_cur_assignment);
@@ -1127,6 +1128,10 @@ Error RALocalAllocator::alloc_branch(InstNode* node, RABlock* target, RABlock* c
       ASMJIT_PROPAGATE(_pass.emit_jump(saved_target));
       _cc.set_cursor(injection_point);
       _cc.bind(trampoline);
+
+      if (_pass._injection_start == nullptr) {
+        _pass._injection_start = injection_point->next();
+      }
     }
 
     _cc.set_cursor(prev_cursor);

src/asmjit/core/rapass.cpp

@@ -204,7 +204,8 @@ Error BaseRAPass::run_on_function(Arena& arena, FuncNode* func, [[maybe_unused]]
   _arena = &arena;
   _func = func;
   _stop = end->next();
-  _extra_block = end;
+  _injection_start = nullptr;
+  _injection_end = end;
 
   RAPass_prepare_for_function(this, &_func->_func_detail);
 
@@ -226,7 +227,8 @@ Error BaseRAPass::run_on_function(Arena& arena, FuncNode* func, [[maybe_unused]]
   _arena = nullptr;
   _func = nullptr;
   _stop = nullptr;
-  _extra_block = nullptr;
+  _injection_start = nullptr;
+  _injection_end = nullptr;
 
   // Reset `Arena` as nothing should persist between `run_on_function()` calls.
   arena.reset();
@@ -877,7 +879,7 @@ ASMJIT_FAVOR_SPEED Error BaseRAPass::build_reg_ids() noexcept {
 // =========================================================
 
 template<typename BitMutator>
-static ASMJIT_INLINE void BaseRAPass_calculateInitialInOut(RABlock* block, size_t live_word_count, Support::FixedStack<RABlock*>& queue, uint32_t pov_index) noexcept {
+static ASMJIT_INLINE void BaseRAPass_calculate_initial_in_out(RABlock* block, size_t live_word_count, Support::FixedStack<RABlock*>& queue, uint32_t pov_index) noexcept {
   using BitWord = Support::BitWord;
 
   // Calculate LIVE-OUT based on LIVE-IN of all successors and then recalculate LIVE-IN based on LIVE-OUT and KILL bits.
@@ -1070,7 +1072,7 @@ static ASMJIT_INLINE Error BaseRAPass_calculateInOutKill(
 
     for (size_t pov_index = 0u; pov_index < pov.size(); pov_index++) {
       RABlock* block = pov[pov_index];
-      BaseRAPass_calculateInitialInOut<BitMutator>(block, multi_work_reg_count_as_bit_words, queue, uint32_t(pov_index));
+      BaseRAPass_calculate_initial_in_out<BitMutator>(block, multi_work_reg_count_as_bit_words, queue, uint32_t(pov_index));
     }
 
     // Iteratively keep recalculating LIVE-IN and LIVE-OUT once there are no more changes to the bits. This is

src/asmjit/core/rapass_p.h

@@ -64,8 +64,10 @@ public:
   FuncNode* _func = nullptr;
   //! Stop node.
   BaseNode* _stop = nullptr;
-  //! Node that is used to insert extra code after the function body.
-  BaseNode* _extra_block = nullptr;
+  //! Start of the code that was injected.
+  BaseNode* _injection_start = nullptr;
+  //! End of the code that was injected.
+  BaseNode* _injection_end = nullptr;
 
   //! Blocks (first block is the entry, always exists).
   ArenaVector<RABlock*> _blocks {};
@@ -185,13 +187,6 @@ public:
   [[nodiscard]]
   ASMJIT_INLINE_NODEBUG BaseNode* stop() const noexcept { return _stop; }
 
-  //! Returns an extra block used by the current function being processed.
-  [[nodiscard]]
-  ASMJIT_INLINE_NODEBUG BaseNode* extra_block() const noexcept { return _extra_block; }
-
-  //! Sets an extra block, see `extra_block()`.
-  ASMJIT_INLINE_NODEBUG void set_extra_block(BaseNode* node) noexcept { _extra_block = node; }
-
   [[nodiscard]]
   ASMJIT_INLINE_NODEBUG uint32_t end_position() const noexcept { return _instruction_count * 2u; }
 
@@ -638,6 +633,42 @@ public:
   //! \name Instruction Rewriter
   //! \{
 
+  template<typename Lambda>
+  ASMJIT_INLINE Error rewrite_iterate(Lambda&& fn) noexcept {
+    // First iteration is not a block - it's possible register/stack changes at the end of the function.
+    RABlock* block = nullptr;
+
+    BaseNode* first = _injection_start;
+    BaseNode* stop = _injection_end;
+
+    Span<RABlock*> pov = _pov.as_span();
+    size_t pov_index = pov.size();
+
+    // If no injection happened, just do basic blocks.
+    if (first == nullptr) {
+      // The size of POV is always greater than zero, because there is always at least one basic block.
+      ASMJIT_ASSERT(pov_index > 0u);
+
+      block = pov[--pov_index];
+      first = block->first();
+      stop = block->last()->next();
+    }
+
+    for (;;) {
+      ASMJIT_PROPAGATE(fn(first, stop, block));
+
+      if (!pov_index) {
+        break;
+      }
+
+      block = pov[--pov_index];
+      first = block->first();
+      stop = block->last()->next();
+    }
+
+    return Error::kOk;
+  }
+
   [[nodiscard]]
   virtual Error rewrite() noexcept;
 

src/asmjit/core/string.cpp

@@ -88,16 +88,23 @@ Error String::clear() noexcept {
 // ================
 
 char* String::prepare(ModifyOp op, size_t size) noexcept {
+  uint8_t type = _type;
+
   char* cur_data;
   size_t cur_size;
   size_t cur_capacity;
 
-  if (is_large_or_external()) {
+  if (is_large_or_external(type)) {
     cur_data = _large.data;
     cur_size = _large.size;
     cur_capacity = _large.capacity;
   }
   else {
+    // For some reason clang's static analysis flags this function having "use-after-free". The step
+    // to get to that is to execute this branch (`is_large_or_external()` returning false) and then
+    // assuming `type == kTypeLarge` in another condition, which contradicts the first condition.
+    ASMJIT_ASSERT(type < kTypeLarge);
+
     cur_data = _small.data;
     cur_size = _small.type;
     cur_capacity = kSSOCapacity;
@@ -117,7 +124,7 @@ char* String::prepare(ModifyOp op, size_t size) noexcept {
         return nullptr;
       }
 
-      if (_type == kTypeLarge) {
+      if (type == kTypeLarge) {
         ::free(cur_data);
       }
 
@@ -159,7 +166,7 @@ char* String::prepare(ModifyOp op, size_t size) noexcept {
 
       memcpy(new_data, cur_data, cur_size);
 
-      if (_type == kTypeLarge) {
+      if (type == kTypeLarge) {
         ::free(cur_data);
       }
 
@@ -183,6 +190,7 @@ char* String::prepare(ModifyOp op, size_t size) noexcept {
 // ===============
 
 Error String::assign(const char* data, size_t size) noexcept {
+  uint8_t type = _type;
   char* dst = nullptr;
 
   // Null terminated string without `size` specified.
@@ -190,7 +198,7 @@ Error String::assign(const char* data, size_t size) noexcept {
     size = data ? strlen(data) : size_t(0);
   }
 
-  if (is_large_or_external()) {
+  if (is_large_or_external(type)) {
     if (size <= _large.capacity) {
       dst = _large.data;
       _large.size = size;
@@ -206,7 +214,7 @@ Error String::assign(const char* data, size_t size) noexcept {
         return make_error(Error::kOutOfMemory);
       }
 
-      if (_type == kTypeLarge) {
+      if (type == kTypeLarge) {
         ::free(_large.data);
       }
 

src/asmjit/core/string.h

@@ -91,6 +91,12 @@ public:
   //! External string (arena allocated or not owned by String).
   static inline constexpr uint8_t kTypeExternal = 0x20u;
 
+  [[nodiscard]]
+  static ASMJIT_INLINE_NODEBUG bool is_external(uint8_t type) noexcept { return type == kTypeExternal; }
+
+  [[nodiscard]]
+  static ASMJIT_INLINE_NODEBUG bool is_large_or_external(uint8_t type) noexcept { return type >= kTypeLarge; }
+
   union Raw {
     uint8_t u8[kLayoutSize];
     uint64_t u64[kLayoutSize / sizeof(uint64_t)];
@@ -167,10 +173,10 @@ public:
   //! \{
 
   [[nodiscard]]
-  ASMJIT_INLINE_NODEBUG bool is_external() const noexcept { return _type == kTypeExternal; }
+  ASMJIT_INLINE_NODEBUG bool is_external() const noexcept { return is_external(_type); }
 
   [[nodiscard]]
-  ASMJIT_INLINE_NODEBUG bool is_large_or_external() const noexcept { return _type >= kTypeLarge; }
+  ASMJIT_INLINE_NODEBUG bool is_large_or_external() const noexcept { return is_large_or_external(_type); }
 
   //! Tests whether the string is empty.
   [[nodiscard]]

src/asmjit/x86/x86rapass.cpp

@@ -1350,13 +1350,9 @@ static InstId transform_vex_to_evex(InstId inst_id) {
 }
 
 ASMJIT_FAVOR_SPEED Error X86RAPass::rewrite() noexcept {
-  size_t virt_count = cc()._virt_regs.size();
-
-  for (RABlock* block : _pov.iterate_reverse()) {
-    BaseNode* node = block->first();
-    BaseNode* stop = block->last();
-
-    for (;;) {
+  const size_t virt_count = cc()._virt_regs.size();
+  return rewrite_iterate([&](BaseNode* node, BaseNode* stop, RABlock* block) noexcept -> Error {
+    while (node != stop) {
       BaseNode* next = node->next();
       if (node->is_inst()) {
         InstNode* inst = node->as<InstNode>();
@@ -1460,7 +1456,8 @@ ASMJIT_FAVOR_SPEED Error X86RAPass::rewrite() noexcept {
             if (operands.size() == 2u) {
               if (operands[0] == operands[1]) {
                 cc().remove_node(node);
-                goto Next;
+                node = next;
+                continue;
               }
             }
           }
@@ -1476,7 +1473,10 @@ ASMJIT_FAVOR_SPEED Error X86RAPass::rewrite() noexcept {
 
               BaseNode* prev = node->prev();
               cc().remove_node(node);
-              block->set_last(prev);
+
+              if (block) {
+                block->set_last(prev);
+              }
             }
           }
         }
@@ -1506,16 +1506,11 @@ ASMJIT_FAVOR_SPEED Error X86RAPass::rewrite() noexcept {
         }
       }
 
-Next:
-      if (node == stop) {
-        break;
-      }
-
       node = next;
     }
-  }
 
-  return Error::kOk;
+    return Error::kOk;
+  });
 }
 
 // x86::X86RAPass - OnEmit