6 files changed, 26 insertions, 19 deletions
diff --git a/llvm/lib/Target/X86/X86.h b/llvm/lib/Target/X86/X86.h
index 8d2805ccc20..e6deda620e2 100644
--- a/llvm/lib/Target/X86/X86.h
+++ b/llvm/lib/Target/X86/X86.h
@@ -44,7 +44,8 @@ FunctionPass *createX86CodePrinterPass(std::ostream &o, X86TargetMachine &tm);
 
 /// createX86CodeEmitterPass - Return a pass that emits the collected X86 code
 /// to the specified MCE object.
-FunctionPass *createX86CodeEmitterPass(MachineCodeEmitter &MCE);
+FunctionPass *createX86CodeEmitterPass(X86TargetMachine &TM,
+                                       MachineCodeEmitter &MCE);
 
 /// addX86ELFObjectWriterPass - Add passes to the FPM that output the generated
 /// code as an ELF object file.
diff --git a/llvm/lib/Target/X86/X86CodeEmitter.cpp b/llvm/lib/Target/X86/X86CodeEmitter.cpp
index 9179890b39b..31b4bdf1852 100644
--- a/llvm/lib/Target/X86/X86CodeEmitter.cpp
+++ b/llvm/lib/Target/X86/X86CodeEmitter.cpp
@@ -35,12 +35,14 @@ namespace {
 namespace {
   class VISIBILITY_HIDDEN Emitter : public MachineFunctionPass {
     const X86InstrInfo  *II;
+    TargetMachine &TM;
     MachineCodeEmitter  &MCE;
-    std::vector<std::pair<MachineBasicBlock *, unsigned> > BBRefs;
   public:
-    explicit Emitter(MachineCodeEmitter &mce) : II(0), MCE(mce) {}
-    Emitter(MachineCodeEmitter &mce, const X86InstrInfo& ii)
-        : II(&ii), MCE(mce) {}
+    explicit Emitter(TargetMachine &tm, MachineCodeEmitter &mce)
+      : II(0), TM(tm), MCE(mce) {}
+    Emitter(TargetMachine &tm, MachineCodeEmitter &mce,
+            const X86InstrInfo& ii)
+      : II(&ii), TM(tm), MCE(mce) {}
 
     bool runOnMachineFunction(MachineFunction &MF);
 
@@ -71,8 +73,9 @@ namespace {
 
 /// createX86CodeEmitterPass - Return a pass that emits the collected X86 code
 /// to the specified MCE object.
-FunctionPass *llvm::createX86CodeEmitterPass(MachineCodeEmitter &MCE) {
-  return new Emitter(MCE);
+FunctionPass *llvm::createX86CodeEmitterPass(X86TargetMachine &TM,
+                                             MachineCodeEmitter &MCE) {
+  return new Emitter(TM, MCE);
 }
 
 bool Emitter::runOnMachineFunction(MachineFunction &MF) {
@@ -82,8 +85,6 @@ bool Emitter::runOnMachineFunction(MachineFunction &MF) {
   II = ((X86TargetMachine&)MF.getTarget()).getInstrInfo();
 
   do {
-    BBRefs.clear();
-
     MCE.startFunction(MF);
     for (MachineFunction::iterator MBB = MF.begin(), E = MF.end(); 
          MBB != E; ++MBB) {
@@ -94,13 +95,6 @@ bool Emitter::runOnMachineFunction(MachineFunction &MF) {
     }
   } while (MCE.finishFunction(MF));
 
-  // Resolve all forward branches now.
-  for (unsigned i = 0, e = BBRefs.size(); i != e; ++i) {
-    unsigned Location = MCE.getMachineBasicBlockAddress(BBRefs[i].first);
-    unsigned Ref = BBRefs[i].second;
-    *((unsigned*)(intptr_t)Ref) = Location-Ref-4;
-  }
-  BBRefs.clear();
   return false;
 }
 
@@ -117,7 +111,7 @@ void Emitter::emitPCRelativeValue(unsigned Address) {
 void Emitter::emitPCRelativeBlockAddress(MachineBasicBlock *MBB) {
   // Remember where this reference was and where it is to so we can
   // deal with it later.
-  BBRefs.push_back(std::make_pair(MBB, MCE.getCurrentPCValue()));
+  TM.getJITInfo()->addBBRef(MBB, MCE.getCurrentPCValue());
   MCE.emitWordLE(0);
 }
 
diff --git a/llvm/lib/Target/X86/X86ELFWriter.cpp b/llvm/lib/Target/X86/X86ELFWriter.cpp
index 0a3a027f1d4..46c46420301 100644
--- a/llvm/lib/Target/X86/X86ELFWriter.cpp
+++ b/llvm/lib/Target/X86/X86ELFWriter.cpp
@@ -35,5 +35,5 @@ void llvm::addX86ELFObjectWriterPass(PassManager &FPM,
                                      std::ostream &O, X86TargetMachine &TM) {
   X86ELFWriter *EW = new X86ELFWriter(O, TM);
   FPM.add(EW);
-  FPM.add(createX86CodeEmitterPass(EW->getMachineCodeEmitter()));
+  FPM.add(createX86CodeEmitterPass(TM, EW->getMachineCodeEmitter()));
 }
diff --git a/llvm/lib/Target/X86/X86JITInfo.cpp b/llvm/lib/Target/X86/X86JITInfo.cpp
index 9bb2a725233..9fd8029962d 100644
--- a/llvm/lib/Target/X86/X86JITInfo.cpp
+++ b/llvm/lib/Target/X86/X86JITInfo.cpp
@@ -203,3 +203,13 @@ void X86JITInfo::relocate(void *Function, MachineRelocation *MR,
     }
   }
 }
+
+void X86JITInfo::resolveBBRefs(MachineCodeEmitter &MCE) {
+  // Resolve all forward branches now.
+  for (unsigned i = 0, e = BBRefs.size(); i != e; ++i) {
+    unsigned Location = MCE.getMachineBasicBlockAddress(BBRefs[i].first);
+    intptr_t Ref = BBRefs[i].second;
+    *((unsigned*)Ref) = Location-Ref-4;
+  }
+  BBRefs.clear();
+}
diff --git a/llvm/lib/Target/X86/X86JITInfo.h b/llvm/lib/Target/X86/X86JITInfo.h
index 02e54af11af..f9e437e41b1 100644
--- a/llvm/lib/Target/X86/X86JITInfo.h
+++ b/llvm/lib/Target/X86/X86JITInfo.h
@@ -51,6 +51,8 @@ namespace llvm {
     /// referenced global symbols.
     virtual void relocate(void *Function, MachineRelocation *MR,
                           unsigned NumRelocs, unsigned char* GOTBase);
+
+    virtual void resolveBBRefs(MachineCodeEmitter &MCE);
   };
 }
 
diff --git a/llvm/lib/Target/X86/X86TargetMachine.cpp b/llvm/lib/Target/X86/X86TargetMachine.cpp
index 2c3c5306847..e4b6a941a46 100644
--- a/llvm/lib/Target/X86/X86TargetMachine.cpp
+++ b/llvm/lib/Target/X86/X86TargetMachine.cpp
@@ -194,7 +194,7 @@ void X86JITInfo::addPassesToJITCompile(FunctionPassManager &PM) {
 
 bool X86TargetMachine::addPassesToEmitMachineCode(FunctionPassManager &PM,
                                                   MachineCodeEmitter &MCE) {
-  PM.add(createX86CodeEmitterPass(MCE));
+  PM.add(createX86CodeEmitterPass(*this, MCE));
   // Delete machine code for this function
   PM.add(createMachineCodeDeleter());
   return false;