diff --git a/clang/lib/Lex/TokenLexer.cpp b/clang/lib/Lex/TokenLexer.cpp
--- a/clang/lib/Lex/TokenLexer.cpp
+++ b/clang/lib/Lex/TokenLexer.cpp
@@ -993,6 +993,16 @@
   llvm::MutableArrayRef<Token> All(begin_tokens, end_tokens);
   llvm::MutableArrayRef<Token> Partition;
 
+  auto NearLast = [&, Last = BeginLoc](SourceLocation Loc) mutable {
+    // The maximum distance between two consecutive tokens in a partition.
+    // This is an important trick to avoid using too much SourceLocation address
+    // space!
+    static constexpr SourceLocation::IntTy MaxDistance = 50;
+    auto Distance = Loc.getRawEncoding() - Last.getRawEncoding();
+    Last = Loc;
+    return Distance <= MaxDistance;
+  };
+
   // Partition the tokens by their FileID.
   // This is a hot function, and calling getFileID can be expensive, the
   // implementation is optimized by reducing the number of getFileID.
@@ -1000,7 +1010,7 @@
     // Consecutive tokens not written in macros must be from the same file.
     // (Neither #include nor eof can occur inside a macro argument.)
     Partition = All.take_while([&](const Token &T) {
-      return T.getLocation().isFileID();
+      return T.getLocation().isFileID() && NearLast(T.getLocation());
     });
   } else {
     // Call getFileID once to calculate the bounds, and use the cheaper
@@ -1009,7 +1019,8 @@
     SourceLocation Limit =
         SM.getComposedLoc(BeginFID, SM.getFileIDSize(BeginFID));
     Partition = All.take_while([&](const Token &T) {
-      return T.getLocation() >= BeginLoc && T.getLocation() < Limit;
+      return T.getLocation() >= BeginLoc && T.getLocation() < Limit &&
+             NearLast(T.getLocation());
     });
   }
   assert(!Partition.empty());
diff --git a/clang/test/Lexer/update_consecutive_macro_address_space.c b/clang/test/Lexer/update_consecutive_macro_address_space.c
new file mode 100644
--- /dev/null
+++ b/clang/test/Lexer/update_consecutive_macro_address_space.c
@@ -0,0 +1,36 @@
+// RUN: %clang -cc1 -print-stats %s 2>&1 | FileCheck %s
+// CHECK: 6 local SLocEntry's allocated
+//
+// Verify that the macro arg expansion is split to two file ids, we have 6 file
+// ids rather than 5:
+//   0: invalid file id
+//   1: main file
+//   2: builtin file
+//   3: macro expansion for X
+//   4: macro arg expansions for 1
+//   5: macro arg expansions for == 2
+#define X(x) (int)(x);
+void func() {
+  X(1
+/*************************************************************************************************/
+/*************************************************************************************************/
+/*************************************************************************************************/
+/*************************************************************************************************/
+/*************************************************************************************************/
+/*************************************************************************************************/
+/*************************************************************************************************/
+/*************************************************************************************************/
+/*************************************************************************************************/
+/*************************************************************************************************/
+/*************************************************************************************************/
+/*************************************************************************************************/
+/*************************************************************************************************/
+/*************************************************************************************************/
+/*************************************************************************************************/
+/*************************************************************************************************/
+/*************************************************************************************************/
+/*************************************************************************************************/
+/*************************************************************************************************/
+/*************************************************************************************************/
+== 2);
+}
\ No newline at end of file