2 files changed, 22 insertions, 3 deletions
diff --git a/lib/Lex/HeaderSearch.cpp b/lib/Lex/HeaderSearch.cpp
index 6403cfb868..4522de5cf3 100644
--- a/lib/Lex/HeaderSearch.cpp
+++ b/lib/Lex/HeaderSearch.cpp
@@ -205,7 +205,7 @@ const FileEntry *DirectoryLookup::LookupFile(
       
       // If there is a module that corresponds to this header, 
       // suggest it.
-      StringRef Module = HS.getModuleForHeader(File);
+      StringRef Module = HS.findModuleForHeader(File);
       if (!Module.empty() && Module != BuildingModule)
         *SuggestedModule = Module;
       
@@ -772,8 +772,10 @@ bool HeaderSearch::hasModuleMap(StringRef FileName,
   return false;
 }
 
-StringRef HeaderSearch::getModuleForHeader(const FileEntry *File) {
-  // FIXME: Actually look for the corresponding module for this header.
+StringRef HeaderSearch::findModuleForHeader(const FileEntry *File) {
+  if (ModuleMap::Module *Module = ModMap.findModuleForHeader(File))
+    return Module->getTopLevelModuleName();
+  
   return StringRef();
 }
 
diff --git a/lib/Lex/ModuleMap.cpp b/lib/Lex/ModuleMap.cpp
index 8adb22014b..7defe01a67 100644
--- a/lib/Lex/ModuleMap.cpp
+++ b/lib/Lex/ModuleMap.cpp
@@ -51,6 +51,14 @@ std::string ModuleMap::Module::getFullModuleName() const {
   return Result;
 }
 
+StringRef ModuleMap::Module::getTopLevelModuleName() const {
+  const Module *Top = this;
+  while (Top->Parent)
+    Top = Top->Parent;
+  
+  return Top->Name;
+}
+
 //----------------------------------------------------------------------------//
 // Module map
 //----------------------------------------------------------------------------//
@@ -67,6 +75,15 @@ ModuleMap::~ModuleMap() {
   delete SourceMgr;
 }
 
+ModuleMap::Module *ModuleMap::findModuleForHeader(const FileEntry *File) {
+  llvm::DenseMap<const FileEntry *, Module *>::iterator Known
+    = Headers.find(File);
+  if (Known != Headers.end())
+    return Known->second;
+  
+  return 0;
+}
+
 static void indent(llvm::raw_ostream &OS, unsigned Spaces) {
   OS << std::string(' ', Spaces);
 }