1 files changed, 24 insertions, 14 deletions

diff --git a/tools/llvm-objdump/MCFunction.h b/tools/llvm-objdump/MCFunction.h
index 023ca39183..4677d91bb7 100644
--- a/tools/llvm-objdump/MCFunction.h
+++ b/tools/llvm-objdump/MCFunction.h
@@ -12,8 +12,11 @@
 //
 //===----------------------------------------------------------------------===//
 
+#ifndef LLVM_OBJECTDUMP_MCFUNCTION_H
+#define LLVM_OBJECTDUMP_MCFUNCTION_H
+
 #include "llvm/ADT/ArrayRef.h"
-#include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/ADT/DenseSet.h"
 #include "llvm/MC/MCInst.h"
 #include <map>
 
@@ -31,15 +34,20 @@ struct MCDecodedInst {
   uint64_t Size;
   MCInst Inst;
 
+  MCDecodedInst() {}
   MCDecodedInst(uint64_t Address, uint64_t Size, MCInst Inst)
     : Address(Address), Size(Size), Inst(Inst) {}
+
+  bool operator<(const MCDecodedInst &RHS) const {
+    return Address < RHS.Address;
+  }
 };
 
 /// MCBasicBlock - Consists of multiple MCDecodedInsts and a list of successing
 /// MCBasicBlocks.
 class MCBasicBlock {
-  SmallVector<MCDecodedInst, 8> Insts;
-  typedef SmallPtrSet<MCBasicBlock*, 8> SetTy;
+  std::vector<MCDecodedInst> Insts;
+  typedef DenseSet<uint64_t> SetTy;
   SetTy Succs;
 public:
   ArrayRef<MCDecodedInst> getInsts() const { return Insts; }
@@ -48,10 +56,14 @@ public:
   succ_iterator succ_begin() const { return Succs.begin(); }
   succ_iterator succ_end() const { return Succs.end(); }
 
-  bool contains(MCBasicBlock *BB) const { return Succs.count(BB); }
+  bool contains(uint64_t Addr) const { return Succs.count(Addr); }
 
   void addInst(const MCDecodedInst &Inst) { Insts.push_back(Inst); }
-  void addSucc(MCBasicBlock *BB) { Succs.insert(BB); }
+  void addSucc(uint64_t Addr) { Succs.insert(Addr); }
+
+  bool operator<(const MCBasicBlock &RHS) const {
+    return Insts.size() < RHS.Insts.size();
+  }
 };
 
 /// MCFunction - Represents a named function in machine code, containing
@@ -59,7 +71,7 @@ public:
 class MCFunction {
   const StringRef Name;
   // Keep BBs sorted by address.
-  typedef std::map<uint64_t, MCBasicBlock> MapTy;
+  typedef std::vector<std::pair<uint64_t, MCBasicBlock> > MapTy;
   MapTy Blocks;
 public:
   MCFunction(StringRef Name) : Name(Name) {}
@@ -68,7 +80,8 @@ public:
   static MCFunction
   createFunctionFromMC(StringRef Name, const MCDisassembler *DisAsm,
                        const MemoryObject &Region, uint64_t Start, uint64_t End,
-                       const MCInstrAnalysis *Ana, raw_ostream &DebugOut);
+                       const MCInstrAnalysis *Ana, raw_ostream &DebugOut,
+                       SmallVectorImpl<uint64_t> &Calls);
 
   typedef MapTy::iterator iterator;
   iterator begin() { return Blocks.begin(); }
@@ -77,14 +90,11 @@ public:
   StringRef getName() const { return Name; }
 
   MCBasicBlock &addBlock(uint64_t Address, const MCBasicBlock &BB) {
-    assert(!Blocks.count(Address) && "Already a BB at address.");
-    return Blocks[Address] = BB;
-  }
-
-  MCBasicBlock &getBlockAtAddress(uint64_t Address) {
-    assert(Blocks.count(Address) && "No BB at address.");
-    return Blocks[Address];
+    Blocks.push_back(std::make_pair(Address, BB));
+    return Blocks.back().second;
   }
 };
 
 }
+
+#endif