Fix MachBuffer branch optimization.

This patch fixes a subtle bug that occurred in the MachBuffer branch
optimization: in tracking labels at the current buffer tail using a
sorted-by-offset array, the code did not update this array properly when
redirecting labels. As a result, the dead-branch removal was unsafe,
because not every label pointing to a branch is guaranteed to be
redirected properly first.

Discovered while doing performance testing: bz2 silently took a wrong
branch and exited compression early. (Eek!)

To address this problem, this patch adopts a slightly simpler data
structure: we only track the labels *at the current buffer tail*, and
*at the start of each branch*, and we're careful to update these
appropriately to maintain the invariants. I'm pretty confident that this
is correct now, but we should (still) fuzz it a bunch, because wrong
control flow scares me a nonzero amount. I should probably also actually
write out a formal proof that these data-structure updates are correct.
The optimizations are important for performance (removing useless empty
blocks, and taking advantage of any fallthrough opportunities at all),
so I don't think we would want to drop them entirely.

Author: cfallin

cranelift/codegen/src/isa/aarch64/inst/emit.rs

@@ -1163,17 +1163,13 @@ impl MachInstEmit for Inst {
             }
             &Inst::Jump { ref dest } => {
                 let off = sink.cur_offset();
-                // Emit the jump itself.
-                sink.put4(enc_jump26(0b000101, dest.as_offset26_or_zero()));
-                // After the jump has been emitted, indicate that it uses a
-                // label, if so, so that a fixup can occur later. This happens
-                // after we emit the bytes because the fixup might occur right
-                // away (so the bytes must actually exist now).
+                // Indicate that the jump uses a label, if so, so that a fixup can occur later.
                 if let Some(l) = dest.as_label() {
                     sink.use_label_at_offset(off, l, LabelUse::Branch26);
-                    let cur_off = sink.cur_offset();
-                    sink.add_uncond_branch(off, cur_off, l);
+                    sink.add_uncond_branch(off, off + 4, l);
                 }
+                // Emit the jump itself.
+                sink.put4(enc_jump26(0b000101, dest.as_offset26_or_zero()));
             }
             &Inst::Ret => {
                 sink.put4(0xd65f03c0);
@@ -1208,28 +1204,27 @@ impl MachInstEmit for Inst {
             } => {
                 // Conditional part first.
                 let cond_off = sink.cur_offset();
-                sink.put4(enc_conditional_br(taken, kind));
                 if let Some(l) = taken.as_label() {
                     sink.use_label_at_offset(cond_off, l, LabelUse::Branch19);
-                    let cur_off = sink.cur_offset();
                     let inverted = enc_conditional_br(taken, kind.invert()).to_le_bytes();
-                    sink.add_cond_branch(cond_off, cur_off, l, &inverted[..]);
+                    sink.add_cond_branch(cond_off, cond_off + 4, l, &inverted[..]);
                 }
-                // Unconditional part.
+                sink.put4(enc_conditional_br(taken, kind));
+
+                // Unconditional part next.
                 let uncond_off = sink.cur_offset();
-                sink.put4(enc_jump26(0b000101, not_taken.as_offset26_or_zero()));
                 if let Some(l) = not_taken.as_label() {
                     sink.use_label_at_offset(uncond_off, l, LabelUse::Branch26);
-                    let cur_off = sink.cur_offset();
-                    sink.add_uncond_branch(uncond_off, cur_off, l);
+                    sink.add_uncond_branch(uncond_off, uncond_off + 4, l);
                 }
+                sink.put4(enc_jump26(0b000101, not_taken.as_offset26_or_zero()));
             }
             &Inst::OneWayCondBr { target, kind } => {
                 let off = sink.cur_offset();
-                sink.put4(enc_conditional_br(target, kind));
                 if let Some(l) = target.as_label() {
                     sink.use_label_at_offset(off, l, LabelUse::Branch19);
                 }
+                sink.put4(enc_conditional_br(target, kind));
             }
             &Inst::IndirectBr { rn, .. } => {
                 sink.put4(enc_br(rn));
@@ -1307,12 +1302,12 @@ impl MachInstEmit for Inst {
                 for &target in targets.iter() {
                     let word_off = sink.cur_offset();
                     let off_into_table = word_off - jt_off;
-                    sink.put4(off_into_table);
                     sink.use_label_at_offset(
                         word_off,
                         target.as_label().unwrap(),
                         LabelUse::PCRel32,
                     );
+                    sink.put4(off_into_table);
                 }
 
                 // Lowering produces an EmitIsland before using a JTSequence, so we can safely

cranelift/codegen/src/isa/x64/inst/emit.rs

@@ -812,14 +812,14 @@ pub(crate) fn emit(inst: &Inst, sink: &mut MachBuffer<Inst>) {
             let disp = dest.as_offset32_or_zero() - 5;
             let disp = disp as u32;
             let br_start = sink.cur_offset();
-            sink.put1(0xE9);
-            let br_disp_off = sink.cur_offset();
-            sink.put4(disp);
-            let br_end = sink.cur_offset();
+            let br_disp_off = br_start + 1;
+            let br_end = br_start + 5;
             if let Some(l) = dest.as_label() {
                 sink.use_label_at_offset(br_disp_off, l, LabelUse::Rel32);
                 sink.add_uncond_branch(br_start, br_end, l);
             }
+            sink.put1(0xE9);
+            sink.put4(disp);
         }
         Inst::JmpCondSymm {
             cc,
@@ -835,31 +835,31 @@ pub(crate) fn emit(inst: &Inst, sink: &mut MachBuffer<Inst>) {
             let taken_disp = taken.as_offset32_or_zero() - 6;
             let taken_disp = taken_disp as u32;
             let cond_start = sink.cur_offset();
-            sink.put1(0x0F);
-            sink.put1(0x80 + cc.get_enc());
-            let cond_disp_off = sink.cur_offset();
-            sink.put4(taken_disp);
-            let cond_end = sink.cur_offset();
+            let cond_disp_off = cond_start + 2;
+            let cond_end = cond_start + 6;
             if let Some(l) = taken.as_label() {
                 sink.use_label_at_offset(cond_disp_off, l, LabelUse::Rel32);
                 let inverted: [u8; 6] =
                     [0x0F, 0x80 + (cc.invert().get_enc()), 0xFA, 0xFF, 0xFF, 0xFF];
                 sink.add_cond_branch(cond_start, cond_end, l, &inverted[..]);
             }
+            sink.put1(0x0F);
+            sink.put1(0x80 + cc.get_enc());
+            sink.put4(taken_disp);
 
             // Unconditional part.
 
             let nt_disp = not_taken.as_offset32_or_zero() - 5;
             let nt_disp = nt_disp as u32;
             let uncond_start = sink.cur_offset();
-            sink.put1(0xE9);
-            let uncond_disp_off = sink.cur_offset();
-            sink.put4(nt_disp);
-            let uncond_end = sink.cur_offset();
+            let uncond_disp_off = uncond_start + 1;
+            let uncond_end = uncond_start + 5;
             if let Some(l) = not_taken.as_label() {
                 sink.use_label_at_offset(uncond_disp_off, l, LabelUse::Rel32);
                 sink.add_uncond_branch(uncond_start, uncond_end, l);
             }
+            sink.put1(0xE9);
+            sink.put4(nt_disp);
         }
         Inst::JmpUnknown { target } => {
             match target {