Skip to content

Commit

Permalink
Merge pull request #57 from ergrelet/6-reconstruct-dependencies-in-a-…
Browse files Browse the repository at this point in the history 
…compiler-friendly-way

Reconstruct dependencies in a compiler friendly way
  • Loading branch information
@ergrelet
ergrelet authored Mar 22, 2024
2 parents 69e2ab2 + 7faecbe commit 5eb8a50
Show file tree
Hide file tree
Showing 29 changed files with 594 additions and 277 deletions.
221 changes: 170 additions & 51 deletions resym_core/src/pdb_file.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ use pdb::FallibleIterator;
use rayon::iter::{IntoParallelRefIterator, ParallelIterator};

use std::{
collections::BTreeSet,
collections::{BTreeMap, HashMap, HashSet, VecDeque},
io::{self, Read, Seek},
path::PathBuf,
sync::{Arc, RwLock},
Expand Down Expand Up @@ -421,7 +421,7 @@ where

let mut result = String::default();
module_info.symbols()?.for_each(|symbol| {
let mut needed_types = pdb_types::TypeSet::new();
let mut needed_types = pdb_types::NeededTypeSet::new();

match symbol.parse()? {
pdb::SymbolData::UserDefinedType(udt) => {
Expand Down Expand Up @@ -515,55 +515,88 @@ where
print_access_specifiers,
};
let mut type_data = pdb_types::Data::new();
let mut needed_types = pdb_types::TypeSet::new();

// Add the requested type first
type_data.add(
type_finder,
&self.forwarder_to_complete_type,
type_index,
&primitives_flavor,
&mut needed_types,
)?;

// If dependencies aren't needed, we're done
// If dependencies aren't needed, only process the given type index and return
if !reconstruct_dependencies {
let mut needed_types = pdb_types::NeededTypeSet::new();
type_data.add(
type_finder,
&self.forwarder_to_complete_type,
type_index,
&primitives_flavor,
&mut needed_types,
)?;

let mut reconstruction_output = String::new();
type_data.reconstruct(&fmt_configuration, &mut reconstruction_output)?;
type_data.reconstruct(
&fmt_configuration,
&Default::default(),
&mut reconstruction_output,
)?;
return Ok(reconstruction_output);
}

// Add all the needed types iteratively until we're done
let mut dependencies_data = pdb_types::Data::new();
let mut processed_types = BTreeSet::from([type_index]);
let dep_start = Instant::now();
loop {
// Get the first element in needed_types without holding an immutable borrow
let first = needed_types.difference(&processed_types).next().copied();
match first {
None => break,
Some(needed_type_index) => {
// Add the type
dependencies_data.add(
type_finder,
&self.forwarder_to_complete_type,
needed_type_index,
&primitives_flavor,
&mut needed_types,
)?;
let mut type_dependency_map: HashMap<pdb::TypeIndex, Vec<(pdb::TypeIndex, bool)>> =
HashMap::new();
{
let dep_start = Instant::now();

// Add the requested type first
let mut types_to_process: VecDeque<pdb::TypeIndex> = VecDeque::from([type_index]);
let mut processed_type_set = HashSet::new();
// Keep processing new types until there's nothing to process
while let Some(needed_type_index) = types_to_process.pop_front() {
if processed_type_set.contains(&needed_type_index) {
// Already processed, continue
continue;
}

processed_types.insert(needed_type_index);
// Add the type
let mut needed_types = pdb_types::NeededTypeSet::new();
type_data.add(
type_finder,
&self.forwarder_to_complete_type,
needed_type_index,
&primitives_flavor,
&mut needed_types,
)?;

for pair in &needed_types {
// Add forward declaration for types referenced by pointers
if pair.1 {
type_data.add_as_forward_declaration(type_finder, pair.0)?;
}

// Update type dependency map
if let Some(type_dependency) = type_dependency_map.get_mut(&needed_type_index) {
type_dependency.push(*pair);
} else {
type_dependency_map.insert(needed_type_index, vec![*pair]);
}
}
// Update the set of processed types
processed_type_set.insert(needed_type_index);
// Update the queue of type to process
types_to_process.extend(needed_types.into_iter().map(|pair| pair.0));
}

log::debug!(
"Dependencies reconstruction took {} ms",
dep_start.elapsed().as_millis()
);
}
log::debug!(
"Dependencies reconstruction took {} ms",
dep_start.elapsed().as_millis()
);

// Deduce type "depth" from the dependency map
let type_depth_map = compute_type_depth_map(&type_dependency_map, &[type_index]);

let mut reconstruction_output = String::new();
dependencies_data.reconstruct(&fmt_configuration, &mut reconstruction_output)?;
type_data.reconstruct(&fmt_configuration, &mut reconstruction_output)?;
type_data.reconstruct(
&fmt_configuration,
&type_depth_map,
&mut reconstruction_output,
)?;

Ok(reconstruction_output)
}

Expand All @@ -573,45 +606,82 @@ where
print_access_specifiers: bool,
) -> Result<String> {
let mut type_data = pdb_types::Data::new();

let mut type_finder = self.type_information.finder();
let mut processed_types = Vec::new();
let mut type_dependency_map: HashMap<pdb::TypeIndex, Vec<(pdb::TypeIndex, bool)>> =
HashMap::new();
{
let mut type_finder = self.type_information.finder();
// Populate our `TypeFinder`
let mut type_iter = self.type_information.iter();
while (type_iter.next()?).is_some() {
type_finder.update(&type_iter);
}

// Add the requested types
type_iter = self.type_information.iter();
let mut type_iter = self.type_information.iter();
while let Some(item) = type_iter.next()? {
let mut needed_types = pdb_types::TypeSet::new();
let type_index = item.index();
let mut needed_types = pdb_types::NeededTypeSet::new();
// Note(ergelet): try to get the complete type's index here.
// This avoids adding empty "forward reference" type index which
// usually have lower type indices
let complete_type_index = self
.forwarder_to_complete_type
.get(&item.index())
.map(|e| *e)
.unwrap_or_else(|| item.index());
let result = type_data.add(
&type_finder,
&self.forwarder_to_complete_type,
type_index,
complete_type_index,
&primitives_flavor,
&mut needed_types,
);

// Process result
if let Err(err) = result {
// Handle error
match err {
ResymCoreError::PdbError(err) => {
// Ignore this kind of error since some particular PDB features might not be supported.
// This allows the recontruction to go through with the correctly reconstructed types.
log::warn!("Failed to reconstruct type with index {type_index}: {err}")
log::warn!("Failed to reconstruct type with index {complete_type_index}: {err}")
}
_ => return Err(err),
}
} else {
// Handle success
processed_types.push(complete_type_index);
for pair in &needed_types {
// Add forward declaration for types referenced by pointers
if pair.1 {
type_data.add_as_forward_declaration(&type_finder, pair.0)?;
}

// Update type dependency map
if let Some(type_dependency) =
type_dependency_map.get_mut(&complete_type_index)
{
type_dependency.push(*pair);
} else {
type_dependency_map.insert(complete_type_index, vec![*pair]);
}
}
}
}
}

let fmt_configuration = DataFormatConfiguration {
print_access_specifiers,
};
// Deduce type "depth" from the dependency map
let type_depth_map = compute_type_depth_map(&type_dependency_map, &processed_types);

let mut reconstruction_output = String::new();
type_data.reconstruct(&fmt_configuration, &mut reconstruction_output)?;
type_data.reconstruct(
&DataFormatConfiguration {
print_access_specifiers,
},
&type_depth_map,
&mut reconstruction_output,
)?;

Ok(reconstruction_output)
}

Expand Down Expand Up @@ -642,7 +712,7 @@ where
let current_type_index = type_item.index();
// Reconstruct type and retrieve referenced types
let mut type_data = pdb_types::Data::new();
let mut needed_types = pdb_types::TypeSet::new();
let mut needed_types = pdb_types::NeededTypeSet::new();
type_data.add(
&type_finder,
&self.forwarder_to_complete_type,
Expand All @@ -651,7 +721,7 @@ where
&mut needed_types,
)?;

par_iter_if_available!(needed_types).for_each(|t| {
par_iter_if_available!(needed_types).for_each(|(t, _)| {
if let Some(mut xref_list) = xref_map.get_mut(t) {
xref_list.push(current_type_index);
} else {
Expand Down Expand Up @@ -695,3 +765,52 @@ where
}
}
}

fn compute_type_depth_map(
type_dependency_map: &HashMap<pdb::TypeIndex, Vec<(pdb::TypeIndex, bool)>>,
root_types: &[pdb::TypeIndex],
) -> BTreeMap<usize, Vec<pdb::TypeIndex>> {
let depth_start = Instant::now();

let mut type_depth_map: HashMap<pdb::TypeIndex, usize> =
HashMap::from_iter(root_types.iter().map(|elem| (*elem, 0)));
// Perform depth-first search to determine the "depth" of each type
let mut types_to_visit: VecDeque<(usize, pdb::TypeIndex)> =
VecDeque::from_iter(root_types.iter().map(|elem| (0, *elem)));
while let Some((current_type_depth, current_type_index)) = types_to_visit.pop_back() {
if let Some(type_dependencies) = type_dependency_map.get(&current_type_index) {
for (child_type_index, child_is_pointer) in type_dependencies {
// Visit child only if it's directly referenced, to avoid infinite loops
if !child_is_pointer && *child_type_index != current_type_index {
let current_child_depth = current_type_depth + 1;
if let Some(child_type_depth) = type_depth_map.get_mut(child_type_index) {
*child_type_depth = std::cmp::max(*child_type_depth, current_child_depth);
} else {
type_depth_map.insert(*child_type_index, current_child_depth);
}
types_to_visit.push_back((current_child_depth, *child_type_index));
}
}
}
}

// Invert type depth map
let inverted_type_depth_map: BTreeMap<usize, Vec<pdb::TypeIndex>> = type_depth_map
.into_iter()
.fold(BTreeMap::new(), |mut acc, (type_index, type_depth)| {
if let Some(type_indices) = acc.get_mut(&type_depth) {
type_indices.push(type_index);
} else {
acc.insert(type_depth, vec![type_index]);
}

acc
});

log::debug!(
"Depth calculation took {} ms",
depth_start.elapsed().as_millis()
);

inverted_type_depth_map
}
19 changes: 13 additions & 6 deletions resym_core/src/pdb_types/class.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,8 @@ use super::{
primitive_types::PrimitiveReconstructionFlavor,
resolve_complete_type_index, type_bitfield_info, type_name, type_size,
union::Union,
DataFormatConfiguration, Field, Method, Result, ResymCoreError, TypeForwarder, TypeSet,
DataFormatConfiguration, Field, Method, NeededTypeSet, ReconstructibleTypeData, Result,
ResymCoreError, TypeForwarder,
};

#[derive(Debug, Clone, PartialEq, Eq)]
Expand Down Expand Up @@ -52,6 +53,7 @@ pub struct BaseClass {

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Class<'p> {
pub index: pdb::TypeIndex,
pub kind: pdb::ClassKind,
pub name: String,
pub size: u64,
Expand All @@ -71,7 +73,7 @@ impl<'p> Class<'p> {
&mut self,
_: &pdb::TypeFinder<'p>,
_: pdb::TypeIndex,
_: &mut TypeSet,
_: &mut NeededTypeSet,
) -> Result<()> {
// TODO
Ok(())
Expand All @@ -83,7 +85,7 @@ impl<'p> Class<'p> {
type_forwarder: &TypeForwarder,
type_index: pdb::TypeIndex,
primitive_flavor: &PrimitiveReconstructionFlavor,
needed_types: &mut TypeSet,
needed_types: &mut NeededTypeSet,
) -> Result<()> {
// Resolve the complete type's index, if present in the PDB
let complete_type_index = resolve_complete_type_index(type_forwarder, type_index);
Expand Down Expand Up @@ -122,6 +124,7 @@ impl<'p> Class<'p> {
};

let mut class = Class {
index: type_index,
kind: data.kind,
name,
size: data.size,
Expand Down Expand Up @@ -162,6 +165,7 @@ impl<'p> Class<'p> {
};

let mut u = Union {
index: type_index,
name,
size: data.size,
fields: Vec::new(),
Expand Down Expand Up @@ -194,6 +198,7 @@ impl<'p> Class<'p> {
};

let mut e = Enum {
index: type_index,
name,
underlying_type_name: type_name(
type_finder,
Expand Down Expand Up @@ -237,7 +242,7 @@ impl<'p> Class<'p> {
type_forwarder: &TypeForwarder,
field: &pdb::TypeData<'p>,
primitive_flavor: &PrimitiveReconstructionFlavor,
needed_types: &mut TypeSet,
needed_types: &mut NeededTypeSet,
) -> Result<()> {
match *field {
pdb::TypeData::Member(ref data) => {
Expand Down Expand Up @@ -410,8 +415,10 @@ impl<'p> Class<'p> {

Ok(())
}
}

pub fn reconstruct(
impl ReconstructibleTypeData for Class<'_> {
fn reconstruct(
&self,
fmt_configuration: &DataFormatConfiguration,
f: &mut impl std::fmt::Write,
Expand Down Expand Up @@ -464,7 +471,7 @@ impl<'p> Class<'p> {
if !self.nested_enums.is_empty() {
writeln!(f, " ")?;
for e in &self.nested_enums {
e.reconstruct(f)?;
e.reconstruct(fmt_configuration, f)?;
}
}

Expand Down
Loading

0 comments on commit 5eb8a50

Please sign in to comment.