MC CFG: Add a few needed methods, mainly MCModule::findFirstAtomAfter.

While there, do some minor cleanup.

llvm-svn: 188880

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.
   /// @{

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;

llvm/lib/MC/MCFunction.cpp

@@ -26,8 +26,9 @@ MCFunction::~MCFunction() {
 }
 
 MCBasicBlock &MCFunction::createBlock(const MCTextAtom &TA) {
-  Blocks.push_back(new MCBasicBlock(TA, this));
+  MCBasicBlock *MCBB = new MCBasicBlock(TA, this);
-  return *Blocks.back();
+  Blocks.push_back(MCBB);
+  return *MCBB;
 }
 
 const MCBasicBlock *MCFunction::find(uint64_t StartAddr) const {

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(),
+  return const_cast<MCAtom*>(
-                                            Addr, AtomComp);
+    const_cast<const MCModule *>(this)->findAtomContaining(Addr));
-  if (I != atom_end() && (*I)->getBeginAddr() <= 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();