Linear scanning in CompileUnit.iter_DIEs() (#625)

CompileUnit.iter_DIEs() is unnecessarily convoluted in terms of parse logic. Its purpose it to return the DIEs as they appear in the info section, but instead it performs recursive tree traversal via CompileUnit.iter_DIE_children(), which contains the more complicated logic of going over immediate children of a DIE. On a siblingless DIE tree (which is rare in practice but allowed by the format), this will contain extra work of parsing the whole subtree only to throw away the children of children.

Physically, DIEs are stored linearly. Logically, they are a tree. iter_DIEs() is expected to return a flat collection, but it does so by building a tree (out of a flat collection) and then flattening it. This PR changes iter_DIEs() to parse sequentially instead.

As a matter of performance, it shaves 5-7% from the linear scan time using iter_DIEs(). Tested on one rather artificial example; the speedup on real life DWARF might vary. Does not fix #623, but somewhat helps them.

As a downside, it duplicates the caching and the tree building logic already present in CompileUnit.iter_DIE_children().

Author: sevaa

elftools/dwarf/compileunit.py

@@ -133,7 +133,42 @@ def iter_DIEs(self):
         """ Iterate over all the DIEs in the CU, in order of their appearance.
             Note that null DIEs will also be returned.
         """
-        return self._iter_DIE_subtree(self.get_top_DIE())
+        stm = self.dwarfinfo.debug_info_sec.stream
+        pos = self.cu_die_offset
+        end_pos = self.cu_offset + self.size
+
+        die = self.get_top_DIE()
+        yield die
+        pos += die.size
+        parent = die
+        i = 1
+        while pos < end_pos:
+            if i < len(self._diemap) and self._diemap[i] == pos: # DIE already cached
+                die = self._dielist[i]
+            else:
+                die = DIE(self, stm, pos)
+                self._dielist.insert(i, die)
+                self._diemap.insert(i, pos)
+            i += 1
+
+            die._parent = parent
+
+            if die.tag is None:
+                parent._terminator = die
+                parent = parent._parent
+
+            if die.has_children:
+                parent = die
+
+            if die.tag == 'DW_TAG_imported_unit' and self.dwarfinfo.supplementary_dwarfinfo:
+                # Falls back to subtree traversal in the supplemental DWARF. Any way to streamline that too?
+                supp_die = die.get_DIE_from_attribute('DW_AT_import')
+                yield from supp_die.cu._iter_DIE_subtree(supp_die)
+            else:
+                yield die
+
+            pos += die.size
+
 
     def iter_DIE_children(self, die):
         """ Given a DIE, yields either its children, without null DIE list

test/test_tree.py

@@ -0,0 +1,55 @@
+#------------------------------------------------------------------------------
+# elftools tests
+#
+# Seva Alekseyev (sevaa@sprynet.com)
+# This code is in the public domain
+#
+# Making sure the two ways of iterating through DIEs (linear and by tree)
+# return an identically built tree (parents, children, terminators).
+# 
+# TODO: find a siblingless file. None in the corpus so far.
+#------------------------------------------------------------------------------
+
+import unittest
+import os
+from elftools.elf.elffile import ELFFile
+
+class TestTree(unittest.TestCase):
+    def test_tree(self):
+        self.run_test_on('dwarf_llpair.elf', 0, True)
+        self.run_test_on('test_debugsup1.debug', 2, False)
+
+    def run_test_on(self, file_name, cu_index, test_cached):
+        def die_summary(die):
+            return (die.offset, die.tag, die._terminator.offset if die._terminator else None, die.get_parent().offset if die.get_parent() else None)
+        
+        with ELFFile.load_from_path(os.path.join('test', 'testfiles_for_unittests', file_name)) as elf:
+            di = elf.get_dwarf_info()
+            cu = next(c for i,c in enumerate(di.iter_CUs()) if i == cu_index)
+            #_terminator is only set on a DIE *after* that DIE is yielded during enumeration
+            DIEs = [d for d in cu.iter_DIEs()]
+            seq_DIEs = [die_summary(d) for d in DIEs]
+
+            if test_cached:
+                sample_offset = DIEs[len(DIEs) // 2].offset
+                # Offset of a random DIE from the middle for later
+
+            # Deliberately erase the CU/DIE cache to force a repeat parse - this time using an explicit tree traversal
+            di = elf.get_dwarf_info()
+            cu = next(c for i,c in enumerate(di.iter_CUs()) if i == cu_index)
+            DIEs = [d for d in cu._iter_DIE_subtree(cu.get_top_DIE())]
+
+            tree_DIEs = [die_summary(d) for d in DIEs]
+            self.assertSequenceEqual(seq_DIEs, tree_DIEs)
+
+            if test_cached:
+                # Another repeat parse, with a nonblank cache
+                di = elf.get_dwarf_info()
+                cu = next(c for i,c in enumerate(di.iter_CUs()) if i == cu_index)
+                cu.get_DIE_from_refaddr(sample_offset) # Cache this random DIE in the middle
+                DIEs = [d for d in cu.iter_DIEs()]
+                seq_DIEs_x = [die_summary(d) for d in DIEs]
+                self.assertSequenceEqual(seq_DIEs, seq_DIEs_x)
+
+if __name__ == '__main__':
+    unittest.main()