3 files changed, 103 insertions, 8 deletions

diff --git a/lib/Frontend/DiagnosticRenderer.cpp b/lib/Frontend/DiagnosticRenderer.cpp
index 3599df82c7..4a332ae82f 100644
--- a/lib/Frontend/DiagnosticRenderer.cpp
+++ b/lib/Frontend/DiagnosticRenderer.cpp
@@ -132,7 +132,7 @@ void DiagnosticRenderer::emitDiagnostic(SourceLocation Loc,
   
     // First, if this diagnostic is not in the main file, print out the
     // "included from" lines.
-    emitIncludeStack(PLoc.getIncludeLoc(), Level, *SM);
+    emitIncludeStack(Loc, PLoc, Level, *SM);
   }
   
   // Next, emit the actual diagnostic message.
@@ -184,21 +184,30 @@ void DiagnosticRenderer::emitStoredDiagnostic(StoredDiagnostic &Diag) {
 /// repeated warnings occur within the same file. It also handles the logic
 /// of customizing the formatting and display of the include stack.
 ///
+/// \param Loc   The diagnostic location.
+/// \param PLoc  The presumed location of the diagnostic location.
 /// \param Level The diagnostic level of the message this stack pertains to.
-/// \param Loc   The include location of the current file (not the diagnostic
-///              location).
 void DiagnosticRenderer::emitIncludeStack(SourceLocation Loc,
+                                          PresumedLoc PLoc,
                                           DiagnosticsEngine::Level Level,
                                           const SourceManager &SM) {
+  SourceLocation IncludeLoc = PLoc.getIncludeLoc();
+
   // Skip redundant include stacks altogether.
-  if (LastIncludeLoc == Loc)
+  if (LastIncludeLoc == IncludeLoc)
     return;
-  LastIncludeLoc = Loc;
+  
+  LastIncludeLoc = IncludeLoc;
   
   if (!DiagOpts->ShowNoteIncludeStack && Level == DiagnosticsEngine::Note)
     return;
-  
-  emitIncludeStackRecursively(Loc, SM);
+
+  if (IncludeLoc.isValid())
+    emitIncludeStackRecursively(IncludeLoc, SM);
+  else {
+    emitModuleBuildPath(SM);
+    emitImportStack(Loc, SM);
+  }
 }
 
 /// \brief Helper to recursivly walk up the include stack and print each layer
@@ -213,7 +222,17 @@ void DiagnosticRenderer::emitIncludeStackRecursively(SourceLocation Loc,
   PresumedLoc PLoc = SM.getPresumedLoc(Loc, DiagOpts->ShowPresumedLoc);
   if (PLoc.isInvalid())
     return;
-  
+
+  // If this source location was imported from a module, print the module
+  // import stack rather than the 
+  // FIXME: We want submodule granularity here.
+  std::pair<SourceLocation, StringRef> Imported = SM.getModuleImportLoc(Loc);
+  if (Imported.first.isValid()) {
+    // This location was imported by a module. Emit the module import stack.
+    emitImportStackRecursively(Imported.first, Imported.second, SM);
+    return;
+  }
+
   // Emit the other include frames first.
   emitIncludeStackRecursively(PLoc.getIncludeLoc(), SM);
   
@@ -221,6 +240,41 @@ void DiagnosticRenderer::emitIncludeStackRecursively(SourceLocation Loc,
   emitIncludeLocation(Loc, PLoc, SM);
 }
 
+/// \brief Emit the module import stack associated with the current location.
+void DiagnosticRenderer::emitImportStack(SourceLocation Loc,
+                                         const SourceManager &SM) {
+  if (Loc.isInvalid()) {
+    emitModuleBuildPath(SM);
+    return;
+  }
+
+  std::pair<SourceLocation, StringRef> NextImportLoc
+    = SM.getModuleImportLoc(Loc);
+  emitImportStackRecursively(NextImportLoc.first, NextImportLoc.second, SM);
+}
+
+/// \brief Helper to recursivly walk up the import stack and print each layer
+/// on the way back down.
+void DiagnosticRenderer::emitImportStackRecursively(SourceLocation Loc,
+                                                    StringRef ModuleName,
+                                                    const SourceManager &SM) {
+  if (Loc.isInvalid()) {
+    return;
+  }
+
+  PresumedLoc PLoc = SM.getPresumedLoc(Loc, DiagOpts->ShowPresumedLoc);
+  if (PLoc.isInvalid())
+    return;
+
+  // Emit the other import frames first.
+  std::pair<SourceLocation, StringRef> NextImportLoc
+    = SM.getModuleImportLoc(Loc);
+  emitImportStackRecursively(NextImportLoc.first, NextImportLoc.second, SM);
+
+  // Emit the inclusion text/note.
+  emitImportLocation(Loc, PLoc, ModuleName, SM);
+}
+
 /// \brief Emit the module build path, for cases where a module is (re-)built
 /// on demand.
 void DiagnosticRenderer::emitModuleBuildPath(const SourceManager &SM) {
@@ -407,6 +461,18 @@ void DiagnosticNoteRenderer::emitIncludeLocation(SourceLocation Loc,
   emitNote(Loc, Message.str(), &SM);
 }
 
+void DiagnosticNoteRenderer::emitImportLocation(SourceLocation Loc,
+                                                PresumedLoc PLoc,
+                                                StringRef ModuleName,
+                                                const SourceManager &SM) {
+  // Generate a note indicating the include location.
+  SmallString<200> MessageStorage;
+  llvm::raw_svector_ostream Message(MessageStorage);
+  Message << "in module '" << ModuleName << "' imported from "
+          << PLoc.getFilename() << ':' << PLoc.getLine() << ":";
+  emitNote(Loc, Message.str(), &SM);
+}
+
 void
 DiagnosticNoteRenderer::emitBuildingModuleLocation(SourceLocation Loc,
                                                    PresumedLoc PLoc,
diff --git a/lib/Frontend/TextDiagnostic.cpp b/lib/Frontend/TextDiagnostic.cpp
index 65df875aea..fe82229c91 100644
--- a/lib/Frontend/TextDiagnostic.cpp
+++ b/lib/Frontend/TextDiagnostic.cpp
@@ -884,6 +884,16 @@ void TextDiagnostic::emitIncludeLocation(SourceLocation Loc,
     OS << "In included file:\n"; 
 }
 
+void TextDiagnostic::emitImportLocation(SourceLocation Loc, PresumedLoc PLoc,
+                                        StringRef ModuleName,
+                                        const SourceManager &SM) {
+  if (DiagOpts->ShowLocation)
+    OS << "In module '" << ModuleName << "' imported from "
+       << PLoc.getFilename() << ':' << PLoc.getLine() << ":\n";
+  else
+    OS << "In module " << ModuleName << "':\n";
+}
+
 void TextDiagnostic::emitBuildingModuleLocation(SourceLocation Loc,
                                                 PresumedLoc PLoc,
                                                 StringRef ModuleName,
diff --git a/lib/Serialization/ASTReader.cpp b/lib/Serialization/ASTReader.cpp
index c1428d5c68..908d8c573c 100644
--- a/lib/Serialization/ASTReader.cpp
+++ b/lib/Serialization/ASTReader.cpp
@@ -1000,6 +1000,25 @@ bool ASTReader::ReadSLocEntry(int ID) {
   return false;
 }
 
+std::pair<SourceLocation, StringRef> ASTReader::getModuleImportLoc(int ID) {
+  if (ID == 0)
+    return std::make_pair(SourceLocation(), "");
+
+  if (unsigned(-ID) - 2 >= getTotalNumSLocs() || ID > 0) {
+    Error("source location entry ID out-of-range for AST file");
+    return std::make_pair(SourceLocation(), "");
+  }
+
+  // Find which module file this entry lands in.
+  ModuleFile *M = GlobalSLocEntryMap.find(-ID)->second;
+  if (M->Kind != MK_Module)
+    return std::make_pair(SourceLocation(), "");
+
+  // FIXME: Can we map this down to a particular submodule? That would be
+  // ideal.
+  return std::make_pair(M->ImportLoc, llvm::sys::path::stem(M->FileName));
+}
+
 /// \brief Find the location where the module F is imported.
 SourceLocation ASTReader::getImportLocation(ModuleFile *F) {
   if (F->ImportLoc.isValid())