[asan-asm-instrumentation] CFI directives are generated for .S files.

Summary: CFI directives are generated for .S files.

Reviewers: eugenis

Subscribers: llvm-commits

Differential Revision: http://reviews.llvm.org/D5520

llvm-svn: 219199

4 files changed, 63 insertions, 10 deletions

diff --git a/llvm/lib/MC/MCStreamer.cpp b/llvm/lib/MC/MCStreamer.cpp
index f6620f07240..2bfb9ba1599 100644
--- a/llvm/lib/MC/MCStreamer.cpp
+++ b/llvm/lib/MC/MCStreamer.cpp
@@ -220,6 +220,16 @@ void MCStreamer::EmitCFIStartProc(bool IsSimple) {
   Frame.IsSimple = IsSimple;
   EmitCFIStartProcImpl(Frame);
 
+  const MCAsmInfo* MAI = Context.getAsmInfo();
+  if (MAI) {
+    for (const MCCFIInstruction& Inst : MAI->getInitialFrameState()) {
+      if (Inst.getOperation() == MCCFIInstruction::OpDefCfa ||
+          Inst.getOperation() == MCCFIInstruction::OpDefCfaRegister) {
+        Frame.CurrentCfaRegister = Inst.getRegister();
+      }
+    }
+  }
+
   DwarfFrameInfos.push_back(Frame);
 }
 
@@ -251,6 +261,7 @@ void MCStreamer::EmitCFIDefCfa(int64_t Register, int64_t Offset) {
     MCCFIInstruction::createDefCfa(Label, Register, Offset);
   MCDwarfFrameInfo *CurFrame = getCurrentDwarfFrameInfo();
   CurFrame->Instructions.push_back(Instruction);
+  CurFrame->CurrentCfaRegister = static_cast<unsigned>(Register);
 }
 
 void MCStreamer::EmitCFIDefCfaOffset(int64_t Offset) {
@@ -275,6 +286,7 @@ void MCStreamer::EmitCFIDefCfaRegister(int64_t Register) {
     MCCFIInstruction::createDefCfaRegister(Label, Register);
   MCDwarfFrameInfo *CurFrame = getCurrentDwarfFrameInfo();
   CurFrame->Instructions.push_back(Instruction);
+  CurFrame->CurrentCfaRegister = static_cast<unsigned>(Register);
 }
 
 void MCStreamer::EmitCFIOffset(int64_t Register, int64_t Offset) {
diff --git a/llvm/lib/Target/X86/AsmParser/X86AsmInstrumentation.cpp b/llvm/lib/Target/X86/AsmParser/X86AsmInstrumentation.cpp
index 68c1b2d8b23..974ee20461a 100644
--- a/llvm/lib/Target/X86/AsmParser/X86AsmInstrumentation.cpp
+++ b/llvm/lib/Target/X86/AsmParser/X86AsmInstrumentation.cpp
@@ -15,6 +15,7 @@
 #include "llvm/ADT/Triple.h"
 #include "llvm/CodeGen/MachineValueType.h"
 #include "llvm/IR/Function.h"
+#include "llvm/MC/MCAsmInfo.h"
 #include "llvm/MC/MCContext.h"
 #include "llvm/MC/MCInst.h"
 #include "llvm/MC/MCInstBuilder.h"
@@ -297,6 +298,13 @@ public:
 
   virtual ~X86AddressSanitizer32() {}
 
+  unsigned GetFrameReg(const MCContext &Ctx, MCStreamer &Out) {
+    unsigned FrameReg = GetFrameRegGeneric(Ctx, Out);
+    if (FrameReg == X86::NoRegister)
+      return FrameReg;
+    return getX86SubSuperRegister(FrameReg, MVT::i32);
+  }
+
   virtual void StoreFlags(MCStreamer &Out) override {
     EmitInstruction(Out, MCInstBuilder(X86::PUSHF32));
   }
@@ -308,7 +316,8 @@ public:
   virtual void InstrumentMemOperandPrologue(const RegisterContext &RegCtx,
                                             MCContext &Ctx,
                                             MCStreamer &Out) override {
-    const MCRegisterInfo* MRI = Ctx.getRegisterInfo();
+    const MCRegisterInfo *MRI = Ctx.getRegisterInfo();
+    unsigned FrameReg = GetFrameReg(Ctx, Out);
     if (MRI && FrameReg != X86::NoRegister) {
       EmitInstruction(
           Out, MCInstBuilder(X86::PUSH32r).addReg(X86::EBP));
@@ -348,6 +357,7 @@ public:
     EmitInstruction(
         Out, MCInstBuilder(X86::POP32r).addReg(RegCtx.addressReg(MVT::i32)));
 
+    unsigned FrameReg = GetFrameReg(Ctx, Out);
     if (Ctx.getRegisterInfo() && FrameReg != X86::NoRegister) {
       EmitInstruction(
           Out, MCInstBuilder(X86::POP32r).addReg(X86::EBP));
@@ -554,6 +564,13 @@ public:
 
   virtual ~X86AddressSanitizer64() {}
 
+  unsigned GetFrameReg(const MCContext &Ctx, MCStreamer &Out) {
+    unsigned FrameReg = GetFrameRegGeneric(Ctx, Out);
+    if (FrameReg == X86::NoRegister)
+      return FrameReg;
+    return getX86SubSuperRegister(FrameReg, MVT::i64);
+  }
+
   virtual void StoreFlags(MCStreamer &Out) override {
     EmitInstruction(Out, MCInstBuilder(X86::PUSHF64));
   }
@@ -565,19 +582,20 @@ public:
   virtual void InstrumentMemOperandPrologue(const RegisterContext &RegCtx,
                                             MCContext &Ctx,
                                             MCStreamer &Out) override {
-    const MCRegisterInfo *RegisterInfo = Ctx.getRegisterInfo();
-    if (RegisterInfo && FrameReg != X86::NoRegister) {
+    const MCRegisterInfo *MRI = Ctx.getRegisterInfo();
+    unsigned FrameReg = GetFrameReg(Ctx, Out);
+    if (MRI && FrameReg != X86::NoRegister) {
       EmitInstruction(Out, MCInstBuilder(X86::PUSH64r).addReg(X86::RBP));
       if (FrameReg == X86::RSP) {
         Out.EmitCFIAdjustCfaOffset(8 /* byte size of the FrameReg */);
         Out.EmitCFIRelOffset(
-            RegisterInfo->getDwarfRegNum(X86::RBP, true /* IsEH */), 0);
+            MRI->getDwarfRegNum(X86::RBP, true /* IsEH */), 0);
       }
       EmitInstruction(
           Out, MCInstBuilder(X86::MOV64rr).addReg(X86::RBP).addReg(FrameReg));
       Out.EmitCFIRememberState();
       Out.EmitCFIDefCfaRegister(
-          RegisterInfo->getDwarfRegNum(X86::RBP, true /* IsEH */));
+          MRI->getDwarfRegNum(X86::RBP, true /* IsEH */));
     }
 
     EmitAdjustRSP(Ctx, Out, -128);
@@ -606,6 +624,7 @@ public:
         Out, MCInstBuilder(X86::POP64r).addReg(RegCtx.shadowReg(MVT::i64)));
     EmitAdjustRSP(Ctx, Out, 128);
 
+    unsigned FrameReg = GetFrameReg(Ctx, Out);
     if (Ctx.getRegisterInfo() && FrameReg != X86::NoRegister) {
       EmitInstruction(
           Out, MCInstBuilder(X86::POP64r).addReg(X86::RBP));
@@ -820,7 +839,7 @@ void X86AddressSanitizer64::InstrumentMOVSImpl(unsigned AccessSize,
 } // End anonymous namespace
 
 X86AsmInstrumentation::X86AsmInstrumentation(const MCSubtargetInfo &STI)
-    : STI(STI), FrameReg(X86::NoRegister) {}
+    : STI(STI), InitialFrameReg(0) {}
 
 X86AsmInstrumentation::~X86AsmInstrumentation() {}
 
@@ -835,6 +854,25 @@ void X86AsmInstrumentation::EmitInstruction(MCStreamer &Out,
   Out.EmitInstruction(Inst, STI);
 }
 
+unsigned X86AsmInstrumentation::GetFrameRegGeneric(const MCContext &Ctx,
+                                                   MCStreamer &Out) {
+  if (!Out.getNumFrameInfos()) // No active dwarf frame
+    return X86::NoRegister;
+  const MCDwarfFrameInfo &Frame = Out.getDwarfFrameInfos().back();
+  if (Frame.End) // Active dwarf frame is closed
+    return X86::NoRegister;
+  const MCRegisterInfo *MRI = Ctx.getRegisterInfo();
+  if (!MRI) // No register info
+    return X86::NoRegister;
+
+  if (InitialFrameReg) {
+    // FrameReg is set explicitly, we're instrumenting a MachineFunction.
+    return InitialFrameReg;
+  }
+
+  return MRI->getLLVMRegNum(Frame.CurrentCfaRegister, true /* IsEH */);
+}
+
 X86AsmInstrumentation *
 CreateX86AsmInstrumentation(const MCTargetOptions &MCOptions,
                             const MCContext &Ctx, const MCSubtargetInfo &STI) {
diff --git a/llvm/lib/Target/X86/AsmParser/X86AsmInstrumentation.h b/llvm/lib/Target/X86/AsmParser/X86AsmInstrumentation.h
index aa2ed03a19b..19ebcc44f61 100644
--- a/llvm/lib/Target/X86/AsmParser/X86AsmInstrumentation.h
+++ b/llvm/lib/Target/X86/AsmParser/X86AsmInstrumentation.h
@@ -34,8 +34,9 @@ class X86AsmInstrumentation {
 public:
   virtual ~X86AsmInstrumentation();
 
-  void SetFrameRegister(unsigned RegNo) {
-    FrameReg = RegNo;
+  // Sets frame register corresponding to a current frame.
+  void SetInitialFrameRegister(unsigned RegNo) {
+    InitialFrameReg = RegNo;
   }
 
   // Tries to instrument and emit instruction.
@@ -51,11 +52,13 @@ protected:
 
   X86AsmInstrumentation(const MCSubtargetInfo &STI);
 
+  unsigned GetFrameRegGeneric(const MCContext &Ctx, MCStreamer &Out);
+
   void EmitInstruction(MCStreamer &Out, const MCInst &Inst);
 
   const MCSubtargetInfo &STI;
 
-  unsigned FrameReg;
+  unsigned InitialFrameReg;
 };
 
 } // End llvm namespace
diff --git a/llvm/lib/Target/X86/AsmParser/X86AsmParser.cpp b/llvm/lib/Target/X86/AsmParser/X86AsmParser.cpp
index 519ada2b4b9..0a2ad104f37 100644
--- a/llvm/lib/Target/X86/AsmParser/X86AsmParser.cpp
+++ b/llvm/lib/Target/X86/AsmParser/X86AsmParser.cpp
@@ -973,7 +973,7 @@ bool X86AsmParser::ParseRegister(unsigned &RegNo,
 }
 
 void X86AsmParser::SetFrameRegister(unsigned RegNo) {
-  Instrumentation->SetFrameRegister(RegNo);
+  Instrumentation->SetInitialFrameRegister(RegNo);
 }
 
 std::unique_ptr<X86Operand> X86AsmParser::DefaultMemSIOperand(SMLoc Loc) {