retracking imports/dependencies

crates/vestal/src/call_sites.rs

@@ -3,6 +3,9 @@ use crate::*;
 impl Module {
     /// Collect all calls that point to imports.
     pub fn load_call_sites (self: Arc<Self>) -> Usually<Arc<Self>> {
+        if self.verbose {
+            println!("  {DIM}(load-call-sites){RESET}");
+        }
         let mut decoder = Decoder::with_ip(64, self.code.as_ref(), 0, 0);
         while decoder.can_decode() {
             let position    = decoder.position();

crates/vestal/src/exports.rs

@@ -2,6 +2,9 @@ use crate::*;
 impl Module {
     /// Collect all exported methods.
     pub fn load_exports (self: Arc<Self>) -> Usually<Arc<Self>> {
+        if self.verbose {
+            println!("  {DIM}(load-exports){RESET}");
+        }
         let directory  = ExportDirectory::parse(self.pe.as_ref())?;
         let export_map = directory.get_export_map(self.pe.as_ref())?;
         let exports    = export_map.into_iter().map(|(k, v)|(k.into(), v)).collect();

crates/vestal/src/imports.rs

@@ -2,22 +2,40 @@ use crate::*;
 impl Module {
     /// Collect all imported modules and methods.
     pub fn load_imports (self: Arc<Self>, recurse: bool) -> Usually<Arc<Self>> {
+        if self.verbose {
+            println!("  {DIM}(load-imports){RESET}");
+        }
         let pe = self.pe.clone();
         let directory = ImportDirectory::parse(pe.as_ref())?;
         let mut new_modules = vec![];
         for descriptor in directory.descriptors.iter() {
-            let (name, imports) = Self::parse_import(pe.as_ref(), descriptor)?;
-            let exists = self.namespace.read().unwrap().contains_key(&name);
-            if !exists {
-                Log::dep(false, &name, &imports);
-                let path = self.find(&name)?.expect("not found");
+            let (module_name, imports) = Self::parse_import(pe.as_ref(), descriptor)?;
+            let create = !self.namespace.read().unwrap().contains_key(&module_name);
+            if create {
+                Log::dep(false, &module_name, &imports);
+                let path = self.find(&module_name)?.expect("not found");
                 let module = Arc::new(Self {
                     namespace:    self.namespace.clone(),
                     search_paths: self.search_paths.clone(),
                     ..Self::from_path(&path, self.verbose)?
                 });
+                assert_eq!(module_name, module.name);
                 self.namespace.write().unwrap().insert(module.name.clone(), module.clone());
-                new_modules.push(module);
+            }
+            let module = self.namespace.read().unwrap().get(&module_name).unwrap().clone();
+            self.dependencies.write().unwrap().insert(module.name.clone(), module.clone());
+            for (call_via, method_name) in imports {
+                self.load_import(call_via, &module_name, &method_name);
+            }
+            if create {
+                new_modules.push(module.clone());
+            }
+            if self.verbose {
+                if create {
+                    println!("    (dep {:?})", &module.name);
+                } else {
+                    println!("    {DIM}(dep {:?}){RESET}", &module.name);
+                }
             }
         }
         if recurse {
@@ -25,38 +43,37 @@ impl Module {
                 module.load(recurse)?;
             }
         }
-        if self.verbose {
-            self.show_dependencies();
-        }
         Ok(self)
     }
 
+    fn load_import (
+        &self, call_via: u32, module_name: &Arc<str>, method_name: &Arc<str>
+    ) {
+        if self.imports.read().unwrap().contains_key(&call_via) {
+            panic!("duplicate address {call_via} from {module_name}");
+        }
+        self.imports.write().unwrap().insert(call_via, (
+            module_name.clone(),
+            method_name.clone(),
+        ));
+    }
+
     /// Collect an imported dependency's descriptor thunks into a temporary [Vec].
     fn parse_import (pe: &VecPE, descriptor: &ImageImportDescriptor)
-        -> Usually<(Arc<str>, Vec<(usize, Arc<str>)>)>
+        -> Usually<(Arc<str>, Vec<(u32, Arc<str>)>)>
     {
-        let mut imports = vec![];
-        for (index, import) in descriptor.get_imports(pe)?.iter().enumerate() {
-            imports.push((index, match import {
-                ImportData::Ordinal(x) => format!("___VESTAL_ORDINAL_{x}"),
-                ImportData::ImportByName(name) => format!("{name}"),
-            }.into()));
-        }
         Ok((
             descriptor.get_name(pe)?.as_str()?.to_lowercase().into(),
-            imports
+            descriptor.get_imports(pe)?.iter().enumerate().map(|(index, import)|(
+                descriptor.first_thunk.0 + index as u32 * 8,
+                match import {
+                    ImportData::Ordinal(x) => format!("___VESTAL_ORDINAL_{x}"),
+                    ImportData::ImportByName(name) => format!("{name}"),
+                }.into()
+            )).collect()
         ))
     }
 
-    fn load_import (
-        &mut self, call_via: u32, module_name: &Arc<str>, method_name: &Arc<str>
-    ) {
-        if self.imports.contains_key(&call_via) {
-            panic!("duplicate address {call_via} from {module_name}");
-        }
-        self.imports.insert(call_via, (module_name.clone(), method_name.clone()));
-    }
-
     fn load_dependency (
         &mut self,
         module_name: &Arc<str>,

crates/vestal/src/main.rs

@@ -20,7 +20,7 @@ fn cli () -> clap::Command {
 /// A DLL that will be linked into the output.
 pub struct Module {
     /// Scope.
-    pub namespace: Arc<RwLock<BTreeMap<Arc<str>, Arc<Module>>>>,
+    pub namespace: Arc<RwLock<BTreeMap<Arc<str>, Arc<Self>>>>,
     /// Search paths for resolving DLL names.
     pub search_paths: Arc<RwLock<BTreeSet<Arc<PathBuf>>>>,
     /// Canonical name like `xxx.dll` (always lowercase)
@@ -40,9 +40,9 @@ pub struct Module {
     /// Size of `.text` section
     pub code_size: usize,
     /// Modules that this one depends on.
-    pub dependencies: BTreeMap<Arc<str>, Arc<Self>>,
+    pub dependencies: RwLock<BTreeMap<Arc<str>, Arc<Self>>>,
     /// Call targets to methods from imported modules.
-    pub imports: BTreeMap<u32, (Arc<str>, Arc<str>)>,
+    pub imports: RwLock<BTreeMap<u32, (Arc<str>, Arc<str>)>>,
     /// Addresses of exported methods by name
     pub exports: RwLock<BTreeMap<Arc<str>, ThunkData>>,
     /// Locations in `.text` section that need to be patched

crates/vestal/src/show.rs

@@ -24,7 +24,7 @@ impl Log {
             println!("{DIM}(load {:?}){RESET}", path.as_ref());
         }
     }
-    pub fn dep (show: bool, name: &impl AsRef<str>, exports: &[(usize, Arc<str>)]) {
+    pub fn dep (show: bool, name: &impl AsRef<str>, exports: &[(u32, Arc<str>)]) {
         if show {
             println!("(dep {:?} {})", name.as_ref(), exports.len());
         }
@@ -62,8 +62,8 @@ impl Show {
 
 impl std::fmt::Debug for Module {
     fn fmt (&self, f: &mut std::fmt::Formatter) -> std::result::Result<(), std::fmt::Error> {
-        write!(f, "(dll {BOLD}{UNDERLINE}{:15}{RESET} [0x{:>08x}] {} (img 0x{:>08x} -> mem 0x{:>08x}))",
-            &self.name,
+        write!(f, "(dll {} [0x{:>08x}] {} (img 0x{:>08x} -> mem 0x{:>08x}))",
+            format!("{BOLD}{UNDERLINE}{:15}{RESET}", self.name),
             self.code_size,
             format!("({:>5})", self.exports.read().unwrap().len()),
             self.code_start,
@@ -72,15 +72,6 @@ impl std::fmt::Debug for Module {
 }
 
 impl Module {
-    pub fn show_dependencies (&self) {
-        for (name, module) in self.dependencies.iter() {
-            print!("    ({module:?}");
-            for (method, addr) in module.exports.read().unwrap().iter() {
-                //print!("\n      (0x{addr:08x} {method})")
-            }
-            println!(")");
-        }
-    }
     pub fn show_instruction (&self, addr: usize) {
         let mut decoder = Decoder::with_ip(64, &self.code[addr..], 0x1000, DecoderOptions::NONE);
         let instruction = decoder.decode();