Provide a way of accessing the ThinLTO module import map in rustc.

Author: michaelwoerister

src/librustc_codegen_llvm/back/lto.rs

@@ -20,12 +20,13 @@ use rustc::hir::def_id::LOCAL_CRATE;
 use rustc::middle::exported_symbols::SymbolExportLevel;
 use rustc::session::config::{self, Lto};
 use rustc::util::common::time_ext;
+use rustc_data_structures::fx::FxHashMap;
 use time_graph::Timeline;
 use {ModuleCodegen, ModuleLlvm, ModuleKind, ModuleSource};
 
 use libc;
 
-use std::ffi::CString;
+use std::ffi::{CString, CStr};
 use std::ptr;
 use std::slice;
 use std::sync::Arc;
@@ -776,3 +777,69 @@ impl ThinModule {
         Ok(module)
     }
 }
+
+
+#[derive(Debug)]
+pub struct ThinLTOImports {
+    // key = llvm name of importing module, value = list of modules it imports from
+    imports: FxHashMap<String, Vec<String>>,
+}
+
+impl ThinLTOImports {
+
+    /// Load the ThinLTO import map from ThinLTOData.
+    unsafe fn from_thin_lto_data(data: *const llvm::ThinLTOData) -> ThinLTOImports {
+        let raw_data: *const llvm::ThinLTOModuleImports =
+            llvm::LLVMRustGetThinLTOModuleImports(data);
+
+        assert!(!raw_data.is_null());
+
+        let mut imports = FxHashMap();
+        let mut module_ptr = raw_data;
+        let mut module_index = 0;
+
+        loop {
+            let mut entry_ptr: *const llvm::ThinLTOModuleName = *module_ptr;
+
+            if entry_ptr.is_null() {
+                break;
+            }
+
+            let importing_module_name = CStr::from_ptr(*entry_ptr)
+                .to_str()
+                .expect("Non-utf8 LLVM module name encountered")
+                .to_owned();
+
+            entry_ptr = entry_ptr.offset(1);
+
+            let mut imported_modules = vec![];
+
+            loop {
+                let imported_module_name = *entry_ptr;
+
+                if imported_module_name.is_null() {
+                    break
+                }
+
+                let imported_module_name = CStr::from_ptr(imported_module_name)
+                    .to_str()
+                    .expect("Non-utf8 LLVM module name encountered")
+                    .to_owned();
+
+                imported_modules.push(imported_module_name);
+                entry_ptr = entry_ptr.offset(1);
+            }
+
+            imports.insert(importing_module_name, imported_modules);
+
+            module_ptr = module_ptr.offset(1);
+            module_index += 1;
+        }
+
+        assert_eq!(module_index, imports.len());
+
+        ThinLTOImports {
+            imports
+        }
+    }
+}

src/librustc_llvm/ffi.rs

@@ -350,6 +350,11 @@ pub enum ThinLTOData {}
 /// LLVMRustThinLTOBuffer
 pub enum ThinLTOBuffer {}
 
+/// LLVMRustThinLTOModuleName
+pub type ThinLTOModuleName = *const c_char;
+/// LLVMRustThinLTOModuleImports
+pub type ThinLTOModuleImports = *const ThinLTOModuleName;
+
 /// LLVMRustThinLTOModule
 #[repr(C)]
 pub struct ThinLTOModule {
@@ -1778,6 +1783,9 @@ extern "C" {
         Data: *const ThinLTOData,
         Module: ModuleRef,
     ) -> bool;
+    pub fn LLVMRustGetThinLTOModuleImports(
+        Data: *const ThinLTOData,
+    ) -> *const ThinLTOModuleImports;
     pub fn LLVMRustFreeThinLTOData(Data: *mut ThinLTOData);
     pub fn LLVMRustParseBitcodeForThinLTO(
         Context: ContextRef,

src/rustllvm/PassWrapper.cpp

@@ -798,6 +798,11 @@ LLVMRustPGOAvailable() {
 #endif
 }
 
+// We encode the ThinLTO module import map as a nested null-terminated list to
+// get it into Rust.
+typedef const char* LLVMRustThinLTOModuleName;
+typedef LLVMRustThinLTOModuleName* LLVMRustThinLTOModuleImports;
+
 #if LLVM_VERSION_GE(4, 0)
 
 // Here you'll find an implementation of ThinLTO as used by the Rust compiler
@@ -1099,6 +1104,52 @@ LLVMRustPrepareThinLTOImport(const LLVMRustThinLTOData *Data, LLVMModuleRef M) {
   return true;
 }
 
+/// Converts the LLVMRustThinLTOData::ImportLists map into a nested list. The
+/// first level is a null-terminated array with an entry for each module. Each
+/// entry is a pointer that points to a null-termined array of module names. The
+/// first entry is always the name of the *importing* module, the following
+/// entries are  the names of the modules it imports from. Each module name is
+/// a regular C string.
+extern "C" LLVMRustThinLTOModuleImports*
+LLVMRustGetThinLTOModuleImports(const LLVMRustThinLTOData *Data) {
+  // Allocate number of module +1. This is a null-terminated array.
+  LLVMRustThinLTOModuleImports* thinLTOModuleImports =
+    new LLVMRustThinLTOModuleImports[Data->ImportLists.size() + 1];
+  size_t module_index = 0;
+
+  for (const auto & module : Data->ImportLists) {
+    StringRef module_id = module.getKey();
+    const auto& imports = module.getValue();
+
+    // Allocate number of imported module + 2, one extra for the name of the
+    // importing module and another one for null-termination.
+    LLVMRustThinLTOModuleImports imports_array =
+      new LLVMRustThinLTOModuleName[imports.size() + 2];
+
+    // The first value is always the name of the *importing* module.
+    imports_array[0] = strndup(module_id.data(), module_id.size());
+
+    size_t imports_array_index = 1;
+    for (const auto imported_module_id : imports.keys()) {
+      // The following values are the names of the imported modules.
+      imports_array[imports_array_index] = strndup(imported_module_id.data(),
+                                                   imported_module_id.size());
+      imports_array_index += 1;
+    }
+
+    assert(imports_array_index == imports.size() + 1);
+    imports_array[imports_array_index] = nullptr;
+
+    thinLTOModuleImports[module_index] = imports_array;
+    module_index += 1;
+  }
+
+  assert(module_index == Data->ImportLists.size());
+  thinLTOModuleImports[module_index] = nullptr;
+
+  return thinLTOModuleImports;
+}
+
 // This struct and various functions are sort of a hack right now, but the
 // problem is that we've got in-memory LLVM modules after we generate and
 // optimize all codegen-units for one compilation in rustc. To be compatible
@@ -1280,6 +1331,11 @@ LLVMRustPrepareThinLTOImport(const LLVMRustThinLTOData *Data, LLVMModuleRef M) {
   report_fatal_error("ThinLTO not available");
 }
 
+extern "C" LLVMRustThinLTOModuleImports
+LLVMRustGetLLVMRustThinLTOModuleImports(const LLVMRustThinLTOData *Data) {
+  report_fatal_error("ThinLTO not available");
+}
+
 extern "C" void
 LLVMRustFreeThinLTOData(LLVMRustThinLTOData *Data) {
   report_fatal_error("ThinLTO not available");