3 files changed, 42 insertions, 17 deletions

diff --git a/lib/Lex/PPCaching.cpp b/lib/Lex/PPCaching.cpp
index 33106591c3..986341b986 100644
--- a/lib/Lex/PPCaching.cpp
+++ b/lib/Lex/PPCaching.cpp
@@ -74,6 +74,8 @@ void Preprocessor::EnterCachingLexMode() {
     return;
 
   PushIncludeMacroStack();
+  if (CurLexerKind != CLK_LexAfterModuleImport)
+    CurLexerKind = CLK_CachingLexer;
 }
 
 
diff --git a/lib/Lex/PPLexerChange.cpp b/lib/Lex/PPLexerChange.cpp
index 758bfc4bb0..d52220798e 100644
--- a/lib/Lex/PPLexerChange.cpp
+++ b/lib/Lex/PPLexerChange.cpp
@@ -113,7 +113,9 @@ void Preprocessor::EnterSourceFileWithLexer(Lexer *TheLexer,
   CurLexer.reset(TheLexer);
   CurPPLexer = TheLexer;
   CurDirLookup = CurDir;
-
+  if (CurLexerKind != CLK_LexAfterModuleImport)
+    CurLexerKind = CLK_Lexer;
+  
   // Notify the client, if desired, that we are in a new source file.
   if (Callbacks && !CurLexer->Is_PragmaLexer) {
     SrcMgr::CharacteristicKind FileType =
@@ -135,7 +137,9 @@ void Preprocessor::EnterSourceFileWithPTH(PTHLexer *PL,
   CurDirLookup = CurDir;
   CurPTHLexer.reset(PL);
   CurPPLexer = CurPTHLexer.get();
-
+  if (CurLexerKind != CLK_LexAfterModuleImport)
+    CurLexerKind = CLK_PTHLexer;
+  
   // Notify the client, if desired, that we are in a new source file.
   if (Callbacks) {
     FileID FID = CurPPLexer->getFileID();
@@ -159,6 +163,8 @@ void Preprocessor::EnterMacro(Token &Tok, SourceLocation ILEnd,
     CurTokenLexer.reset(TokenLexerCache[--NumCachedTokenLexers]);
     CurTokenLexer->Init(Tok, ILEnd, Args);
   }
+  if (CurLexerKind != CLK_LexAfterModuleImport)
+    CurLexerKind = CLK_TokenLexer;
 }
 
 /// EnterTokenStream - Add a "macro" context to the top of the include stack,
@@ -188,6 +194,8 @@ void Preprocessor::EnterTokenStream(const Token *Toks, unsigned NumToks,
     CurTokenLexer.reset(TokenLexerCache[--NumCachedTokenLexers]);
     CurTokenLexer->Init(Toks, NumToks, DisableMacroExpansion, OwnsTokens);
   }
+  if (CurLexerKind != CLK_LexAfterModuleImport)
+    CurLexerKind = CLK_TokenLexer;
 }
 
 /// HandleEndOfFile - This callback is invoked when the lexer hits the end of
diff --git a/lib/Lex/Preprocessor.cpp b/lib/Lex/Preprocessor.cpp
index d33d9c513d..db3479e8f3 100644
--- a/lib/Lex/Preprocessor.cpp
+++ b/lib/Lex/Preprocessor.cpp
@@ -144,8 +144,6 @@ void Preprocessor::Initialize(const TargetInfo &Target) {
   // We haven't read anything from the external source.
   ReadMacrosFromExternalSource = false;
   
-  LexDepth = 0;
-  
   // "Poison" __VA_ARGS__, which can only appear in the expansion of a macro.
   // This gets unpoisoned where it is allowed.
   (Ident__VA_ARGS__ = getIdentifierInfo("__VA_ARGS__"))->setIsPoisoned();
@@ -524,27 +522,44 @@ void Preprocessor::HandleIdentifier(Token &Identifier) {
   // like "#define TY typeof", "TY(1) x".
   if (II.isExtensionToken() && !DisableMacroExpansion)
     Diag(Identifier, diag::ext_token_used);
+  
+  // If this is the '__import_module__' keyword, note that the next token
+  // indicates a module name.
+  if (II.getTokenID() == tok::kw___import_module__ &&
+      !InMacroArgs && !DisableMacroExpansion) {
+    ModuleImportLoc = Identifier.getLocation();
+    CurLexerKind = CLK_LexAfterModuleImport;
+  }
 }
 
-void Preprocessor::HandleModuleImport(Token &Import) {
+/// \brief Lex a token following the __import_module__ keyword.
+void Preprocessor::LexAfterModuleImport(Token &Result) {
+  // Figure out what kind of lexer we actually have.
+  if (CurLexer)
+    CurLexerKind = CLK_Lexer;
+  else if (CurPTHLexer)
+    CurLexerKind = CLK_PTHLexer;
+  else if (CurTokenLexer)
+    CurLexerKind = CLK_TokenLexer;
+  else 
+    CurLexerKind = CLK_CachingLexer;
+  
+  // Lex the next token.
+  Lex(Result);
+
   // The token sequence 
   //
   //   __import_module__ identifier
   //
-  // indicates a module import directive. We load the module and then 
-  // leave the token sequence for the parser.
-  Token ModuleNameTok = LookAhead(0);
-  if (ModuleNameTok.getKind() != tok::identifier)
+  // indicates a module import directive. We already saw the __import_module__
+  // keyword, so now we're looking for the identifier.
+  if (Result.getKind() != tok::identifier)
     return;
   
-  (void)TheModuleLoader.loadModule(Import.getLocation(),
-                                   *ModuleNameTok.getIdentifierInfo(), 
-                                   ModuleNameTok.getLocation());
-  
-  // FIXME: Transmogrify __import_module__ into some kind of AST-only 
-  // __import_module__ that is not recognized by the preprocessor but is 
-  // recognized by the parser. It would also be useful to stash the ModuleKey
-  // somewhere, so we don't try to load the module twice.
+  // Load the module.
+  (void)TheModuleLoader.loadModule(ModuleImportLoc,
+                                   *Result.getIdentifierInfo(), 
+                                   Result.getLocation());
 }
 
 void Preprocessor::AddCommentHandler(CommentHandler *Handler) {