4 files changed, 30 insertions, 5 deletions
diff --git a/llvm/include/llvm/MC/MCFunction.h b/llvm/include/llvm/MC/MCFunction.h
index a87b9488d36..697b5006a9a 100644
--- a/llvm/include/llvm/MC/MCFunction.h
+++ b/llvm/include/llvm/MC/MCFunction.h
@@ -97,6 +97,12 @@ public:
 
   StringRef getName() const { return Name; }
 
+  /// \name Get the owning MC Module.
+  /// @{
+  const MCModule *getParent() const { return ParentModule; }
+        MCModule *getParent()       { return ParentModule; }
+  /// @}
+
   /// \name Access to the function's basic blocks. No ordering is enforced,
   /// except that the first block is the entry block.
   /// @{
diff --git a/llvm/include/llvm/MC/MCModule.h b/llvm/include/llvm/MC/MCModule.h
index a145653af79..c0c242c21a6 100644
--- a/llvm/include/llvm/MC/MCModule.h
+++ b/llvm/include/llvm/MC/MCModule.h
@@ -43,7 +43,9 @@ class MCModule {
   typedef std::vector<MCAtom*> AtomListTy;
   AtomListTy Atoms;
 
+  // For access to map/remap.
   friend class MCAtom;
+
   /// \brief Remap \p Atom to the given range, and update its Begin/End fields.
   /// \param Atom An atom belonging to this module.
   /// An atom should always use this method to update its bounds, because this
@@ -83,6 +85,8 @@ public:
   /// @{
   const MCAtom *findAtomContaining(uint64_t Addr) const;
         MCAtom *findAtomContaining(uint64_t Addr);
+  const MCAtom *findFirstAtomAfter(uint64_t Addr) const;
+        MCAtom *findFirstAtomAfter(uint64_t Addr);
 
   typedef AtomListTy::const_iterator const_atom_iterator;
   typedef AtomListTy::      iterator       atom_iterator;
diff --git a/llvm/lib/MC/MCFunction.cpp b/llvm/lib/MC/MCFunction.cpp
index 5011d5fd6b8..300ab5b1a0e 100644
--- a/llvm/lib/MC/MCFunction.cpp
+++ b/llvm/lib/MC/MCFunction.cpp
@@ -26,8 +26,9 @@ MCFunction::~MCFunction() {
 }
 
 MCBasicBlock &MCFunction::createBlock(const MCTextAtom &TA) {
-  Blocks.push_back(new MCBasicBlock(TA, this));
-  return *Blocks.back();
+  MCBasicBlock *MCBB = new MCBasicBlock(TA, this);
+  Blocks.push_back(MCBB);
+  return *MCBB;
 }
 
 const MCBasicBlock *MCFunction::find(uint64_t StartAddr) const {
diff --git a/llvm/lib/MC/MCModule.cpp b/llvm/lib/MC/MCModule.cpp
index 9a9d90e5b6a..bdd5cc6b1cd 100644
--- a/llvm/lib/MC/MCModule.cpp
+++ b/llvm/lib/MC/MCModule.cpp
@@ -18,6 +18,10 @@ static bool AtomComp(const MCAtom *L, uint64_t Addr) {
   return L->getEndAddr() < Addr;
 }
 
+static bool AtomCompInv(uint64_t Addr, const MCAtom *R) {
+  return Addr < R->getEndAddr();
+}
+
 void MCModule::map(MCAtom *NewAtom) {
   uint64_t Begin = NewAtom->Begin;
 
@@ -77,13 +81,23 @@ const MCAtom *MCModule::findAtomContaining(uint64_t Addr) const {
 }
 
 MCAtom *MCModule::findAtomContaining(uint64_t Addr) {
-  AtomListTy::iterator I = std::lower_bound(atom_begin(), atom_end(),
-                                            Addr, AtomComp);
-  if (I != atom_end() && (*I)->getBeginAddr() <= Addr)
+  return const_cast<MCAtom*>(
+    const_cast<const MCModule *>(this)->findAtomContaining(Addr));
+}
+
+const MCAtom *MCModule::findFirstAtomAfter(uint64_t Addr) const {
+  AtomListTy::const_iterator I = std::upper_bound(atom_begin(), atom_end(),
+                                                  Addr, AtomCompInv);
+  if (I != atom_end())
     return *I;
   return 0;
 }
 
+MCAtom *MCModule::findFirstAtomAfter(uint64_t Addr) {
+  return const_cast<MCAtom*>(
+    const_cast<const MCModule *>(this)->findFirstAtomAfter(Addr));
+}
+
 MCFunction *MCModule::createFunction(StringRef Name) {
   Functions.push_back(new MCFunction(Name, this));
   return Functions.back();