1 files changed, 13 insertions, 9 deletions
diff --git a/llvm/lib/Transforms/Utils/AddDiscriminators.cpp b/llvm/lib/Transforms/Utils/AddDiscriminators.cpp
index 45b4fab07ed..a9b5d7206b7 100644
--- a/llvm/lib/Transforms/Utils/AddDiscriminators.cpp
+++ b/llvm/lib/Transforms/Utils/AddDiscriminators.cpp
@@ -168,7 +168,8 @@ static bool addDiscriminators(Function &F) {
   DIBuilder Builder(*M, /*AllowUnresolved*/ false);
 
   typedef std::pair<StringRef, unsigned> Location;
-  typedef DenseMap<const BasicBlock *, Metadata *> BBScopeMap;
+  typedef SmallDenseMap<DIScope *, DILexicalBlockFile *, 1> ScopeMap;
+  typedef DenseMap<const BasicBlock *, ScopeMap> BBScopeMap;
   typedef DenseMap<Location, BBScopeMap> LocationBBMap;
   typedef DenseMap<Location, unsigned> LocationDiscriminatorMap;
   typedef DenseSet<Location> LocationSet;
@@ -186,20 +187,23 @@ static bool addDiscriminators(Function &F) {
       const DILocation *DIL = I.getDebugLoc();
       if (!DIL)
         continue;
+      DIScope *Scope = DIL->getScope();
       Location L = std::make_pair(DIL->getFilename(), DIL->getLine());
       auto &BBMap = LBM[L];
-      auto R = BBMap.insert(std::make_pair(&B, (Metadata *)nullptr));
+      auto R = BBMap.insert({&B, ScopeMap()});
       if (BBMap.size() == 1)
         continue;
       bool InsertSuccess = R.second;
-      Metadata *&NewScope = R.first->second;
-      // If we could insert a different block in the same location, a
+      ScopeMap &Scopes = R.first->second; 
+      // If we could insert more than one block with the same line+file, a
       // discriminator is needed to distinguish both instructions.
-      if (InsertSuccess) {
-        auto *Scope = DIL->getScope();
-        auto *File =
-            Builder.createFile(DIL->getFilename(), Scope->getDirectory());
-        NewScope = Builder.createLexicalBlockFile(Scope, File, ++LDM[L]);
+      auto R1 = Scopes.insert({Scope, nullptr});
+      DILexicalBlockFile *&NewScope = R1.first->second;
+      unsigned Discriminator = InsertSuccess ? ++LDM[L] : LDM[L];
+      if (!NewScope) {
+        auto *File = Builder.createFile(DIL->getFilename(),
+                                        Scope->getDirectory());
+        NewScope = Builder.createLexicalBlockFile(Scope, File, Discriminator);
       }
       I.setDebugLoc(DILocation::get(Ctx, DIL->getLine(), DIL->getColumn(),
                                     NewScope, DIL->getInlinedAt()));