Enhance SymbolFileDWARF::ParseDeclsForContext performance

[guiandrade]

Summary

This implements DWARFASTParserClang::EnsureAllDIEsInDeclContextHaveBeenParsed so as to provide a faster way to ensure all DIEs linked to a certain declaration context have been parsed.

Currently, we rely on SymbolFileDWARF::ParseDeclsForContext calling DWARFASTParserClang::GetDIEForDeclContext, and only then DWARFASTParserClang::GetDeclForUIDFromDWARF. This change shortcuts that logic and removes redundant calls to DWARFASTParserClang:: GetClangDeclForDIE by deleting DIEs from the m_decl_ctx_to_die map once they have been parsed.

[guiandrade]

Hey guys,

This change is more for me to get to know what you think. I've noticed that the GetDeclForUIDFromDWARF() calls inside SymbolFileDWARF::ParseDeclsForContext (https://github.com/llvm/llvm-project/blob/ef82098a800178a1f973abb8af86eaa690a29734/lldb/source/Plugins/SymbolFile/DWARF/SymbolFileDWARF.cpp#L1216) often end up having no side effect besides generating a llvm::DenseMap::find invocation (https://github.com/llvm/llvm-project/blob/f07b4aff06d83c6ad25d95f456fbc12b2d2a0a0c/lldb/source/Plugins/SymbolFile/DWARF/DWARFASTParserClang.cpp#L3322), especially when we evaluate multiple expressions inside the same scope. So I was wondering if there's a way to improve that. Do you guys think we could do something like what this change proposes, or am I missing something important? Thanks!

[labath]

I am not very familiar with this code, but I don't see a reason why what you're doing could not work. However, it looks like the implementation of it could be done in a better way. For instance, this function seems to be the only caller of GetDIEForDeclContext. So, we could replace it with something like ast_parser->ParseAllDiesForContext(decl_ctx). This would avoid materializing the std::vector, and the ast parser could store the list of processed dies in a more intelligent (memory friendly) fashion.

The thing that's not clear to me is under what circumstances (if any) can a new DIE appear in a decl_context after ParseDeclsForContext has been called. @clayborg, do you have any idea?

[clayborg]

In D67022#1653248, @guiandrade wrote:

Hey guys,

This change is more for me to get to know what you think. I've noticed that the GetDeclForUIDFromDWARF() calls inside SymbolFileDWARF::ParseDeclsForContext (https://github.com/llvm/llvm-project/blob/ef82098a800178a1f973abb8af86eaa690a29734/lldb/source/Plugins/SymbolFile/DWARF/SymbolFileDWARF.cpp#L1216) often end up having no side effect besides generating a llvm::DenseMap::find invocation (https://github.com/llvm/llvm-project/blob/f07b4aff06d83c6ad25d95f456fbc12b2d2a0a0c/lldb/source/Plugins/SymbolFile/DWARF/DWARFASTParserClang.cpp#L3322), especially when we evaluate multiple expressions inside the same scope. So I was wondering if there's a way to improve that. Do you guys think we could do something like what this change proposes, or am I missing something important? Thanks!

So this function gets called when clang wants to make sure that all decls inside of a function, class, struct, union are parsed. Parsed means DWARF has been converted into clang AST in the TypeSystem for the current DWARF file. This means if the DeclContext that is being used is something the clang expression parser wants to use and search, we can make sure all decls inside of a DeclContext are parsed so they will be visible. It will cause all DIEs to be parsed and all clang AST counterparts to be created for each DIE that gets returned. I really am not following what this patch does. The fact that you see:

clang::Decl *DWARFASTParserClang::GetClangDeclForDIE(const DWARFDIE &die)

seem like it is always just grabbing an entry from the cache (https://github.com/llvm/llvm-project/blob/f07b4aff06d83c6ad25d95f456fbc12b2d2a0a0c/lldb/source/Plugins/SymbolFile/DWARF/DWARFASTParserClang.cpp#L3322) is only because this DWARFDIE has converted DWARF into a clang AST type already. The first time a decl gets parsed, it will fall through to the code below which converts the DWARF to clang AST types.

So, original code is doing what is intended. More comments below about how we can make this more efficient.

In D67022#1654754, @labath wrote:

I am not very familiar with this code, but I don't see a reason why what you're doing could not work. However, it looks like the implementation of it could be done in a better way. For instance, this function seems to be the only caller of GetDIEForDeclContext. So, we could replace it with something like ast_parser->ParseAllDiesForContext(decl_ctx). This would avoid materializing the std::vector, and the ast parser could store the list of processed dies in a more intelligent (memory friendly) fashion.

The thing that's not clear to me is under what circumstances (if any) can a new DIE appear in a decl_context after ParseDeclsForContext has been called. @clayborg, do you have any idea?

No new DIEs should appear in a decl context in the future.

Speaking to efficiency we can do a few things:
1 - keep a cache of CompilerDeclContext's that have already been expanded. This way we can avoid the need to get a list of DIE, just to iterate through them and do nothing.

We can declare a new ivar in SymbolFileDWARF:

std::set<CompilerDeclContext> m_parsed_decls_for_decl_ctx;

And then use this std::set. If the pair that is returned has "second" set to "false", then it has already been parsed, else it will set it and fall through and parse all decls one time:

void SymbolFileDWARF::ParseDeclsForContext(CompilerDeclContext decl_ctx) {
  if (!m_parsed_decls_for_decl_ctx.insert(decl_ctx).second)
    return;
  // Fall through and do what we did before

2 - Change GetDIEForDeclContext to take a std::function or llvm::function that will get called with each DWARFDIE we need. This will avoid making a std::vector which we just iterate through. If we go this route we should rename this to ForEachDeclContextDIE(...).

Change:

std::vector<DWARFDIE> DWARFASTParserClang::GetDIEForDeclContext(lldb_private::CompilerDeclContext decl_context);

to:

void DWARFASTParserClang::ForEachDIEInDeclContext(lldb_private::CompilerDeclContext decl_context, std::function<void(DWARFDIE die)> const &callback);

"callback" will get called with each DIE instead of us making a vector.

lldb/source/Plugins/SymbolFile/DWARF/SymbolFileDWARF.cpp

1210–1211

delete these two lines

1212

Why the reference to 'offset'? Also we shouldn't need to make a key, just use decl_ctx:

uint32_t offset = m_decl_context_to_offset_map[decl_ctx];

Also, m_decl_context_to_offset_map is not populated anywhere. Not sure what this will return as this was never assigned. It will probably just return a default constructed "uint32_t"??

lldb/source/Plugins/SymbolFile/DWARF/SymbolFileDWARF.h

493

Any reason we are just using a lldb_private::CompilerDeclContext? Remove this line?

494

typedef std::map<CompilerDeclContext, uint32_t> DeclContextToOffsetMap;

495

I don't see 'm_decl_context_to_offset_map' being populated anywhere?

[guiandrade]

Thank you so much for the thorough explanation and ideas, @labath and @clayborg.

I forgot to mention that the initial motivation behind this change was to reduce those find calls as I saw they add up to a lot (10^5 calls debugging a Unreal engine sample at certain breakpoints).

lldb/source/Plugins/SymbolFile/DWARF/SymbolFileDWARF.cpp

1212

The idea was to use that reference to populate the map at line 1222. However, I was only using this offset map because I thought DIEs could appear in a decl context. Given that this is not the case, I agree that using a set as you suggested is much better.

[guiandrade]

Renaming the test

lldb/unittests/SymbolFile/DWARF/DWARFASTParserClangTests.cpp

48–49

I think we should not expect any particular order, right?

[labath]

Thanks for the clarification Greg.

Functionally, this patch seems fine, but I am still wondering if we shouldn't make this more efficient. std::set is not the most memory-efficient structure, and so creating a std::set entry just to store one bit seems wasteful, particularly as there is already a container which uses the context as a key (m_decl_ctx_to_die). Could just shove this bit into the m_decl_ctx_to_die map (probably by changing it to something functionally equivalent to map<DeclContext, pair<bool, vector<DWARFDie>>>)? Given that the sole "raison d´être" of this map is to enable the ParseDeclsForContext functionality, I don't think that having that flag be stored there should be awkward. Quite the opposite, it would remove the awkward back-and-forth between the SymbolFileDWARF and the DWARFASTParsers (symbol file calls ForEachDIEInDeclContext, passing it a callback; then ast parser calls the callback; finally the callback immediately re-enters the parser via GetDeclForUIDFromDWARF) -- instead we could just have the SymbolFileDWARF do ast_parser->PleaseParseAllDiesInThisContextIfTheyHaventBeenParsedAlready(ctx) and have everything happen inside the parser.

So, overall, it seems to me that this way, we could improve the code readability while reducing the time, and avoiding a bigger memory increase. In fact, since we now will not be using the list of dies after they have been parsed once, we could even free up the vector<DWARFDie> after the parsing is complete, and thereby reduce the memory footprint as well. That would be a win-win-win scenario. :)

WDYT?

lldb/unittests/SymbolFile/DWARF/DWARFASTParserClangTests.cpp

48–49

Probably not. A more succinct way of expressing that would be EXPECT_THAT(die_list, testing::UnorderedElementsAre(die2, die3))

[clayborg]

In D67022#1656874, @labath wrote:

Thanks for the clarification Greg.

Functionally, this patch seems fine, but I am still wondering if we shouldn't make this more efficient. std::set is not the most memory-efficient structure, and so creating a std::set entry just to store one bit seems wasteful, particularly as there is already a container which uses the context as a key (m_decl_ctx_to_die). Could just shove this bit into the m_decl_ctx_to_die map (probably by changing it to something functionally equivalent to map<DeclContext, pair<bool, vector<DWARFDie>>>)? Given that the sole "raison d´être" of this map is to enable the ParseDeclsForContext functionality, I don't think that having that flag be stored there should be awkward. Quite the opposite, it would remove the awkward back-and-forth between the SymbolFileDWARF and the DWARFASTParsers (symbol file calls ForEachDIEInDeclContext, passing it a callback; then ast parser calls the callback; finally the callback immediately re-enters the parser via GetDeclForUIDFromDWARF) -- instead we could just have the SymbolFileDWARF do ast_parser->PleaseParseAllDiesInThisContextIfTheyHaventBeenParsedAlready(ctx) and have everything happen inside the parser.

So, overall, it seems to me that this way, we could improve the code readability while reducing the time, and avoiding a bigger memory increase. In fact, since we now will not be using the list of dies after they have been parsed once, we could even free up the vector<DWARFDie> after the parsing is complete, and thereby reduce the memory footprint as well. That would be a win-win-win scenario. :)

WDYT?

Great idea, but this idea gave me the idea to use "m_decl_ctx_to_die" as is, and just remove the entry once we have parse all decls. Then we free the memory and don't need a bit. If there is no entry in the m_decl_ctx_to_die map, then ForEachDIEInDeclContext will just not iterate at all?

[labath]

In D67022#1658057, @clayborg wrote:

Great idea, but this idea gave me the idea to use "m_decl_ctx_to_die" as is, and just remove the entry once we have parse all decls. Then we free the memory and don't need a bit. If there is no entry in the m_decl_ctx_to_die map, then ForEachDIEInDeclContext will just not iterate at all?

Yes, that sounds even better!

[guiandrade]

Migrating to EnsureAllDIEsInDeclContextHaveBeenParsed

[guiandrade]

In D67022#1658057, @clayborg wrote:

In D67022#1656874, @labath wrote:

Thanks for the clarification Greg.

Functionally, this patch seems fine, but I am still wondering if we shouldn't make this more efficient. std::set is not the most memory-efficient structure, and so creating a std::set entry just to store one bit seems wasteful, particularly as there is already a container which uses the context as a key (m_decl_ctx_to_die). Could just shove this bit into the m_decl_ctx_to_die map (probably by changing it to something functionally equivalent to map<DeclContext, pair<bool, vector<DWARFDie>>>)? Given that the sole "raison d´être" of this map is to enable the ParseDeclsForContext functionality, I don't think that having that flag be stored there should be awkward. Quite the opposite, it would remove the awkward back-and-forth between the SymbolFileDWARF and the DWARFASTParsers (symbol file calls ForEachDIEInDeclContext, passing it a callback; then ast parser calls the callback; finally the callback immediately re-enters the parser via GetDeclForUIDFromDWARF) -- instead we could just have the SymbolFileDWARF do ast_parser->PleaseParseAllDiesInThisContextIfTheyHaventBeenParsedAlready(ctx) and have everything happen inside the parser.

So, overall, it seems to me that this way, we could improve the code readability while reducing the time, and avoiding a bigger memory increase. In fact, since we now will not be using the list of dies after they have been parsed once, we could even free up the vector<DWARFDie> after the parsing is complete, and thereby reduce the memory footprint as well. That would be a win-win-win scenario. :)

WDYT?

Great idea, but this idea gave me the idea to use "m_decl_ctx_to_die" as is, and just remove the entry once we have parse all decls. Then we free the memory and don't need a bit. If there is no entry in the m_decl_ctx_to_die map, then ForEachDIEInDeclContext will just not iterate at all?

Thanks for the ideas, I think they greatly simplify the code. Let me know what you guys think of my implementation. I think the main issue with it is that it makes it hard to unit test it...

[guiandrade]

lldb/unittests/SymbolFile/DWARF/DWARFASTParserClangTests.cpp

50

Once we agree on how to move forward, the idea is to delete/modify this, not leave it commented out

[labath]

Thanks. The code looks fine to me. Losing the ability to unit test is a bit of a pity, but since test was pretty icky to begin with, I can't say I'm going to miss it too much...
One way we could possibly test this is with a full scale test via the "target modules dump ast" command. So, the idea would be to run to process to some point, evaluate some expressions, run "target modules dump ast" and check the things are in the output (or that they are *not* there).

It's hard to say what exactly you should be checking for, since the assumption is that we are not changing the behavior here, and so the existing tests should cover that in theory, but the "dump ast" command is a new thing, we don't have too many of tests for it, and so any tests you add will be useful, even if they're not specifically targeting the thing you fix in this patch. If you are able to test the thing that the previous patch fixed in this way, than that would be great though. I'm not 100% sure it can be done, but since (IIUC) the bug/problem was about parsing too many things, I would expect that would manifest itself in some additional entries showing up in the parsed ast.

[guiandrade]

In D67022#1659155, @labath wrote:

Thanks. The code looks fine to me. Losing the ability to unit test is a bit of a pity, but since test was pretty icky to begin with, I can't say I'm going to miss it too much...
One way we could possibly test this is with a full scale test via the "target modules dump ast" command. So, the idea would be to run to process to some point, evaluate some expressions, run "target modules dump ast" and check the things are in the output (or that they are *not* there).

It's hard to say what exactly you should be checking for, since the assumption is that we are not changing the behavior here, and so the existing tests should cover that in theory, but the "dump ast" command is a new thing, we don't have too many of tests for it, and so any tests you add will be useful, even if they're not specifically targeting the thing you fix in this patch. If you are able to test the thing that the previous patch fixed in this way, than that would be great though. I'm not 100% sure it can be done, but since (IIUC) the bug/problem was about parsing too many things, I would expect that would manifest itself in some additional entries showing up in the parsed ast.

I'm trying to play with that command to see if anything changes there, but I haven't been able to do so yet. I'm using the following snippet.

// main.cpp
#include "Foo.h"
int main() { while (true) foo(); /*Breakpoint A*/ /*Breakpoint C*/ }

// Foo.h
#pragma once
void foo();

// Foo.cpp
#include "Foo.h"
namespace {
	struct C { void f() {} };
	struct C0 : public C {
		int c00, c01, c02;
	};
};
void foo() {
	C0 c0; // Breakpoint B
	c0.f();
}

Removing this patch and the other one, these are the GetClangDeclForDIE calls we make https://pastebin.com/AcKGCBz7.
After applying this patch, we cut that down to https://pastebin.com/BjwmbmE0.
Nevertheless, in either case our AST ends up looking the same at breakpoint 'C'.

TranslationUnitDecl 0x55fe81dff238 <<invalid sloc>> <invalid sloc>
|-FunctionDecl 0x55fe81dffb28 <<invalid sloc>> <invalid sloc> main 'int ()' extern
|-FunctionDecl 0x55fe81dffbf8 <<invalid sloc>> <invalid sloc> foo 'void ()' extern
|-NamespaceDecl 0x55fe81dffc98 <<invalid sloc>> <invalid sloc>
| `-CXXRecordDecl 0x55fe81dffd20 <<invalid sloc>> <invalid sloc> <undeserialized declarations> struct C0
`-UsingDirectiveDecl 0x55fe81dffec0 <<invalid sloc>> <invalid sloc> Namespace 0x55fe81dffc98 ''

I'm not sure if there's a way to have the AST reflect the extra calls (but I also know very little about how that part of the code works). Can you guys think of a piece of code that could do that?

Thanks!

[clayborg]

This looks like a good approach to me. Pavel?

[guiandrade]

In D67022#1660901, @clayborg wrote:

This looks like a good approach to me. Pavel?

Even without tests? Just to clarify my last comment, I could not find a way to use the command "target modules dump as" to ensure we won't have a regression in the future. Apparently, with or without the extra calls, that command's output looks the same.

[clayborg]

In D67022#1660945, @guiandrade wrote:

In D67022#1660901, @clayborg wrote:

This looks like a good approach to me. Pavel?

Even without tests? Just to clarify my last comment, I could not find a way to use the command "target modules dump as" to ensure we won't have a regression in the future. Apparently, with or without the extra calls, that command's output looks the same.

I would prefer to have tests. I was just saying that the code looked good, and it is now time to proceed with finding a way to test.

[guiandrade]

In D67022#1660979, @clayborg wrote:

In D67022#1660945, @guiandrade wrote:

In D67022#1660901, @clayborg wrote:

This looks like a good approach to me. Pavel?

Even without tests? Just to clarify my last comment, I could not find a way to use the command "target modules dump as" to ensure we won't have a regression in the future. Apparently, with or without the extra calls, that command's output looks the same.

I would prefer to have tests. I was just saying that the code looked good, and it is now time to proceed with finding a way to test.

Oh, got it. Thanks for the clarification.

[guiandrade]

Falling back to the previous workaround to test EnsureAllDIEsInDeclContextHaveBeenParsed.

[guiandrade]

Maybe we could use the previous workaround until we find something better?

[clayborg]

LGTM. Test seems about as good as we can do. Pavel, you ok with this now?

[labath]

looks fine to me

[guiandrade]

In D67022#1665800, @labath wrote:

looks fine to me

Thank you, guys! Could you please submit it for me?

[labath]

In D67022#1669298, @guiandrade wrote:

Thank you, guys! Could you please submit it for me?

Landed as r370374. Sorry about the delay.

lldb/source/Plugins/SymbolFile/DWARF/DWARFASTParser.h

public:
   GetDeclForUIDFromDWARF(const DWARFDIE &die) = 0;
 
   virtual lldb_private::CompilerDeclContext
   GetDeclContextForUIDFromDWARF(const DWARFDIE &die) = 0;
 
   virtual lldb_private::CompilerDeclContext
   GetDeclContextContainingUIDFromDWARF(const DWARFDIE &die) = 0;
 
-  virtual std::vector<DWARFDIE>
-  GetDIEForDeclContext(lldb_private::CompilerDeclContext decl_context) = 0;
+  virtual void ForEachDIEInDeclContext(
+      lldb_private::CompilerDeclContext decl_context,
+      std::function<void(DWARFDIE die)> const &callback) = 0;
 
   static llvm::Optional<lldb_private::SymbolFile::ArrayInfo>
   ParseChildArrayInfo(const DWARFDIE &parent_die,
                       const lldb_private::ExecutionContext *exe_ctx = nullptr);
 };
 
 #endif // SymbolFileDWARF_DWARFASTParser_h_

lldb/source/Plugins/SymbolFile/DWARF/DWARFASTParserClang.h

public:
 
   bool
   CompleteTypeFromDWARF(const DWARFDIE &die, lldb_private::Type *type,
                         lldb_private::CompilerType &compiler_type) override;
 
   lldb_private::CompilerDecl
   GetDeclForUIDFromDWARF(const DWARFDIE &die) override;
 
-  std::vector<DWARFDIE>
-  GetDIEForDeclContext(lldb_private::CompilerDeclContext decl_context) override;
+  void ForEachDIEInDeclContext(
+      lldb_private::CompilerDeclContext decl_context,
+      std::function<void(DWARFDIE die)> const &callback) override;
 
   lldb_private::CompilerDeclContext
   GetDeclContextForUIDFromDWARF(const DWARFDIE &die) override;
 
   lldb_private::CompilerDeclContext
   GetDeclContextContainingUIDFromDWARF(const DWARFDIE &die) override;
 
   lldb_private::ClangASTImporter &GetClangASTImporter();

lldb/source/Plugins/SymbolFile/DWARF/DWARFASTParserClang.cpp

#endif
   default:
     assert(false && "not a forward clang type decl!");
     break;
   }
 
   return false;
 }
 
-std::vector<DWARFDIE> DWARFASTParserClang::GetDIEForDeclContext(
-    lldb_private::CompilerDeclContext decl_context) {
-  std::vector<DWARFDIE> result;
+void DWARFASTParserClang::ForEachDIEInDeclContext(
+    lldb_private::CompilerDeclContext decl_context,
+    std::function<void(DWARFDIE die)> const &callback) {
   auto opaque_decl_ctx =
       (clang::DeclContext *)decl_context.GetOpaqueDeclContext();
   for (auto it = m_decl_ctx_to_die.find(opaque_decl_ctx);
        it != m_decl_ctx_to_die.end() && it->first == opaque_decl_ctx; it++)
-    result.push_back(it->second);
-  return result;
+    callback(it->second);
 }
 
 CompilerDecl DWARFASTParserClang::GetDeclForUIDFromDWARF(const DWARFDIE &die) {
   clang::Decl *clang_decl = GetClangDeclForDIE(die);
   if (clang_decl != nullptr)
     return CompilerDecl(&m_ast, clang_decl);
   return CompilerDecl();
 }

lldb/source/Plugins/SymbolFile/DWARF/SymbolFileDWARF.h

protected:
   UniqueDWARFASTTypeMap m_unique_ast_type_map;
   DIEToTypePtr m_die_to_type;
   DIEToVariableSP m_die_to_variable_sp;
   DIEToClangType m_forward_decl_die_to_clang_type;
   ClangTypeToDIE m_forward_decl_clang_type_to_die;
   llvm::DenseMap<dw_offset_t, lldb_private::FileSpecList>
       m_type_unit_support_files;
   std::vector<uint32_t> m_lldb_cu_to_dwarf_unit;
+
+  std::set<lldb_private::CompilerDeclContext> m_parsed_decls_for_decl_ctx;
 };
 
 #endif // SymbolFileDWARF_SymbolFileDWARF_h_

lldb/source/Plugins/SymbolFile/DWARF/SymbolFileDWARF.cpp

bool SymbolFileDWARF::ClassOrStructIsVirtual(const DWARFDIE &parent_die) {
   }
   return false;
 }
 
 void SymbolFileDWARF::ParseDeclsForContext(CompilerDeclContext decl_ctx) {
   auto *type_system = decl_ctx.GetTypeSystem();
   if (!type_system)
     return;
-  DWARFASTParser *ast_parser = type_system->GetDWARFParser();
-  std::vector<DWARFDIE> decl_ctx_die_list =
-      ast_parser->GetDIEForDeclContext(decl_ctx);
+  if (!m_parsed_decls_for_decl_ctx.insert(decl_ctx).second)
+    return;
 
-  for (DWARFDIE decl_ctx_die : decl_ctx_die_list)
-    for (DWARFDIE decl = decl_ctx_die.GetFirstChild(); decl;
-         decl = decl.GetSibling())
+  DWARFASTParser *ast_parser = type_system->GetDWARFParser();
+  ast_parser->ForEachDIEInDeclContext(decl_ctx, [ast_parser](DWARFDIE die) {
+    for (DWARFDIE decl = die.GetFirstChild(); decl; decl = decl.GetSibling())
       ast_parser->GetDeclForUIDFromDWARF(decl);
+  });
 }
 
 user_id_t SymbolFileDWARF::GetUID(DIERef ref) {
   if (GetDebugMapSymfile())
     return GetID() | ref.die_offset();
 
   return user_id_t(GetDwoNum().getValueOr(0x7fffffff)) << 32 |
          ref.die_offset() |

lldb/unittests/SymbolFile/DWARF/DWARFASTParserClangTests.cpp

public:
   using DWARFASTParserClang::DWARFASTParserClang;
   using DWARFASTParserClang::LinkDeclContextToDIE;
 };
 } // namespace
 
 // If your implementation needs to dereference the dummy pointers we are
 // defining here, causing this test to fail, feel free to delete it.
 TEST(DWARFASTParserClangTests,
-     TestGetDIEForDeclContextReturnsOnlyMatchingEntries) {
+     TestForEachDIEInDeclContextReturnsOnlyMatchingEntries) {
   ClangASTContext ast_ctx;
   DWARFASTParserClangStub ast_parser(ast_ctx);
 
   DWARFUnit *unit = nullptr;
   DWARFDIE die1(unit, (DWARFDebugInfoEntry *)1LL);
   DWARFDIE die2(unit, (DWARFDebugInfoEntry *)2LL);
   DWARFDIE die3(unit, (DWARFDebugInfoEntry *)3LL);
   DWARFDIE die4(unit, (DWARFDebugInfoEntry *)4LL);
   ast_parser.LinkDeclContextToDIE((clang::DeclContext *)1LL, die1);
   ast_parser.LinkDeclContextToDIE((clang::DeclContext *)2LL, die2);
   ast_parser.LinkDeclContextToDIE((clang::DeclContext *)2LL, die3);
   ast_parser.LinkDeclContextToDIE((clang::DeclContext *)3LL, die4);
 
-  auto die_list = ast_parser.GetDIEForDeclContext(
-      CompilerDeclContext(nullptr, (clang::DeclContext *)2LL));
+  std::vector<DWARFDIE> die_list;
+  ast_parser.ForEachDIEInDeclContext(
+      CompilerDeclContext(nullptr, (clang::DeclContext *)2LL),
+      [&die_list](DWARFDIE die) { die_list.push_back(die); });
+
   ASSERT_EQ(2u, die_list.size());
-  ASSERT_EQ(die2, die_list[0]);
-  ASSERT_EQ(die3, die_list[1]);
+  ASSERT_NE(die_list.end(), std::find(die_list.begin(), die_list.end(), die2));
+  ASSERT_NE(die_list.end(), std::find(die_list.begin(), die_list.end(), die3));
 }