DWARF Import : Introduce fastpath for compilation unit handling

4 files changed, 205 insertions, 58 deletions

diff --git a/rust/examples/dwarf/dwarf_import/src/dwarfdebuginfo.rs b/rust/examples/dwarf/dwarf_import/src/dwarfdebuginfo.rs
index 7872d8e2..d4a2794d 100644
--- a/rust/examples/dwarf/dwarf_import/src/dwarfdebuginfo.rs
+++ b/rust/examples/dwarf/dwarf_import/src/dwarfdebuginfo.rs
@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use crate::helpers::{get_uid, resolve_specification};
+use crate::helpers::{get_uid, resolve_specification, DieReference};
 
 use binaryninja::{
     binaryview::{BinaryView, BinaryViewBase},
@@ -214,9 +214,19 @@ impl DebugInfoBuilder {
         unit: &Unit<R>,
         entry: &DebuggingInformationEntry<R>,
     ) -> Option<CString> {
-        let (entry_unit, entry_offset) = resolve_specification(dwarf, unit, entry);
-        let entry = entry_unit.entry(entry_offset).unwrap();
-        self.names.get(&get_uid(&entry_unit, &entry)).cloned()
+        match resolve_specification(dwarf, unit, entry) {
+            DieReference::Offset(entry_offset) => self
+                .names
+                .get(&get_uid(unit, &unit.entry(entry_offset).unwrap()))
+                .cloned(),
+            DieReference::UnitAndOffset((entry_unit, entry_offset)) => self
+                .names
+                .get(&get_uid(
+                    &entry_unit,
+                    &entry_unit.entry(entry_offset).unwrap(),
+                ))
+                .cloned(),
+        }
     }
 
     fn commit_types(&self, debug_info: &mut DebugInfo) {
diff --git a/rust/examples/dwarf/dwarf_import/src/helpers.rs b/rust/examples/dwarf/dwarf_import/src/helpers.rs
index 8670e6f4..f3b7dd81 100644
--- a/rust/examples/dwarf/dwarf_import/src/helpers.rs
+++ b/rust/examples/dwarf/dwarf_import/src/helpers.rs
@@ -33,6 +33,11 @@ pub(crate) fn get_uid<R: Reader<Offset = usize>>(
 ////////////////////////////////////
 // DIE attr convenience functions
 
+pub(crate) enum DieReference<R: Reader<Offset = usize>> {
+    Offset(UnitOffset),
+    UnitAndOffset((Unit<R>, UnitOffset)),
+}
+
 fn get_unit_copy<'a, R: Reader<Offset = usize>>(dwarf: &'a Dwarf<R>, unit: &'a Unit<R>) -> Unit<R> {
     let mut iter = dwarf.units();
     while let Ok(Some(header)) = iter.next() {
@@ -45,17 +50,20 @@ fn get_unit_copy<'a, R: Reader<Offset = usize>>(dwarf: &'a Dwarf<R>, unit: &'a U
 
 pub(crate) fn get_attr_die<'a, R: Reader<Offset = usize>>(
     dwarf: &'a Dwarf<R>,
-    unit: &'a Unit<R>,
+    _unit: &'a Unit<R>,
     entry: &'a DebuggingInformationEntry<R>,
     attr: constants::DwAt,
-) -> Option<(Unit<R>, UnitOffset)> {
+) -> Option<DieReference<R>> {
     match entry.attr_value(attr) {
-        Ok(Some(UnitRef(offset))) => Some((get_unit_copy(dwarf, unit), offset)),
+        Ok(Some(UnitRef(offset))) => Some(DieReference::Offset(offset)),
         Ok(Some(DebugInfoRef(offset))) | Ok(Some(DebugInfoRefSup(offset))) => {
             let mut iter = dwarf.units();
             while let Ok(Some(header)) = iter.next() {
                 if let Some(new_offset) = offset.to_unit_offset(&header) {
-                    return Some((dwarf.unit(header).unwrap(), new_offset));
+                    return Some(DieReference::UnitAndOffset((
+                        dwarf.unit(header).unwrap(),
+                        new_offset,
+                    )));
                 }
             }
             unreachable!() //None
@@ -68,17 +76,75 @@ pub(crate) fn resolve_specification<'a, R: Reader<Offset = usize>>(
     dwarf: &'a Dwarf<R>,
     unit: &'a Unit<R>,
     entry: &'a DebuggingInformationEntry<R>,
-) -> (Unit<R>, UnitOffset) {
-    if let Some((entry_unit, entry_offset)) =
-        get_attr_die(dwarf, unit, entry, constants::DW_AT_specification)
+) -> DieReference<R> {
+    if let Some(die_reference) = get_attr_die(dwarf, unit, entry, constants::DW_AT_specification) {
+        match die_reference {
+            DieReference::Offset(entry_offset) => {
+                resolve_specification(dwarf, unit, &unit.entry(entry_offset).unwrap())
+            }
+            DieReference::UnitAndOffset((entry_unit, entry_offset)) => {
+                resolve_specification_slowpath(
+                    dwarf,
+                    &entry_unit,
+                    &entry_unit.entry(entry_offset).unwrap(),
+                )
+            }
+        }
+    } else if let Some(die_reference) =
+        get_attr_die(dwarf, unit, entry, constants::DW_AT_abstract_origin)
     {
-        resolve_specification(dwarf, &entry_unit, &entry_unit.entry(entry_offset).unwrap())
-    } else if let Some((entry_unit, entry_offset)) =
+        match die_reference {
+            DieReference::Offset(entry_offset) => {
+                resolve_specification(dwarf, unit, &unit.entry(entry_offset).unwrap())
+            }
+            DieReference::UnitAndOffset((entry_unit, entry_offset)) => {
+                resolve_specification_slowpath(
+                    dwarf,
+                    &entry_unit,
+                    &entry_unit.entry(entry_offset).unwrap(),
+                )
+            }
+        }
+    } else {
+        DieReference::Offset(entry.offset())
+    }
+}
+
+fn resolve_specification_slowpath<'a, R: Reader<Offset = usize>>(
+    dwarf: &'a Dwarf<R>,
+    unit: &'a Unit<R>,
+    entry: &'a DebuggingInformationEntry<R>,
+) -> DieReference<R> {
+    if let Some(die_reference) = get_attr_die(dwarf, unit, entry, constants::DW_AT_specification) {
+        match die_reference {
+            DieReference::Offset(entry_offset) => {
+                resolve_specification_slowpath(dwarf, unit, &unit.entry(entry_offset).unwrap())
+            }
+            DieReference::UnitAndOffset((entry_unit, entry_offset)) => {
+                resolve_specification_slowpath(
+                    dwarf,
+                    &entry_unit,
+                    &entry_unit.entry(entry_offset).unwrap(),
+                )
+            }
+        }
+    } else if let Some(die_reference) =
         get_attr_die(dwarf, unit, entry, constants::DW_AT_abstract_origin)
     {
-        resolve_specification(dwarf, &entry_unit, &entry_unit.entry(entry_offset).unwrap())
+        match die_reference {
+            DieReference::Offset(entry_offset) => {
+                resolve_specification_slowpath(dwarf, unit, &unit.entry(entry_offset).unwrap())
+            }
+            DieReference::UnitAndOffset((entry_unit, entry_offset)) => {
+                resolve_specification_slowpath(
+                    dwarf,
+                    &entry_unit,
+                    &entry_unit.entry(entry_offset).unwrap(),
+                )
+            }
+        }
     } else {
-        (get_unit_copy(dwarf, unit), entry.offset())
+        DieReference::UnitAndOffset((get_unit_copy(dwarf, unit), entry.offset()))
     }
 }
 
@@ -88,17 +154,36 @@ pub(crate) fn get_name<R: Reader<Offset = usize>>(
     unit: &Unit<R>,
     entry: &DebuggingInformationEntry<R>,
 ) -> Option<CString> {
-    let (entry_unit, entry_offset) = resolve_specification(dwarf, unit, entry);
-    let entry = entry_unit.entry(entry_offset).unwrap();
-
-    if let Ok(Some(attr_val)) = entry.attr_value(constants::DW_AT_name) {
-        if let Ok(attr_string) = dwarf.attr_string(&entry_unit, attr_val) {
-            if let Ok(attr_string) = attr_string.to_string() {
-                return Some(CString::new(attr_string.to_string()).unwrap());
+    match resolve_specification(dwarf, unit, entry) {
+        DieReference::Offset(entry_offset) => {
+            if let Ok(Some(attr_val)) = unit
+                .entry(entry_offset)
+                .unwrap()
+                .attr_value(constants::DW_AT_name)
+            {
+                if let Ok(attr_string) = dwarf.attr_string(unit, attr_val) {
+                    if let Ok(attr_string) = attr_string.to_string() {
+                        return Some(CString::new(attr_string.to_string()).unwrap());
+                    }
+                }
             }
+            None
+        }
+        DieReference::UnitAndOffset((entry_unit, entry_offset)) => {
+            if let Ok(Some(attr_val)) = entry_unit
+                .entry(entry_offset)
+                .unwrap()
+                .attr_value(constants::DW_AT_name)
+            {
+                if let Ok(attr_string) = dwarf.attr_string(&entry_unit, attr_val) {
+                    if let Ok(attr_string) = attr_string.to_string() {
+                        return Some(CString::new(attr_string.to_string()).unwrap());
+                    }
+                }
+            }
+            None
         }
     }
-    None
 }
 
 // Get raw name from DIE, or referenced dependencies
diff --git a/rust/examples/dwarf/dwarf_import/src/lib.rs b/rust/examples/dwarf/dwarf_import/src/lib.rs
index 396aa0cf..cc87e1c6 100644
--- a/rust/examples/dwarf/dwarf_import/src/lib.rs
+++ b/rust/examples/dwarf/dwarf_import/src/lib.rs
@@ -20,7 +20,7 @@ mod types;
 
 use crate::dwarfdebuginfo::DebugInfoBuilder;
 use crate::functions::parse_function_entry;
-use crate::helpers::{get_attr_die, get_name, get_uid};
+use crate::helpers::{get_attr_die, get_name, get_uid, DieReference};
 use crate::types::parse_data_variable;
 
 use binaryninja::{
@@ -42,6 +42,7 @@ use std::ffi::CString;
 fn recover_names<R: Reader<Offset = usize>>(
     dwarf: &Dwarf<R>,
     debug_info_builder: &mut DebugInfoBuilder,
+    progress: &dyn Fn(usize, usize) -> Result<(), ()>,
 ) -> usize {
     let mut total_die_count = 0;
     let mut iter = dwarf.units();
@@ -59,6 +60,11 @@ fn recover_names<R: Reader<Offset = usize>>(
 
         while let Ok(Some((delta_depth, entry))) = entries.next_dfs() {
             total_die_count += 1;
+
+            if (*progress)(0, total_die_count).is_err() {
+                return 0; // Parsing canceled
+            };
+
             depth += delta_depth;
             assert!(depth >= 0);
 
@@ -76,17 +82,27 @@ fn recover_names<R: Reader<Offset = usize>>(
                     ) {
                         if let Some(namespace_qualifier) = get_name(dwarf, unit, entry) {
                             namespace_qualifiers.push((depth, namespace_qualifier));
-                        } else if let Some((entry_unit, entry_offset)) =
+                        } else if let Some(die_reference) =
                             get_attr_die(dwarf, unit, entry, constants::DW_AT_extension)
                         {
-                            let entry = entry_unit.entry(entry_offset).unwrap();
-                            resolve_namespace_name(
-                                dwarf,
-                                &entry_unit,
-                                &entry,
-                                namespace_qualifiers,
-                                depth,
-                            );
+                            match die_reference {
+                                DieReference::Offset(entry_offset) => resolve_namespace_name(
+                                    dwarf,
+                                    unit,
+                                    &unit.entry(entry_offset).unwrap(),
+                                    namespace_qualifiers,
+                                    depth,
+                                ),
+                                DieReference::UnitAndOffset((entry_unit, entry_offset)) => {
+                                    resolve_namespace_name(
+                                        dwarf,
+                                        &entry_unit,
+                                        &entry_unit.entry(entry_offset).unwrap(),
+                                        namespace_qualifiers,
+                                        depth,
+                                    )
+                                }
+                            }
                         } else {
                             namespace_qualifiers
                                 .push((depth, CString::new("anonymous_namespace").unwrap()));
@@ -181,9 +197,7 @@ fn parse_unit<R: Reader<Offset = usize>>(
     // There's a lot of junk we don't care about in DWARF info, so we choose a couple DIEs and mutate state (add functions (which adds the types it uses) and keep track of what namespace we're in)
     while let Ok(Some((_, entry))) = entries.next_dfs() {
         *current_die_number += 1;
-        if *current_die_number % 1000 == 0
-            && (*progress)(*current_die_number, total_die_count).is_err()
-        {
+        if (*progress)(*current_die_number, total_die_count).is_err() {
             return; // Parsing canceled
         }
 
@@ -232,7 +246,10 @@ fn parse_dwarf(
     if (*progress)(0, 1).is_err() {
         return debug_info_builder; // Parsing canceled
     };
-    let total_die_count = recover_names(&dwarf, &mut debug_info_builder);
+    let total_die_count = recover_names(&dwarf, &mut debug_info_builder, &progress);
+    if total_die_count == 0 {
+        return debug_info_builder;
+    }
 
     // Parse all the compilation units
     let mut iter = dwarf.units();
diff --git a/rust/examples/dwarf/dwarf_import/src/types.rs b/rust/examples/dwarf/dwarf_import/src/types.rs
index ca750e19..d8edf5ab 100644
--- a/rust/examples/dwarf/dwarf_import/src/types.rs
+++ b/rust/examples/dwarf/dwarf_import/src/types.rs
@@ -203,27 +203,62 @@ pub(crate) fn get_type<R: Reader<Offset = usize>>(
         return None;
     }
 
-    let entry_type = if let Some((entry_unit, entry_offset)) =
-        get_attr_die(dwarf, unit, entry, constants::DW_AT_type)
-    {
-        // This needs to recurse first (before the early return below) to ensure all sub-types have been parsed
-        let entry = entry_unit.entry(entry_offset).unwrap();
-        get_type(dwarf, &entry_unit, &entry, debug_info_builder)
-    } else if let Some((entry_unit, entry_offset)) =
-        get_attr_die(dwarf, unit, entry, constants::DW_AT_specification)
-    {
-        // This needs to recurse first (before the early return below) to ensure all sub-types have been parsed
-        let entry = entry_unit.entry(entry_offset).unwrap();
-        get_type(dwarf, &entry_unit, &entry, debug_info_builder)
-    } else if let Some((entry_unit, entry_offset)) =
-        get_attr_die(dwarf, unit, entry, constants::DW_AT_abstract_origin)
-    {
-        // This needs to recurse first (before the early return below) to ensure all sub-types have been parsed
-        let entry = entry_unit.entry(entry_offset).unwrap();
-        get_type(dwarf, &entry_unit, &entry, debug_info_builder)
-    } else {
-        None
-    };
+    let entry_type =
+        if let Some(die_reference) = get_attr_die(dwarf, unit, entry, constants::DW_AT_type) {
+            // This needs to recurse first (before the early return below) to ensure all sub-types have been parsed
+            match die_reference {
+                DieReference::Offset(entry_offset) => get_type(
+                    dwarf,
+                    unit,
+                    &unit.entry(entry_offset).unwrap(),
+                    debug_info_builder,
+                ),
+                DieReference::UnitAndOffset((entry_unit, entry_offset)) => get_type(
+                    dwarf,
+                    &entry_unit,
+                    &entry_unit.entry(entry_offset).unwrap(),
+                    debug_info_builder,
+                ),
+            }
+        } else if let Some(die_reference) =
+            get_attr_die(dwarf, unit, entry, constants::DW_AT_specification)
+        {
+            // This needs to recurse first (before the early return below) to ensure all sub-types have been parsed
+            match die_reference {
+                DieReference::Offset(entry_offset) => get_type(
+                    dwarf,
+                    unit,
+                    &unit.entry(entry_offset).unwrap(),
+                    debug_info_builder,
+                ),
+                DieReference::UnitAndOffset((entry_unit, entry_offset)) => get_type(
+                    dwarf,
+                    &entry_unit,
+                    &entry_unit.entry(entry_offset).unwrap(),
+                    debug_info_builder,
+                ),
+            }
+        } else if let Some(die_reference) =
+            get_attr_die(dwarf, unit, entry, constants::DW_AT_abstract_origin)
+        {
+            // This needs to recurse first (before the early return below) to ensure all sub-types have been parsed
+            match die_reference {
+                DieReference::Offset(entry_offset) => get_type(
+                    dwarf,
+                    unit,
+                    &unit.entry(entry_offset).unwrap(),
+                    debug_info_builder,
+                ),
+                DieReference::UnitAndOffset((entry_unit, entry_offset)) => get_type(
+                    dwarf,
+                    &entry_unit,
+                    &entry_unit.entry(entry_offset).unwrap(),
+                    debug_info_builder,
+                ),
+            }
+        } else {
+            None
+        };
 
     // If this node (and thus all its referenced nodes) has already been processed, just return the offset
     // This check is not redundant because this type might have been processes in the recursive calls above