From 2010dae4c7846b22d30de7746176efad67431f3f Mon Sep 17 00:00:00 2001
From: Rintaro Ishizaki <rishizaki@apple.com>
Date: Wed, 10 Jul 2024 10:44:37 -0700
Subject: [PATCH] [Macros] Optimize 'rangeContainsTokenLocWithGeneratedSource'

* Early return for the direct 'true' case.
* Instead of comparing the buffer IDs, check f the entire buffer
  contains the target range. This keeps 'LocCache.lastBufferID' and
  improves the hit rate.

(cherry picked from commit 126ad5ce79c63a33751f3d381f8b1153b80f98f1)
---
 lib/AST/TypeRefinementContext.cpp | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/lib/AST/TypeRefinementContext.cpp b/lib/AST/TypeRefinementContext.cpp
index 6bfe42fc92a96..9df647423513a 100644
--- a/lib/AST/TypeRefinementContext.cpp
+++ b/lib/AST/TypeRefinementContext.cpp
@@ -181,10 +181,13 @@ TypeRefinementContext::createForWhileStmtBody(ASTContext &Ctx, WhileStmt *S,
 /// range.
 static bool rangeContainsTokenLocWithGeneratedSource(
     SourceManager &sourceMgr, SourceRange parentRange, SourceLoc childLoc) {
-  auto parentBuffer = sourceMgr.findBufferContainingLoc(parentRange.Start);
+  if (sourceMgr.rangeContainsTokenLoc(parentRange, childLoc))
+    return true;
+
   auto childBuffer = sourceMgr.findBufferContainingLoc(childLoc);
-  while (parentBuffer != childBuffer) {
-    auto info = sourceMgr.getGeneratedSourceInfo(childBuffer);
+  while (!sourceMgr.rangeContainsTokenLoc(
+      sourceMgr.getRangeForBuffer(childBuffer), parentRange.Start)) {
+    auto *info = sourceMgr.getGeneratedSourceInfo(childBuffer);
     if (!info)
       return false;