diff --git a/lib/Frontend/ASTUnit.cpp b/lib/Frontend/ASTUnit.cpp
index 48a6983..2b4eecd 100644
--- a/lib/Frontend/ASTUnit.cpp
+++ b/lib/Frontend/ASTUnit.cpp
@@ -359,7 +359,7 @@
 
   Act.reset(new TopLevelDeclTrackerAction(*AST));
   if (!Act->BeginSourceFile(Clang, Clang.getFrontendOpts().Inputs[0].second,
-                           /*IsAST=*/false))
+                            Clang.getFrontendOpts().Inputs[0].first))
     goto error;
 
   Act->Execute();
diff --git a/lib/Frontend/CompilerInstance.cpp b/lib/Frontend/CompilerInstance.cpp
index 4f6ed56..0af355c 100644
--- a/lib/Frontend/CompilerInstance.cpp
+++ b/lib/Frontend/CompilerInstance.cpp
@@ -509,7 +509,7 @@
       createPreprocessor();
     }
 
-    if (Act.BeginSourceFile(*this, InFile, IsAST)) {
+    if (Act.BeginSourceFile(*this, InFile, getFrontendOpts().Inputs[i].first)) {
       Act.Execute();
       Act.EndSourceFile();
     }
diff --git a/lib/Frontend/FrontendAction.cpp b/lib/Frontend/FrontendAction.cpp
index 87fc122..66bec7c 100644
--- a/lib/Frontend/FrontendAction.cpp
+++ b/lib/Frontend/FrontendAction.cpp
@@ -32,7 +32,7 @@
 
 bool FrontendAction::BeginSourceFile(CompilerInstance &CI,
                                      llvm::StringRef Filename,
-                                     bool IsAST) {
+                                     InputKind InputKind) {
   assert(!Instance && "Already processing a source file!");
   assert(!Filename.empty() && "Unexpected empty filename!");
   setCurrentFile(Filename);
@@ -40,6 +40,7 @@
 
   // AST files follow a very different path, since they share objects via the
   // AST unit.
+  bool IsAST = InputKind == IK_AST;
   if (IsAST) {
     assert(!usesPreprocessorOnly() &&
            "Attempt to pass AST file to preprocessor only action!");