Use callback-based interface to load ThinLTO import data into rustc.

4 files changed, 58 insertions, 96 deletions

diff --git a/src/librustc_codegen_llvm/back/lto.rs b/src/librustc_codegen_llvm/back/lto.rs
index a68f22b2651..55499552f5c 100644
--- a/src/librustc_codegen_llvm/back/lto.rs
+++ b/src/librustc_codegen_llvm/back/lto.rs
@@ -805,7 +805,7 @@ pub struct ThinLTOImports {
 
 impl ThinLTOImports {
 
-    pub fn new_empty() -> ThinLTOImports {
+    pub fn new() -> ThinLTOImports {
         ThinLTOImports {
             imports: FxHashMap(),
         }
@@ -813,58 +813,46 @@ 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;
+        fn module_name_to_str(c_str: &CStr) -> &str {
+            match c_str.to_str() {
+                Ok(s) => s,
+                Err(e) => {
+                    bug!("Encountered non-utf8 LLVM module name `{}`: {}",
+                        c_str.to_string_lossy(),
+                        e)
+                }
             }
+        }
 
-            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
-                }
+        unsafe extern "C" fn imported_module_callback(payload: *mut libc::c_void,
+                                                      importing_module_name: *const libc::c_char,
+                                                      imported_module_name: *const libc::c_char) {
+            let map = &mut* (payload as *mut ThinLTOImports);
 
-                let imported_module_name = CStr::from_ptr(imported_module_name)
-                    .to_str()
-                    .expect("Non-utf8 LLVM module name encountered")
-                    .to_owned();
+            let importing_module_name = CStr::from_ptr(importing_module_name);
+            let importing_module_name = module_name_to_str(&importing_module_name);
+            let imported_module_name = CStr::from_ptr(imported_module_name);
+            let imported_module_name = module_name_to_str(&imported_module_name);
 
-                imported_modules.push(imported_module_name);
-                entry_ptr = entry_ptr.offset(1);
+            if !map.imports.contains_key(importing_module_name) {
+                map.imports.insert(importing_module_name.to_owned(), vec![]);
             }
 
-            imports.insert(importing_module_name, imported_modules);
-
-            module_ptr = module_ptr.offset(1);
-            module_index += 1;
+            map.imports
+               .get_mut(importing_module_name)
+               .unwrap()
+               .push(imported_module_name.to_owned());
         }
 
-        assert_eq!(module_index, imports.len());
+        let mut map = ThinLTOImports {
+            imports: FxHashMap(),
+        };
 
-        ThinLTOImports {
-            imports
-        }
+        llvm::LLVMRustGetThinLTOModuleImports(data,
+                                              imported_module_callback,
+                                              &mut map as *mut _ as *mut libc::c_void);
+        map
     }
 
     pub fn save_to_file(&self, path: &Path) -> io::Result<()> {
diff --git a/src/librustc_codegen_llvm/base.rs b/src/librustc_codegen_llvm/base.rs
index b4ba06c30be..e984b73471a 100644
--- a/src/librustc_codegen_llvm/base.rs
+++ b/src/librustc_codegen_llvm/base.rs
@@ -1360,7 +1360,7 @@ fn load_thin_lto_imports(sess: &Session) -> lto::ThinLTOImports {
     );
 
     if !path.exists() {
-        return lto::ThinLTOImports::new_empty();
+        return lto::ThinLTOImports::new();
     }
 
     match lto::ThinLTOImports::load_from_file(&path) {
diff --git a/src/librustc_llvm/ffi.rs b/src/librustc_llvm/ffi.rs
index b3a3893bfc3..551f1cda6f0 100644
--- a/src/librustc_llvm/ffi.rs
+++ b/src/librustc_llvm/ffi.rs
@@ -350,10 +350,9 @@ pub enum ThinLTOData {}
 /// LLVMRustThinLTOBuffer
 pub enum ThinLTOBuffer {}
 
-/// LLVMRustThinLTOModuleName
-pub type ThinLTOModuleName = *const c_char;
-/// LLVMRustThinLTOModuleImports
-pub type ThinLTOModuleImports = *const ThinLTOModuleName;
+// LLVMRustModuleNameCallback
+pub type ThinLTOModuleNameCallback =
+    unsafe extern "C" fn(*mut c_void, *const c_char, *const c_char);
 
 /// LLVMRustThinLTOModule
 #[repr(C)]
@@ -1785,7 +1784,9 @@ extern "C" {
     ) -> bool;
     pub fn LLVMRustGetThinLTOModuleImports(
         Data: *const ThinLTOData,
-    ) -> *const ThinLTOModuleImports;
+        ModuleNameCallback: ThinLTOModuleNameCallback,
+        CallbackPayload: *mut c_void,
+    );
     pub fn LLVMRustFreeThinLTOData(Data: *mut ThinLTOData);
     pub fn LLVMRustParseBitcodeForThinLTO(
         Context: ContextRef,
diff --git a/src/rustllvm/PassWrapper.cpp b/src/rustllvm/PassWrapper.cpp
index 30f585efedc..1f96b9042ba 100644
--- a/src/rustllvm/PassWrapper.cpp
+++ b/src/rustllvm/PassWrapper.cpp
@@ -798,11 +798,6 @@ 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
@@ -1104,50 +1099,28 @@ 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;
+extern "C" typedef void (*LLVMRustModuleNameCallback)(void*, // payload
+                                                      const char*, // importing module name
+                                                      const char*); // imported module name
 
-    thinLTOModuleImports[module_index] = imports_array;
-    module_index += 1;
+// Calls `module_name_callback` for each module import done by ThinLTO.
+// The callback is provided with regular null-terminated C strings.
+extern "C" void
+LLVMRustGetThinLTOModuleImports(const LLVMRustThinLTOData *data,
+                                LLVMRustModuleNameCallback module_name_callback,
+                                void* callback_payload) {
+  for (const auto& importing_module : data->ImportLists) {
+    const std::string importing_module_id = importing_module.getKey().str();
+
+    const auto& imports = importing_module.getValue();
+
+    for (const auto& imported_module : imports) {
+      const std::string imported_module_id = imported_module.getKey().str();
+      module_name_callback(callback_payload,
+                           importing_module_id.c_str(),
+                           imported_module_id.c_str());
+    }
   }
-
-  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