[JITLink][RISCV] Adjust offsets of non-relaxable edges

[jobnoorman]

Summary

The relaxation algorithm used to only update offsets of relaxable edges.
This caused non-relaxable edges that appear after a relaxed instruction
to have an incorrect offset and be applied at the wrong location. This
patch fixes this by updating the offsets of all edges.

Note that this bug was caused by an incorrect translation of LLD's
relaxation algorithm. LLD always uses all edges during relaxation while
I decided to filter-out relaxable edges to prevent having to iterate
non-relaxable edges at each step. However, this had the side-effect of
only updating offsets of relaxable edges. This patch leaves the
filtering of relaxable edges as-is but iterates all edges when updating
offsets.

[StephenFan]

LGTM.

llvm/lib/ExecutionEngine/JITLink/ELF_riscv.cpp

for (auto [I, E] : llvm::enumerate(Aux.RelaxEdges)) {
     Dest += Skip;
     Offset = E->getOffset() + Skip + Remove;
   }
 
   std::memmove(Dest, Contents.data() + Offset, Contents.size() - Offset);
 
   // Fixup edge offsets and kinds.
   Delta = 0;
-  for (auto [I, E] : llvm::enumerate(Aux.RelaxEdges)) {
-    E->setOffset(E->getOffset() - Delta);
+  size_t I = 0;
+  for (auto &E : Block.edges()) {
+    E.setOffset(E.getOffset() - Delta);
 
+    if (I < Aux.RelaxEdges.size() && Aux.RelaxEdges[I] == &E) {
       if (Aux.EdgeKinds[I] != Edge::Invalid)
-      E->setKind(Aux.EdgeKinds[I]);
+        E.setKind(Aux.EdgeKinds[I]);
 
       Delta = Aux.RelocDeltas[I];
+      ++I;
+    }
   }
 
   // Remove AlignRelaxable edges: all other relaxable edges got modified and
   // will be used later while linking. Alignment is entirely handled here so we
   // don't need these edges anymore.
   for (auto *B : G.blocks()) {
     for (auto IE = B->edges().begin(); IE != B->edges().end();) {
       if (IE->getKind() == AlignRelaxable)

llvm/test/ExecutionEngine/JITLink/RISCV/ELF_relax_nonrelaxable.s

+## Test that non-relaxable edges have their offset adjusted by relaxation
+
+# RUN: rm -rf %t && mkdir %t && cd %t
+
+# RUN: llvm-mc -filetype=obj -triple=riscv32 -mattr=+relax %s -o %t.rv32
+# RUN: llvm-jitlink -noexec \
+# RUN:     -slab-allocate 100Kb -slab-address 0x0 -slab-page-size 4096 \
+# RUN:     -check %s %t.rv32
+
+# RUN: llvm-mc -filetype=obj -triple=riscv64 -mattr=+relax %s -o %t.rv64
+# RUN: llvm-jitlink -noexec \
+# RUN:     -slab-allocate 100Kb -slab-address 0x0 -slab-page-size 4096 \
+# RUN:     -check %s %t.rv64
+
+    .globl main,nonrelaxable,nonrelaxable_target
+    .size nonrelaxable, 4
+    .size nonrelaxable_target, 4
+main:
+    call f
+nonrelaxable:
+    ## Non-relaxable R_RISCV_BRANCH edge after a relaxable R_RISCV_CALL edge.
+    ## Even though this edge isn't relaxable, its offset should still be
+    ## adjusted.
+    beq zero, zero, nonrelaxable_target
+nonrelaxable_target:
+    ret
+    .size main, .-main
+
+    .globl f
+f:
+    ret
+    .size f, .-f
+
+# jitlink-check: decode_operand(nonrelaxable, 2) = (nonrelaxable_target - nonrelaxable)
+