1 files changed, 25 insertions, 9 deletions

diff --git a/llvm/lib/IR/DIBuilder.cpp b/llvm/lib/IR/DIBuilder.cpp
index ae1a1a3..242ca44 100644
--- a/llvm/lib/IR/DIBuilder.cpp
+++ b/llvm/lib/IR/DIBuilder.cpp
@@ -29,7 +29,7 @@ DIBuilder::DIBuilder(Module &m, bool AllowUnresolvedNodes, DICompileUnit *CU)
       AllowUnresolvedNodes(AllowUnresolvedNodes) {
   if (CUNode) {
     if (const auto &ETs = CUNode->getEnumTypes())
-      AllEnumTypes.assign(ETs.begin(), ETs.end());
+      EnumTypes.assign(ETs.begin(), ETs.end());
     if (const auto &RTs = CUNode->getRetainedTypes())
       AllRetainTypes.assign(RTs.begin(), RTs.end());
     if (const auto &GVs = CUNode->getGlobalVariables())
@@ -66,10 +66,10 @@ void DIBuilder::finalize() {
     return;
   }
 
-  if (!AllEnumTypes.empty())
+  if (!EnumTypes.empty())
     CUNode->replaceEnumTypes(MDTuple::get(
-        VMContext, SmallVector<Metadata *, 16>(AllEnumTypes.begin(),
-                                               AllEnumTypes.end())));
+        VMContext, SmallVector<Metadata *, 16>(EnumTypes.begin(),
+                                               EnumTypes.end())));
 
   SmallVector<Metadata *, 16> RetainValues;
   // Declarations and definitions of the same type may be retained. Some
@@ -334,10 +334,13 @@ DIDerivedType *DIBuilder::createTypedef(DIType *Ty, StringRef Name,
                                         DIScope *Context, uint32_t AlignInBits,
                                         DINode::DIFlags Flags,
                                         DINodeArray Annotations) {
-  return DIDerivedType::get(VMContext, dwarf::DW_TAG_typedef, Name, File,
-                            LineNo, getNonCompileUnitScope(Context), Ty, 0,
-                            AlignInBits, 0, std::nullopt, Flags, nullptr,
-                            Annotations);
+  auto *T =
+      DIDerivedType::get(VMContext, dwarf::DW_TAG_typedef, Name, File, LineNo,
+                         getNonCompileUnitScope(Context), Ty, 0, AlignInBits, 0,
+                         std::nullopt, Flags, nullptr, Annotations);
+  if (isa_and_nonnull<DILocalScope>(Context))
+    getSubprogramNodesTrackingVector(Context).emplace_back(T);
+  return T;
 }
 
 DIDerivedType *DIBuilder::createFriend(DIType *Ty, DIType *FriendTy) {
@@ -485,6 +488,8 @@ DICompositeType *DIBuilder::createClassType(
       OffsetInBits, Flags, Elements, 0, VTableHolder,
       cast_or_null<MDTuple>(TemplateParams), UniqueIdentifier);
   trackIfUnresolved(R);
+  if (isa_and_nonnull<DILocalScope>(Context))
+    getSubprogramNodesTrackingVector(Context).emplace_back(R);
   return R;
 }
 
@@ -498,6 +503,8 @@ DICompositeType *DIBuilder::createStructType(
       getNonCompileUnitScope(Context), DerivedFrom, SizeInBits, AlignInBits, 0,
       Flags, Elements, RunTimeLang, VTableHolder, nullptr, UniqueIdentifier);
   trackIfUnresolved(R);
+  if (isa_and_nonnull<DILocalScope>(Context))
+    getSubprogramNodesTrackingVector(Context).emplace_back(R);
   return R;
 }
 
@@ -510,6 +517,8 @@ DICompositeType *DIBuilder::createUnionType(
       getNonCompileUnitScope(Scope), nullptr, SizeInBits, AlignInBits, 0, Flags,
       Elements, RunTimeLang, nullptr, nullptr, UniqueIdentifier);
   trackIfUnresolved(R);
+  if (isa_and_nonnull<DILocalScope>(Scope))
+    getSubprogramNodesTrackingVector(Scope).emplace_back(R);
   return R;
 }
 
@@ -542,7 +551,10 @@ DICompositeType *DIBuilder::createEnumerationType(
       getNonCompileUnitScope(Scope), UnderlyingType, SizeInBits, AlignInBits, 0,
       IsScoped ? DINode::FlagEnumClass : DINode::FlagZero, Elements, 0, nullptr,
       nullptr, UniqueIdentifier);
-  AllEnumTypes.emplace_back(CTy);
+  if (isa_and_nonnull<DILocalScope>(Scope))
+    getSubprogramNodesTrackingVector(Scope).emplace_back(CTy);
+  else
+    EnumTypes.emplace_back(CTy);
   trackIfUnresolved(CTy);
   return CTy;
 }
@@ -640,6 +652,8 @@ DIBuilder::createForwardDecl(unsigned Tag, StringRef Name, DIScope *Scope,
       SizeInBits, AlignInBits, 0, DINode::FlagFwdDecl, nullptr, RuntimeLang,
       nullptr, nullptr, UniqueIdentifier);
   trackIfUnresolved(RetTy);
+  if (isa_and_nonnull<DILocalScope>(Scope))
+    getSubprogramNodesTrackingVector(Scope).emplace_back(RetTy);
   return RetTy;
 }
 
@@ -656,6 +670,8 @@ DICompositeType *DIBuilder::createReplaceableCompositeType(
           nullptr, Annotations)
           .release();
   trackIfUnresolved(RetTy);
+  if (isa_and_nonnull<DILocalScope>(Scope))
+    getSubprogramNodesTrackingVector(Scope).emplace_back(RetTy);
   return RetTy;
 }