[clang] [clang][dataflow] Fix crash when `ConstantExpr` is used in conditional operator. (PR #90112)

[via cfe-commits]

https://github.com/martinboehme created https://github.com/llvm/llvm-project/pull/90112

`ConstantExpr` does not appear as a `CFGStmt` in the CFG, so
`StmtToEnvMap::getEnvironment()` was not finding an entry for it in the map,
causing a crash when we tried to access the iterator resulting from the map
lookup.

The fix is to make `ignoreCFGOmittedNodes()` ignore `ConstantExpr`, but in
addition, I'm hardening `StmtToEnvMap::getEnvironment()` to make sure release
builds don't crash in similar situations in the future.


>From 41917c5465f48022f57412849aa7e1735cc70b17 Mon Sep 17 00:00:00 2001
From: Martin Braenne <mboehme at google.com>
Date: Thu, 25 Apr 2024 19:50:48 +0000
Subject: [PATCH] [clang][dataflow] Fix crash when `ConstantExpr` is used in
 conditional operator.

`ConstantExpr` does not appear as a `CFGStmt` in the CFG, so
`StmtToEnvMap::getEnvironment()` was not finding an entry for it in the map,
causing a crash when we tried to access the iterator resulting from the map
lookup.

The fix is to make `ignoreCFGOmittedNodes()` ignore `ConstantExpr`, but in
addition, I'm hardening `StmtToEnvMap::getEnvironment()` to make sure release
builds don't crash in similar situations in the future.
---
 clang/lib/Analysis/FlowSensitive/ASTOps.cpp   | 16 +++++++---
 clang/lib/Analysis/FlowSensitive/Transfer.cpp |  6 +++-
 .../Analysis/FlowSensitive/TransferTest.cpp   | 32 +++++++++++++++++++
 3 files changed, 49 insertions(+), 5 deletions(-)

diff --git a/clang/lib/Analysis/FlowSensitive/ASTOps.cpp b/clang/lib/Analysis/FlowSensitive/ASTOps.cpp
index 619bf772bba5ee..bd1676583ecccd 100644
--- a/clang/lib/Analysis/FlowSensitive/ASTOps.cpp
+++ b/clang/lib/Analysis/FlowSensitive/ASTOps.cpp
@@ -33,12 +33,20 @@ namespace clang::dataflow {
 
 const Expr &ignoreCFGOmittedNodes(const Expr &E) {
   const Expr *Current = &E;
-  if (auto *EWC = dyn_cast<ExprWithCleanups>(Current)) {
-    Current = EWC->getSubExpr();
+  const Expr *Last = nullptr;
+  while (Current != Last) {
+    Last = Current;
+    if (auto *EWC = dyn_cast<ExprWithCleanups>(Current)) {
+      Current = EWC->getSubExpr();
+      assert(Current != nullptr);
+    }
+    if (auto *CE = dyn_cast<ConstantExpr>(Current)) {
+      Current = CE->getSubExpr();
+      assert(Current != nullptr);
+    }
+    Current = Current->IgnoreParens();
     assert(Current != nullptr);
   }
-  Current = Current->IgnoreParens();
-  assert(Current != nullptr);
   return *Current;
 }
 
diff --git a/clang/lib/Analysis/FlowSensitive/Transfer.cpp b/clang/lib/Analysis/FlowSensitive/Transfer.cpp
index 43fdfa5abcbb51..fd224aeb79b151 100644
--- a/clang/lib/Analysis/FlowSensitive/Transfer.cpp
+++ b/clang/lib/Analysis/FlowSensitive/Transfer.cpp
@@ -41,7 +41,11 @@ namespace dataflow {
 
 const Environment *StmtToEnvMap::getEnvironment(const Stmt &S) const {
   auto BlockIt = ACFG.getStmtToBlock().find(&ignoreCFGOmittedNodes(S));
-  assert(BlockIt != ACFG.getStmtToBlock().end());
+  if (BlockIt == ACFG.getStmtToBlock().end()) {
+    assert(false);
+    // Return null to avoid dereferencing the end iterator in non-assert builds.
+    return nullptr;
+  }
   if (!ACFG.isBlockReachable(*BlockIt->getSecond()))
     return nullptr;
   if (BlockIt->getSecond()->getBlockID() == CurBlockID)
diff --git a/clang/unittests/Analysis/FlowSensitive/TransferTest.cpp b/clang/unittests/Analysis/FlowSensitive/TransferTest.cpp
index d204700919d315..301bec32c0cf1d 100644
--- a/clang/unittests/Analysis/FlowSensitive/TransferTest.cpp
+++ b/clang/unittests/Analysis/FlowSensitive/TransferTest.cpp
@@ -5357,6 +5357,38 @@ TEST(TransferTest, ConditionalOperatorLocation) {
       });
 }
 
+TEST(TransferTest, ConditionalOperatorOnConstantExpr) {
+  // This is a regression test: We used to crash when a `ConstantExpr` was used
+  // in the branches of a conditional operator.
+  std::string Code = R"cc(
+    consteval bool identity(bool B) { return B; }
+    void target(bool Cond) {
+      bool JoinTrueTrue = Cond ? identity(true) : identity(true);
+      bool JoinTrueFalse = Cond ? identity(true) : identity(false);
+      // [[p]]
+    }
+  )cc";
+  runDataflow(
+      Code,
+      [](const llvm::StringMap<DataflowAnalysisState<NoopLattice>> &Results,
+         ASTContext &ASTCtx) {
+        Environment Env = getEnvironmentAtAnnotation(Results, "p").fork();
+
+        auto &JoinTrueTrue =
+            getValueForDecl<BoolValue>(ASTCtx, Env, "JoinTrueTrue");
+        // FIXME: This test documents the current behavior, namely that we
+        // don't actually use the constant result of the `ConstantExpr` and
+        // instead treat it like a normal function call.
+        EXPECT_EQ(JoinTrueTrue.formula().kind(), Formula::Kind::AtomRef);
+        // EXPECT_TRUE(JoinTrueTrue.formula().literal());
+
+        auto &JoinTrueFalse =
+            getValueForDecl<BoolValue>(ASTCtx, Env, "JoinTrueFalse");
+        EXPECT_EQ(JoinTrueFalse.formula().kind(), Formula::Kind::AtomRef);
+      },
+      LangStandard::lang_cxx20);
+}
+
 TEST(TransferTest, IfStmtBranchExtendsFlowCondition) {
   std::string Code = R"(
     void target(bool Foo) {