1 files changed, 24 insertions, 21 deletions

diff --git a/clang/lib/Analysis/ThreadSafety.cpp b/clang/lib/Analysis/ThreadSafety.cpp
index cee98d5..d19f86a 100644
--- a/clang/lib/Analysis/ThreadSafety.cpp
+++ b/clang/lib/Analysis/ThreadSafety.cpp
@@ -1668,13 +1668,13 @@ void ThreadSafetyAnalyzer::getEdgeLockset(FactSet& Result,
   const CFGBlockInfo *PredBlockInfo = &BlockInfo[PredBlock->getBlockID()];
   const LocalVarContext &LVarCtx = PredBlockInfo->ExitContext;
 
-  // Temporarily set the lookup context for SExprBuilder.
-  SxBuilder.setLookupLocalVarExpr([&](const NamedDecl *D) -> const Expr * {
-    if (!Handler.issueBetaWarnings())
-      return nullptr;
-    auto Ctx = LVarCtx;
-    return LocalVarMap.lookupExpr(D, Ctx);
-  });
+  if (Handler.issueBetaWarnings()) {
+    // Temporarily set the lookup context for SExprBuilder.
+    SxBuilder.setLookupLocalVarExpr(
+        [this, Ctx = LVarCtx](const NamedDecl *D) mutable -> const Expr * {
+          return LocalVarMap.lookupExpr(D, Ctx);
+        });
+  }
   auto Cleanup = llvm::make_scope_exit(
       [this] { SxBuilder.setLookupLocalVarExpr(nullptr); });
 
@@ -1722,6 +1722,19 @@ class BuildLockset : public ConstStmtVisitor<BuildLockset> {
   LocalVariableMap::Context LVarCtx;
   unsigned CtxIndex;
 
+  // To update and adjust the context.
+  void updateLocalVarMapCtx(const Stmt *S) {
+    if (S)
+      LVarCtx = Analyzer->LocalVarMap.getNextContext(CtxIndex, S, LVarCtx);
+    if (!Analyzer->Handler.issueBetaWarnings())
+      return;
+    // The lookup closure needs to be reconstructed with the refreshed LVarCtx.
+    Analyzer->SxBuilder.setLookupLocalVarExpr(
+        [this, Ctx = LVarCtx](const NamedDecl *D) mutable -> const Expr * {
+          return Analyzer->LocalVarMap.lookupExpr(D, Ctx);
+        });
+  }
+
   // helper functions
 
   void checkAccess(const Expr *Exp, AccessKind AK,
@@ -1747,13 +1760,7 @@ public:
       : ConstStmtVisitor<BuildLockset>(), Analyzer(Anlzr), FSet(Info.EntrySet),
         FunctionExitFSet(FunctionExitFSet), LVarCtx(Info.EntryContext),
         CtxIndex(Info.EntryIndex) {
-    Analyzer->SxBuilder.setLookupLocalVarExpr(
-        [this](const NamedDecl *D) -> const Expr * {
-          if (!Analyzer->Handler.issueBetaWarnings())
-            return nullptr;
-          auto Ctx = LVarCtx;
-          return Analyzer->LocalVarMap.lookupExpr(D, Ctx);
-        });
+    updateLocalVarMapCtx(nullptr);
   }
 
   ~BuildLockset() { Analyzer->SxBuilder.setLookupLocalVarExpr(nullptr); }
@@ -2259,9 +2266,7 @@ void BuildLockset::VisitBinaryOperator(const BinaryOperator *BO) {
   if (!BO->isAssignmentOp())
     return;
 
-  // adjust the context
-  LVarCtx = Analyzer->LocalVarMap.getNextContext(CtxIndex, BO, LVarCtx);
-
+  updateLocalVarMapCtx(BO);
   checkAccess(BO->getLHS(), AK_Written);
 }
 
@@ -2307,8 +2312,7 @@ void BuildLockset::examineArguments(const FunctionDecl *FD,
 }
 
 void BuildLockset::VisitCallExpr(const CallExpr *Exp) {
-  // adjust the context
-  LVarCtx = Analyzer->LocalVarMap.getNextContext(CtxIndex, Exp, LVarCtx);
+  updateLocalVarMapCtx(Exp);
 
   if (const auto *CE = dyn_cast<CXXMemberCallExpr>(Exp)) {
     const auto *ME = dyn_cast<MemberExpr>(CE->getCallee());
@@ -2404,8 +2408,7 @@ static const Expr *UnpackConstruction(const Expr *E) {
 }
 
 void BuildLockset::VisitDeclStmt(const DeclStmt *S) {
-  // adjust the context
-  LVarCtx = Analyzer->LocalVarMap.getNextContext(CtxIndex, S, LVarCtx);
+  updateLocalVarMapCtx(S);
 
   for (auto *D : S->getDeclGroup()) {
     if (auto *VD = dyn_cast_or_null<VarDecl>(D)) {