Properly preserve and restore CFA state in FDE (#1373)

* Properly preserve and restore CFA state in FDE

cranelift/codegen/src/isa/x86/abi.rs

@@ -835,25 +835,7 @@ fn insert_common_epilogues(
         pos.goto_last_inst(ebb);
         if let Some(inst) = pos.current_inst() {
             if pos.func.dfg[inst].opcode().is_return() {
-                if let (Some(ref mut frame_layout), ref func_layout) =
+                let is_last = pos.func.layout.last_ebb() == Some(ebb);
-                    (pos.func.frame_layout.as_mut(), &pos.func.layout)
-                {
-                    // Figure out if we need to insert end-of-function-aware frame layout information.
-                    let following_inst = func_layout
-                        .next_ebb(ebb)
-                        .and_then(|next_ebb| func_layout.first_inst(next_ebb));
-
-                    if let Some(following_inst) = following_inst {
-                        frame_layout
-                            .instructions
-                            .insert(inst, vec![FrameLayoutChange::Preserve].into_boxed_slice());
-                        frame_layout.instructions.insert(
-                            following_inst,
-                            vec![FrameLayoutChange::Restore].into_boxed_slice(),
-                        );
-                    }
-                }
-
                 insert_common_epilogue(
                     inst,
                     stack_size,
@@ -861,6 +843,7 @@ fn insert_common_epilogues(
                     reg_type,
                     csrs,
                     isa,
+                    is_last,
                     cfa_state.clone(),
                 );
             }
@@ -877,6 +860,7 @@ fn insert_common_epilogue(
     reg_type: ir::types::Type,
     csrs: &RegisterSet,
     isa: &dyn TargetIsa,
+    is_last: bool,
     mut cfa_state: Option<CFAState>,
 ) {
     let word_size = isa.pointer_bytes() as isize;
@@ -931,10 +915,16 @@ fn insert_common_epilogue(
         assert_eq!(cfa_state.current_depth, -word_size);
         assert_eq!(cfa_state.cf_ptr_offset, word_size);
 
+        // Inserting preserve CFA state operation after FP pop instructions.
         let new_cfa = FrameLayoutChange::CallFrameAddressAt {
             reg: cfa_state.cf_ptr_reg,
             offset: cfa_state.cf_ptr_offset,
         };
+        let new_cfa = if is_last {
+            vec![new_cfa]
+        } else {
+            vec![FrameLayoutChange::Preserve, new_cfa]
+        };
 
         frame_layout
             .instructions
@@ -943,10 +933,17 @@ fn insert_common_epilogue(
                 *insts = insts
                     .iter()
                     .cloned()
-                    .chain(std::iter::once(new_cfa))
+                    .chain(new_cfa.clone().into_iter())
                     .collect::<Box<[_]>>();
             })
-            .or_insert_with(|| Box::new([new_cfa]));
+            .or_insert_with(|| new_cfa.into_boxed_slice());
+
+        if !is_last {
+            // Inserting restore CFA state operation after each return.
+            frame_layout
+                .instructions
+                .insert(inst, vec![FrameLayoutChange::Restore].into_boxed_slice());
+        }
     }
 }
 

cranelift/codegen/src/isa/x86/fde.rs

@@ -163,9 +163,8 @@ fn to_cfi(
             let cfa_offset = *cfa_offset as i32;
             CallFrameInstruction::Offset(X86_64::RA, cfa_offset)
         }
-        _ => {
+        FrameLayoutChange::Preserve => CallFrameInstruction::RememberState,
-            return None;
+        FrameLayoutChange::Restore => CallFrameInstruction::RestoreState,
-        }
     })
 }
 
@@ -265,7 +264,10 @@ mod tests {
     use super::*;
     use crate::binemit::{FrameUnwindOffset, Reloc};
     use crate::cursor::{Cursor, FuncCursor};
-    use crate::ir::{ExternalName, InstBuilder, Signature, StackSlotData, StackSlotKind};
+    use crate::ir::{
+        types, AbiParam, ExternalName, InstBuilder, Signature, StackSlotData, StackSlotKind,
+        TrapCode,
+    };
     use crate::isa::{lookup, CallConv};
     use crate::settings::{builder, Flags};
     use crate::Context;
@@ -352,4 +354,79 @@ mod tests {
 
         func
     }
+
+    #[test]
+    fn test_multi_return_func() {
+        let isa = lookup(triple!("x86_64"))
+            .expect("expect x86 ISA")
+            .finish(Flags::new(builder()));
+
+        let mut context = Context::for_function(create_multi_return_function(CallConv::SystemV));
+        context.func.collect_frame_layout_info();
+
+        context.compile(&*isa).expect("expected compilation");
+
+        let mut sink = SimpleUnwindSink(Vec::new(), 0, Vec::new());
+        emit_fde(&context.func, &*isa, &mut sink);
+
+        assert_eq!(
+            sink.0,
+            vec![
+                20, 0, 0, 0, // CIE len
+                0, 0, 0, 0,   // CIE marker
+                1,   // version
+                0,   // augmentation string
+                1,   // code aligment = 1
+                120, // data alignment = -8
+                16,  // RA = r16
+                0x0c, 0x07, 0x08, // DW_CFA_def_cfa r7, 8
+                0x90, 0x01, //  DW_CFA_offset r16, -8 * 1
+                0, 0, 0, 0, 0, 0, // padding
+                36, 0, 0, 0, // FDE len
+                28, 0, 0, 0, // CIE offset
+                0, 0, 0, 0, 0, 0, 0, 0, // addr reloc
+                15, 0, 0, 0, 0, 0, 0, 0,    // function length
+                0x42, // DW_CFA_advance_loc 2
+                0x0e, 0x10, // DW_CFA_def_cfa_offset 16
+                0x86, 0x02, // DW_CFA_offset r6, -8 * 2
+                0x43, // DW_CFA_advance_loc 3
+                0x0d, 0x06, // DW_CFA_def_cfa_register
+                0x47, // DW_CFA_advance_loc 10
+                0x0a, // DW_CFA_preserve_state
+                0x0c, 0x07, 0x08, // DW_CFA_def_cfa r7, 8
+                0x41, // DW_CFA_advance_loc 1
+                0x0b, // DW_CFA_restore_state
+                // NOTE: no additional CFA directives -- DW_CFA_restore_state
+                // is done before trap and it is last instruction in the function.
+                0, // padding
+                0, 0, 0, 0, // End of FDEs
+            ]
+        );
+        assert_eq!(sink.1, 24);
+        assert_eq!(sink.2.len(), 1);
+    }
+
+    fn create_multi_return_function(call_conv: CallConv) -> Function {
+        let mut sig = Signature::new(call_conv);
+        sig.params.push(AbiParam::new(types::I32));
+        let mut func = Function::with_name_signature(ExternalName::user(0, 0), sig);
+
+        let ebb0 = func.dfg.make_ebb();
+        let v0 = func.dfg.append_ebb_param(ebb0, types::I32);
+        let ebb1 = func.dfg.make_ebb();
+        let ebb2 = func.dfg.make_ebb();
+
+        let mut pos = FuncCursor::new(&mut func);
+        pos.insert_ebb(ebb0);
+        pos.ins().brnz(v0, ebb2, &[]);
+        pos.ins().jump(ebb1, &[]);
+
+        pos.insert_ebb(ebb1);
+        pos.ins().return_(&[]);
+
+        pos.insert_ebb(ebb2);
+        pos.ins().trap(TrapCode::User(0));
+
+        func
+    }
 }