Index: lib/Format/Format.cpp
===================================================================
--- lib/Format/Format.cpp
+++ lib/Format/Format.cpp
@@ -1436,6 +1436,55 @@
          llvm::Regex(IncludeRegexPattern).match(Replace.getReplacementText());
 }
 
+void skipComments(Lexer &Lex, Token &Tok) {
+  while (Tok.is(tok::comment))
+    if (Lex.LexFromRawLexer(Tok))
+      return;
+}
+
+// Check if a sequence of tokens is like "#<Name> <raw_identifier>". If it is,
+// \p Tok will be the token after this directive; otherwise, it can be any token
+// after the given \p Tok (including \p Tok).
+bool checkAndConsumeDirectiveWithName(Lexer &Lex, StringRef Name, Token &Tok) {
+  bool Matched = Tok.is(tok::hash) && !Lex.LexFromRawLexer(Tok) &&
+                 Tok.is(tok::raw_identifier) &&
+                 Tok.getRawIdentifier() == Name && !Lex.LexFromRawLexer(Tok) &&
+                 Tok.is(tok::raw_identifier);
+  if (Matched)
+    Lex.LexFromRawLexer(Tok);
+  return Matched;
+}
+
+unsigned getOffsetAfterHeaderGuardsAndComments(StringRef FileName,
+                                               StringRef Code,
+                                               FormatStyle Style) {
+  std::unique_ptr<Environment> Env =
+      Environment::CreateVirtualEnvironment(Code, FileName, /*Ranges=*/{});
+  const SourceManager &SourceMgr = Env->getSourceManager();
+  Lexer Lex(Env->getFileID(), SourceMgr.getBuffer(Env->getFileID()), SourceMgr,
+            getFormattingLangOpts(Style));
+  unsigned Ret = Code.size();
+  Token Tok;
+  // Get the first token.
+  Lex.LexFromRawLexer(Tok);
+  skipComments(Lex, Tok);
+  Token TokAfterComments = Tok;
+  bool HeaderGuardFound = false;
+  if (checkAndConsumeDirectiveWithName(Lex, "ifndef", Tok)) {
+    skipComments(Lex, Tok);
+    if (checkAndConsumeDirectiveWithName(Lex, "define", Tok))
+      HeaderGuardFound = true;
+  }
+  if (!HeaderGuardFound)
+    Tok = TokAfterComments;
+  return Tok.isNot(tok::eof) ? SourceMgr.getFileOffset(Tok.getLocation())
+                             : Ret;
+}
+
+// FIXME: we always +1 when we calculate the offset for each new line; however,
+// the last line does not necessarily end with '\n', which makes offset exceed
+// code size. Program crashes when we insert at the end of code which does not
+// ends with '\n'.
 // FIXME: do not insert headers into conditional #include blocks, e.g. #includes
 // surrounded by compile condition "#if...".
 // FIXME: do not insert existing headers.
@@ -1469,20 +1518,6 @@
   StringRef FileName = Replaces.begin()->getFilePath();
   IncludeCategoryManager Categories(Style, FileName);
 
-  std::unique_ptr<Environment> Env =
-      Environment::CreateVirtualEnvironment(Code, FileName, /*Ranges=*/{});
-  const SourceManager &SourceMgr = Env->getSourceManager();
-  Lexer Lex(Env->getFileID(), SourceMgr.getBuffer(Env->getFileID()), SourceMgr,
-            getFormattingLangOpts(Style));
-  Token Tok;
-  // All new headers should be inserted after this offset.
-  int MinInsertOffset = Code.size();
-  while (!Lex.LexFromRawLexer(Tok)) {
-    if (Tok.isNot(tok::comment)) {
-      MinInsertOffset = SourceMgr.getFileOffset(Tok.getLocation());
-      break;
-    }
-  }
   // Record the offset of the end of the last include in each category.
   std::map<int, int> CategoryEndOffsets;
   // All possible priorities.
@@ -1491,7 +1526,9 @@
   for (const auto &Category : Style.IncludeCategories)
     Priorities.insert(Category.Priority);
   int FirstIncludeOffset = -1;
-  bool HeaderGuardFound = false;
+  // All new headers should be inserted after this offset.
+  unsigned MinInsertOffset =
+      getOffsetAfterHeaderGuardsAndComments(FileName, Code, Style);
   StringRef TrimmedCode = Code.drop_front(MinInsertOffset);
   SmallVector<StringRef, 32> Lines;
   TrimmedCode.split(Lines, '\n');
@@ -1506,11 +1543,6 @@
         FirstIncludeOffset = Offset;
     }
     Offset += Line.size() + 1;
-    // FIXME: make header guard matching stricter, e.g. consider #ifndef.
-    if (!HeaderGuardFound && DefineRegex.match(Line)) {
-      HeaderGuardFound = true;
-      MinInsertOffset = Offset;
-    }
   }
 
   // Populate CategoryEndOfssets:
Index: unittests/Format/CleanupTest.cpp
===================================================================
--- unittests/Format/CleanupTest.cpp
+++ unittests/Format/CleanupTest.cpp
@@ -608,6 +608,77 @@
   EXPECT_EQ(Expected, apply(Code, Replaces));
 }
 
+TEST_F(CleanUpReplacementsTest, FakeHeaderGuardIfDef) {
+  std::string Code = "// comment \n"
+                     "#ifdef X\n"
+                     "#define X\n";
+  std::string Expected = "// comment \n"
+                         "#include <vector>\n"
+                         "#ifdef X\n"
+                         "#define X\n";
+  tooling::Replacements Replaces = {createInsertion("#include <vector>")};
+  EXPECT_EQ(Expected, apply(Code, Replaces));
+}
+
+TEST_F(CleanUpReplacementsTest, RealHeaderGuardAfterComments) {
+  std::string Code = "// comment \n"
+                     "#ifndef X\n"
+                     "#define X\n"
+                     "int x;\n"
+                     "#define Y 1\n";
+  std::string Expected = "// comment \n"
+                         "#ifndef X\n"
+                         "#define X\n"
+                         "#include <vector>\n"
+                         "int x;\n"
+                         "#define Y 1\n";
+  tooling::Replacements Replaces = {createInsertion("#include <vector>")};
+  EXPECT_EQ(Expected, apply(Code, Replaces));
+}
+
+TEST_F(CleanUpReplacementsTest, IfNDefWithNoDefine) {
+  std::string Code = "// comment \n"
+                     "#ifndef X\n"
+                     "int x;\n"
+                     "#define Y 1\n";
+  std::string Expected = "// comment \n"
+                         "#include <vector>\n"
+                         "#ifndef X\n"
+                         "int x;\n"
+                         "#define Y 1\n";
+  tooling::Replacements Replaces = {createInsertion("#include <vector>")};
+  EXPECT_EQ(Expected, apply(Code, Replaces));
+}
+
+TEST_F(CleanUpReplacementsTest, HeaderGuardWithComment) {
+  std::string Code = "// comment \n"
+                     "#ifndef X // comment\n"
+                     "// comment\n"
+                     "/* comment\n"
+                     "*/\n"
+                     "/* comment */ #define X\n"
+                     "int x;\n"
+                     "#define Y 1\n";
+  std::string Expected = "// comment \n"
+                         "#ifndef X // comment\n"
+                         "// comment\n"
+                         "/* comment\n"
+                         "*/\n"
+                         "/* comment */ #define X\n"
+                         "#include <vector>\n"
+                         "int x;\n"
+                         "#define Y 1\n";
+  tooling::Replacements Replaces = {createInsertion("#include <vector>")};
+  EXPECT_EQ(Expected, apply(Code, Replaces));
+}
+
+TEST_F(CleanUpReplacementsTest, EmptyCode) {
+  std::string Code = "";
+  std::string Expected = "#include <vector>\n";
+  tooling::Replacements Replaces = {createInsertion("#include <vector>")};
+  EXPECT_EQ(Expected, apply(Code, Replaces));
+}
+
 } // end namespace
 } // end namespace format
 } // end namespace clang