Linear scanning in CompileUnit.iter_DIEs()

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 ([email protected])
+# 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()