1 files changed, 69 insertions, 0 deletions
diff --git a/clang/lib/Lex/Preprocessor.cpp b/clang/lib/Lex/Preprocessor.cpp
index fe9adb5..924a7e6 100644
--- a/clang/lib/Lex/Preprocessor.cpp
+++ b/clang/lib/Lex/Preprocessor.cpp
@@ -1462,6 +1462,75 @@ void Preprocessor::emitFinalMacroWarning(const Token &Identifier,
   Diag(*A.FinalAnnotationLoc, diag::note_pp_macro_annotation) << 2;
 }
 
+bool Preprocessor::isSafeBufferOptOut(const SourceManager &SourceMgr,
+                                           const SourceLocation &Loc) const {
+  // Try to find a region in `SafeBufferOptOutMap` where `Loc` is in:
+  auto FirstRegionEndingAfterLoc = llvm::partition_point(
+      SafeBufferOptOutMap,
+      [&SourceMgr,
+       &Loc](const std::pair<SourceLocation, SourceLocation> &Region) {
+        return SourceMgr.isBeforeInTranslationUnit(Region.second, Loc);
+      });
+
+  if (FirstRegionEndingAfterLoc != SafeBufferOptOutMap.end()) {
+    // To test if the start location of the found region precedes `Loc`:
+    return SourceMgr.isBeforeInTranslationUnit(FirstRegionEndingAfterLoc->first,
+                                               Loc);
+  }
+  // If we do not find a region whose end location passes `Loc`, we want to
+  // check if the current region is still open:
+  if (!SafeBufferOptOutMap.empty() &&
+      SafeBufferOptOutMap.back().first == SafeBufferOptOutMap.back().second)
+    return SourceMgr.isBeforeInTranslationUnit(SafeBufferOptOutMap.back().first,
+                                               Loc);
+  return false;
+}
+
+bool Preprocessor::enterOrExitSafeBufferOptOutRegion(
+    bool isEnter, const SourceLocation &Loc) {
+  if (isEnter) {
+    if (isPPInSafeBufferOptOutRegion())
+      return true; // invalid enter action
+    InSafeBufferOptOutRegion = true;
+    CurrentSafeBufferOptOutStart = Loc;
+
+    // To set the start location of a new region:
+
+    if (!SafeBufferOptOutMap.empty()) {
+      auto *PrevRegion = &SafeBufferOptOutMap.back();
+      assert(PrevRegion->first != PrevRegion->second &&
+             "Shall not begin a safe buffer opt-out region before closing the "
+             "previous one.");
+    }
+    // If the start location equals to the end location, we call the region a
+    // open region or a unclosed region (i.e., end location has not been set
+    // yet).
+    SafeBufferOptOutMap.emplace_back(Loc, Loc);
+  } else {
+    if (!isPPInSafeBufferOptOutRegion())
+      return true; // invalid enter action
+    InSafeBufferOptOutRegion = false;
+
+    // To set the end location of the current open region:
+
+    assert(!SafeBufferOptOutMap.empty() &&
+           "Misordered safe buffer opt-out regions");
+    auto *CurrRegion = &SafeBufferOptOutMap.back();
+    assert(CurrRegion->first == CurrRegion->second &&
+           "Set end location to a closed safe buffer opt-out region");
+    CurrRegion->second = Loc;
+  }
+  return false;
+}
+
+bool Preprocessor::isPPInSafeBufferOptOutRegion() {
+  return InSafeBufferOptOutRegion;
+}
+bool Preprocessor::isPPInSafeBufferOptOutRegion(SourceLocation &StartLoc) {
+  StartLoc = CurrentSafeBufferOptOutStart;
+  return InSafeBufferOptOutRegion;
+}
+
 ModuleLoader::~ModuleLoader() = default;
 
 CommentHandler::~CommentHandler() = default;