Add support for DebugInfoRef during DWARF transform (#853)

crates/debug/src/transform/attr.rs

@@ -1,13 +1,10 @@
 use super::address_transform::AddressTransform;
 use super::expression::{compile_expression, CompiledExpression, FunctionFrameInfo};
 use super::range_info_builder::RangeInfoBuilder;
-use super::unit::PendingDieRef;
+use super::refs::{PendingDebugInfoRefs, PendingUnitRefs};
 use super::{DebugInputContext, Reader, TransformError};
 use anyhow::Error;
-use gimli::{
-    write, AttributeValue, DebugLineOffset, DebugStr, DebuggingInformationEntry, UnitOffset,
-};
-use std::collections::HashMap;
+use gimli::{write, AttributeValue, DebugLineOffset, DebugStr, DebuggingInformationEntry};
 
 pub(crate) enum FileAttributeContext<'a> {
     Root(Option<DebugLineOffset>),
@@ -41,8 +38,8 @@ pub(crate) fn clone_die_attributes<'a, R>(
     scope_ranges: Option<&Vec<(u64, u64)>>,
     cu_low_pc: u64,
     out_strings: &mut write::StringTable,
-    die_ref_map: &HashMap<UnitOffset, write::UnitEntryId>,
-    pending_die_refs: &mut Vec<PendingDieRef>,
+    pending_die_refs: &mut PendingUnitRefs,
+    pending_di_refs: &mut PendingDebugInfoRefs,
     file_context: FileAttributeContext<'a>,
 ) -> Result<(), Error>
 where
@@ -243,17 +240,14 @@ where
             AttributeValue::CallingConvention(e) => write::AttributeValue::CallingConvention(e),
             AttributeValue::Inline(e) => write::AttributeValue::Inline(e),
             AttributeValue::Ordering(e) => write::AttributeValue::Ordering(e),
-            AttributeValue::UnitRef(ref offset) => {
-                if let Some(unit_id) = die_ref_map.get(offset) {
-                    write::AttributeValue::ThisUnitEntryRef(*unit_id)
-                } else {
-                    pending_die_refs.push((current_scope_id, attr.name(), *offset));
-                    continue;
-                }
+            AttributeValue::UnitRef(offset) => {
+                pending_die_refs.insert(current_scope_id, attr.name(), offset);
+                continue;
+            }
+            AttributeValue::DebugInfoRef(offset) => {
+                pending_di_refs.insert(current_scope_id, attr.name(), offset);
+                continue;
             }
-            // AttributeValue::DebugInfoRef(_) => {
-            //     continue;
-            // }
             _ => panic!(), //write::AttributeValue::StringRef(out_strings.add("_")),
         };
         let current_scope = out_unit.get_mut(current_scope_id);

crates/debug/src/transform/mod.rs

@@ -1,3 +1,6 @@
+use self::refs::DebugInfoRefsMap;
+use self::simulate::generate_simulated_dwarf;
+use self::unit::clone_unit;
 use crate::gc::build_dependencies;
 use crate::DebugInfoData;
 use anyhow::Error;
@@ -5,10 +8,8 @@ use gimli::{
     write, DebugAddr, DebugAddrBase, DebugLine, DebugStr, LocationLists, RangeLists,
     UnitSectionOffset,
 };
-use simulate::generate_simulated_dwarf;
 use std::collections::HashSet;
 use thiserror::Error;
-use unit::clone_unit;
 use wasmtime_environ::isa::TargetFrontendConfig;
 use wasmtime_environ::{ModuleAddressMap, ModuleVmctxInfo, ValueLabelsRanges};
 
@@ -19,6 +20,7 @@ mod attr;
 mod expression;
 mod line_program;
 mod range_info_builder;
+mod refs;
 mod simulate;
 mod unit;
 mod utils;
@@ -76,12 +78,14 @@ pub fn transform_dwarf(
     let mut out_units = write::UnitTable::default();
 
     let out_line_strings = write::LineStringTable::default();
+    let mut pending_di_refs = Vec::new();
+    let mut di_ref_map = DebugInfoRefsMap::new();
 
     let mut translated = HashSet::new();
     let mut iter = di.dwarf.debug_info.units();
-    while let Some(unit) = iter.next().unwrap_or(None) {
-        let unit = di.dwarf.unit(unit)?;
-        clone_unit(
+    while let Some(header) = iter.next().unwrap_or(None) {
+        let unit = di.dwarf.unit(header)?;
+        if let Some((id, ref_map, pending_refs)) = clone_unit(
             unit,
             &context,
             &addr_tr,
@@ -91,8 +95,12 @@ pub fn transform_dwarf(
             &mut out_units,
             &mut out_strings,
             &mut translated,
-        )?;
+        )? {
+            di_ref_map.insert(&header, id, ref_map);
+            pending_di_refs.push((id, pending_refs));
+        }
     }
+    di_ref_map.patch(pending_di_refs.into_iter(), &mut out_units);
 
     generate_simulated_dwarf(
         &addr_tr,

crates/debug/src/transform/refs.rs

@@ -0,0 +1,111 @@
+//! Helper utils for tracking and patching intra unit or section references.
+
+use gimli::write;
+use gimli::{CompilationUnitHeader, DebugInfoOffset, Reader, UnitOffset};
+use std::collections::HashMap;
+
+/// Stores compiled unit references: UnitEntryId+DwAt denotes a patch location
+/// and UnitOffset is a location in original DWARF.
+pub struct PendingUnitRefs {
+    refs: Vec<(write::UnitEntryId, gimli::DwAt, UnitOffset)>,
+}
+
+impl PendingUnitRefs {
+    pub fn new() -> Self {
+        Self { refs: Vec::new() }
+    }
+    pub fn insert(&mut self, entry_id: write::UnitEntryId, attr: gimli::DwAt, offset: UnitOffset) {
+        self.refs.push((entry_id, attr, offset));
+    }
+}
+
+/// Stores .debug_info references: UnitEntryId+DwAt denotes a patch location
+/// and DebugInfoOffset is a location in original DWARF.
+pub struct PendingDebugInfoRefs {
+    refs: Vec<(write::UnitEntryId, gimli::DwAt, DebugInfoOffset)>,
+}
+
+impl PendingDebugInfoRefs {
+    pub fn new() -> Self {
+        Self { refs: Vec::new() }
+    }
+    pub fn insert(
+        &mut self,
+        entry_id: write::UnitEntryId,
+        attr: gimli::DwAt,
+        offset: DebugInfoOffset,
+    ) {
+        self.refs.push((entry_id, attr, offset));
+    }
+}
+
+/// Stores map between read and written references of DWARF entries of
+/// a compiled unit.
+pub struct UnitRefsMap {
+    map: HashMap<UnitOffset, write::UnitEntryId>,
+}
+
+impl UnitRefsMap {
+    pub fn new() -> Self {
+        Self {
+            map: HashMap::new(),
+        }
+    }
+    pub fn insert(&mut self, offset: UnitOffset, entry_id: write::UnitEntryId) {
+        self.map.insert(offset, entry_id);
+    }
+    pub fn patch(&self, refs: PendingUnitRefs, comp_unit: &mut write::Unit) {
+        for (die_id, attr_name, offset) in refs.refs {
+            let die = comp_unit.get_mut(die_id);
+            if let Some(unit_id) = self.map.get(&offset) {
+                die.set(attr_name, write::AttributeValue::ThisUnitEntryRef(*unit_id));
+            }
+        }
+    }
+}
+
+/// Stores map between read and written references of DWARF entries of
+/// the entire .debug_info.
+pub struct DebugInfoRefsMap {
+    map: HashMap<DebugInfoOffset, (write::UnitId, write::UnitEntryId)>,
+}
+
+impl DebugInfoRefsMap {
+    pub fn new() -> Self {
+        Self {
+            map: HashMap::new(),
+        }
+    }
+    pub fn insert<R>(
+        &mut self,
+        unit: &CompilationUnitHeader<R>,
+        unit_id: write::UnitId,
+        unit_map: UnitRefsMap,
+    ) where
+        R: Reader<Offset = usize>,
+    {
+        self.map
+            .extend(unit_map.map.into_iter().map(|(off, entry_id)| {
+                let off = off.to_debug_info_offset(unit);
+                (off, (unit_id, entry_id))
+            }));
+    }
+    pub fn patch(
+        &self,
+        refs: impl Iterator<Item = (write::UnitId, PendingDebugInfoRefs)>,
+        units: &mut write::UnitTable,
+    ) {
+        for (id, refs) in refs {
+            let unit = units.get_mut(id);
+            for (die_id, attr_name, offset) in refs.refs {
+                let die = unit.get_mut(die_id);
+                if let Some((id, entry_id)) = self.map.get(&offset) {
+                    die.set(
+                        attr_name,
+                        write::AttributeValue::AnyUnitEntryRef((*id, *entry_id)),
+                    );
+                }
+            }
+        }
+    }
+}

crates/debug/src/transform/unit.rs

@@ -3,17 +3,16 @@ use super::attr::{clone_die_attributes, FileAttributeContext};
 use super::expression::compile_expression;
 use super::line_program::clone_line_program;
 use super::range_info_builder::RangeInfoBuilder;
+use super::refs::{PendingDebugInfoRefs, PendingUnitRefs, UnitRefsMap};
 use super::utils::{add_internal_types, append_vmctx_info, get_function_frame_info};
 use super::{DebugInputContext, Reader, TransformError};
 use anyhow::Error;
 use gimli::write;
-use gimli::{AttributeValue, DebuggingInformationEntry, Unit, UnitOffset};
-use std::collections::{HashMap, HashSet};
+use gimli::{AttributeValue, DebuggingInformationEntry, Unit};
+use std::collections::HashSet;
 use wasmtime_environ::entity::EntityRef;
 use wasmtime_environ::{ModuleVmctxInfo, ValueLabelsRanges};
 
-pub(crate) type PendingDieRef = (write::UnitEntryId, gimli::DwAt, UnitOffset);
-
 struct InheritedAttr<T> {
     stack: Vec<(usize, T)>,
 }
@@ -90,7 +89,7 @@ fn replace_pointer_type<R>(
     unit: &Unit<R, R::Offset>,
     context: &DebugInputContext<R>,
     out_strings: &mut write::StringTable,
-    pending_die_refs: &mut Vec<(write::UnitEntryId, gimli::DwAt, UnitOffset)>,
+    pending_die_refs: &mut PendingUnitRefs,
 ) -> Result<write::UnitEntryId, Error>
 where
     R: Reader,
@@ -119,7 +118,7 @@ where
         write::AttributeValue::ThisUnitEntryRef(wp_die_id),
     );
     if let Some(AttributeValue::UnitRef(ref offset)) = entry.attr_value(gimli::DW_AT_type)? {
-        pending_die_refs.push((p_die_id, gimli::DW_AT_type, *offset))
+        pending_die_refs.insert(p_die_id, gimli::DW_AT_type, *offset);
     }
 
     let m_die_id = comp_unit.add(die_id, gimli::DW_TAG_member);
@@ -149,17 +148,18 @@ pub(crate) fn clone_unit<'a, R>(
     out_units: &mut write::UnitTable,
     out_strings: &mut write::StringTable,
     translated: &mut HashSet<u32>,
-) -> Result<(), Error>
+) -> Result<Option<(write::UnitId, UnitRefsMap, PendingDebugInfoRefs)>, Error>
 where
     R: Reader,
 {
-    let mut die_ref_map = HashMap::new();
-    let mut pending_die_refs = Vec::new();
+    let mut die_ref_map = UnitRefsMap::new();
+    let mut pending_die_refs = PendingUnitRefs::new();
+    let mut pending_di_refs = PendingDebugInfoRefs::new();
     let mut stack = Vec::new();
 
     // Iterate over all of this compilation unit's entries.
     let mut entries = unit.entries();
-    let (mut comp_unit, file_map, cu_low_pc, wp_die_id, vmctx_die_id) =
+    let (mut comp_unit, unit_id, file_map, cu_low_pc, wp_die_id, vmctx_die_id) =
         if let Some((depth_delta, entry)) = entries.next_dfs()? {
             assert_eq!(depth_delta, 0);
             let (out_line_program, debug_line_offset, file_map) = clone_line_program(
@@ -200,8 +200,8 @@ where
                     None,
                     cu_low_pc,
                     out_strings,
-                    &die_ref_map,
                     &mut pending_die_refs,
+                    &mut pending_di_refs,
                     FileAttributeContext::Root(Some(debug_line_offset)),
                 )?;
 
@@ -209,12 +209,19 @@ where
                     add_internal_types(comp_unit, root_id, out_strings, module_info);
 
                 stack.push(root_id);
-                (comp_unit, file_map, cu_low_pc, wp_die_id, vmctx_die_id)
+                (
+                    comp_unit,
+                    unit_id,
+                    file_map,
+                    cu_low_pc,
+                    wp_die_id,
+                    vmctx_die_id,
+                )
             } else {
                 return Err(TransformError("Unexpected unit header").into());
             }
         } else {
-            return Ok(()); // empty
+            return Ok(None); // empty
         };
     let mut skip_at_depth = None;
     let mut current_frame_base = InheritedAttr::new();
@@ -333,8 +340,8 @@ where
             current_scope_ranges.top(),
             cu_low_pc,
             out_strings,
-            &die_ref_map,
             &mut pending_die_refs,
+            &mut pending_di_refs,
             FileAttributeContext::Children(&file_map, current_frame_base.top()),
         )?;
 
@@ -350,13 +357,6 @@ where
             )?;
         }
     }
-    for (die_id, attr_name, offset) in pending_die_refs {
-        let die = comp_unit.get_mut(die_id);
-        if let Some(unit_id) = die_ref_map.get(&offset) {
-            die.set(attr_name, write::AttributeValue::ThisUnitEntryRef(*unit_id));
-        } else {
-            // TODO check why loosing DIEs
-        }
-    }
-    Ok(())
+    die_ref_map.patch(pending_die_refs, comp_unit);
+    Ok(Some((unit_id, die_ref_map, pending_di_refs)))
 }