[clang][dataflow] Handle return statements

[samestep]

Summary

This patch adds a ReturnLoc field to the Environment, serving a similar to the ThisPointeeLoc field in the DataflowAnalysisContext. It then uses that (along with a new VisitReturnStmt method in TransferVisitor) to handle non-void-returning functions in context-sensitive analysis.

[gribozavr2]

clang/include/clang/Analysis/FlowSensitive/DataflowAnalysisContext.h

158

168–169

clang/include/clang/Analysis/FlowSensitive/DataflowEnvironment.h

225–226

clang/lib/Analysis/FlowSensitive/Transfer.cpp

347

Please add a TODO about modeling NRVO.

[samestep]

Address Dmitri's review

clang/include/clang/Analysis/FlowSensitive/DataflowAnalysisContext.h

158

Done.

168–169

Done.

clang/include/clang/Analysis/FlowSensitive/DataflowEnvironment.h

225–226

Done.

clang/lib/Analysis/FlowSensitive/Transfer.cpp

347

Added, thanks! I hadn't realized this would be necessary.

[sgatev]

clang/lib/Analysis/FlowSensitive/Transfer.cpp

338–339

Let's make this a FIXME to set a storage location for the outermost context too.

[samestep]

clang/lib/Analysis/FlowSensitive/Transfer.cpp

338–339

@sgatev I could add a FIXME for that, or I could just do it in this same patch; do you have a preference between those two options?

[samestep]

Set a storage location for the outermost context too

[li.zhe.hua]

clang/include/clang/Analysis/FlowSensitive/DataflowAnalysisContext.h

162

Fix this as well? A reader shouldn't need to root around in private implementation details to understand the requirements for calling a function.

clang/lib/Analysis/FlowSensitive/DataflowEnvironment.cpp

207–208

I'm a little unclear on what "this" is. Is it this entire function, or just the call to getDirectCallee()? Can this comment be moved somewhere more appropriate, and specifically, so it is touching the code that is most relevant to it? It is currently floating in the middle of the function, and it's unclear to me why new code is being added above it vs. below it.

clang/unittests/Analysis/FlowSensitive/TransferTest.cpp

3908

Why is this change to the test necessary?

[samestep]

clang/include/clang/Analysis/FlowSensitive/DataflowAnalysisContext.h

162

Could you clarify what you mean? I was simply copying the signature and docstring from setThisPointeeStorageLocation.

clang/lib/Analysis/FlowSensitive/DataflowEnvironment.cpp

207–208

Yes, this comment is not meant to be present in this patch; I am trying to update the parent patch to remove this comment, but am currently unable to do export it to Phabricator for technical reasons. I should be able to do that tomorrow.

clang/unittests/Analysis/FlowSensitive/TransferTest.cpp

3908

This is mentioned in the patch description (updated earlier today); it's not necessary, but I added it to get a bit of extra coverage for some cases in the VisitReturnStmt method this patch adds.

[li.zhe.hua]

clang/include/clang/Analysis/FlowSensitive/DataflowAnalysisContext.h

162

You've marked return in backticks. There is no parameter named return and it is unclear what return refers to. My best guess is that this is a typo of ReturnLoc, which is a private data member. So this is a public interface with a requirement that a private data member has some property. This should instead reframe the requirement as properties from the external reader's perspective.

clang/unittests/Analysis/FlowSensitive/TransferTest.cpp

3908

Please add the coverage as a separate test. Separate behaviors should be tested as separate tests. go/unit-testing-practices#behavior-testing

[sgatev]

clang/include/clang/Analysis/FlowSensitive/DataflowAnalysisContext.h

162

That was my guess initially too, but my next best guess is that return in backticks stands for the keyword/AST node. In any case, let's make it less ambiguous and let's not add requirements based on implementation details. How about: The return value must not be assigned a storage location.?

clang/lib/Analysis/FlowSensitive/Transfer.cpp

338–339

Same patch works!

[samestep]

clang/include/clang/Analysis/FlowSensitive/DataflowAnalysisContext.h

162

Ah sorry, I understand now; you simply meant that I should make the same change here that @gribozavr2 suggested I make to the other docstrings? I'll go ahead and do that (once I am able to re-export this patch again).

clang/lib/Analysis/FlowSensitive/Transfer.cpp

338–339

Cool, I have an update to this patch which does that, so I'll export it here as soon as I am able to do so.

clang/unittests/Analysis/FlowSensitive/TransferTest.cpp

3908

Makes sense, I'll do that. Thanks!

[samestep]

Rebase and address all comments

[ymandel]

clang/include/clang/Analysis/FlowSensitive/DataflowAnalysisContext.h

158

This looks optional (since it is a pointer). If so, please comment to explain.

clang/include/clang/Analysis/FlowSensitive/DataflowEnvironment.h

225–226

This implies a timing/initialization aspect to its state. Is this necessary? If not, can we restructure to avoid it? Alternatively, if it is valid to have it unset, please update the comment to reflect.

[ymandel]

Please ignore my previous comments from an earlier revision.

[samestep]

Rebase

[sgatev]

clang/lib/Analysis/FlowSensitive/Transfer.cpp

345

Let's do if (Val->getKind() == Value::Kind::Reference) return;. Otherwise the framework will be unusable in practice.

348

Let's do if (Loc == nullptr) return;

564

We use stable storage locations to ensure convergence. In that spirit, shouldn't we assign ReturnLoc's value to S's storage location instead of changing the storage location? Alternatively, we can pass S's storage location to pushCall so that it can store it as ReturnLoc.

[samestep]

clang/lib/Analysis/FlowSensitive/Transfer.cpp

345

Makes sense, will do.

348

I don't think we want to do that, right? Shouldn't the return storage location always be set? Or is this about the "analyzing fragments rather than full functions" thing we discussed yesterday?

564

Could you clarify how this hurts convergence? My understanding is that ReturnLoc here is already stable, so this would make S's storage location stable too.

[sgatev]

clang/lib/Analysis/FlowSensitive/Transfer.cpp

348

I think it's related. If we are going with always initializing the return storage location then I guess at some point we should be able to make Environment::getReturnStorageLocation return a reference? In that case I'm fine with keeping the assert around in the meantime.

564

If I follow correctly, ReturnLoc here is the result of Env.createStorageLocation(ReturnType) which isn't stable. Each call to createStorageLocation returns a fresh storage location.

[samestep]

clang/lib/Analysis/FlowSensitive/Transfer.cpp

348

OK; yeah, I think the intention is that we're always initializing it. I'll leave this code as is for now, then.

564

Ah I see, you're right. Is there a way to make a stable storage location for the return? My intuition is that we can't just pass S's storage location to pushCall, because we want the storage location for the return to be the same across analysis of different callsites to the callee (similar to how we currently use the same storage location for a given parameter of the callee, regardless of how many times we analyze it). But maybe it would be fine; @ymandel do you have any thoughts on this?

[samestep]

Don't assert returns to not be references

[samestep]

Use CallExpr location for ReturnLoc

[samestep]

Revert an unnecessary little change

[sgatev]

clang/lib/Analysis/FlowSensitive/DataflowEnvironment.cpp

208

Let's add a FIXME to support references here.

clang/lib/Analysis/FlowSensitive/Transfer.cpp

542–544

Now there's a hidden connection - ReturnLoc gets assigned a value in Env because we implicitly use the same storage location in Env.pushCall(S). I suggest adding a comment about this.

[samestep]

Add Stanislav's suggested comments

clang/include/clang/Analysis/FlowSensitive/DataflowAnalysisContext.h

public:
   }
 
   /// Returns the storage location assigned to the `this` pointee or null if the
   /// `this` pointee has no assigned storage location.
   StorageLocation *getThisPointeeStorageLocation() const {
     return ThisPointeeLoc;
   }
 
   /// Returns a pointer value that represents a null pointer. Calls with
   /// `PointeeType` that are canonically equivalent will return the same result.
   /// A null `PointeeType` can be used for the pointee of `std::nullptr_t`.
   PointerValue &getOrCreateNullPointerValue(QualType PointeeType);
 
   /// Returns a symbolic boolean value that models a boolean literal equal to
   /// `Value`.
   AtomicBoolValue &getBoolLiteralValue(bool Value) const {
     return Value ? TrueVal : FalseVal;
   }
 
   /// Creates an atomic boolean value.
   AtomicBoolValue &createAtomicBoolValue() {
     return takeOwnership(std::make_unique<AtomicBoolValue>());
   }
 
   /// Returns a boolean value that represents the conjunction of `LHS` and
   /// `RHS`. Subsequent calls with the same arguments, regardless of their
   /// order, will return the same result. If the given boolean values represent
   /// the same value, the result will be the value itself.
private:
   // Maps from program declarations and statements to storage locations that are
   // assigned to them. These assignments are global (aggregated across all basic
   // blocks) and are used to produce stable storage locations when the same
   // basic blocks are evaluated multiple times. The storage locations that are
   // in scope for a particular basic block are stored in `Environment`.
   llvm::DenseMap<const ValueDecl *, StorageLocation *> DeclToLoc;
   llvm::DenseMap<const Expr *, StorageLocation *> ExprToLoc;
 
+  // FIXME: Move this into `Environment`.
   StorageLocation *ThisPointeeLoc = nullptr;
 
   // Null pointer values, keyed by the canonical pointee type.
   //
   // FIXME: The pointer values are indexed by the pointee types which are
   // required to initialize the `PointeeLoc` field in `PointerValue`. Consider
   // creating a type-independent `NullPointerValue` without a `PointeeLoc`
   // field.

clang/include/clang/Analysis/FlowSensitive/DataflowEnvironment.h

public:
   /// assigned a storage location in the environment.
   StorageLocation *getStorageLocation(const Expr &E, SkipPast SP) const;
 
   /// Returns the storage location assigned to the `this` pointee in the
   /// environment or null if the `this` pointee has no assigned storage location
   /// in the environment.
   StorageLocation *getThisPointeeStorageLocation() const;
 
+  /// Returns the storage location of the return value or null, if unset.
+  StorageLocation *getReturnStorageLocation() const;
+
   /// Returns a pointer value that represents a null pointer. Calls with
   /// `PointeeType` that are canonically equivalent will return the same result.
   PointerValue &getOrCreateNullPointerValue(QualType PointeeType);
 
   /// Creates a value appropriate for `Type`, if `Type` is supported, otherwise
   /// return null. If `Type` is a pointer or reference type, creates all the
   /// necessary storage locations and values for indirections until it finds a
   /// non-pointer/non-reference type.
Value *createValueUnlessSelfReferential(QualType Type,
                                           int Depth, int &CreatedValuesCount);
 
   StorageLocation &skip(StorageLocation &Loc, SkipPast SP) const;
   const StorageLocation &skip(const StorageLocation &Loc, SkipPast SP) const;
 
   // `DACtx` is not null and not owned by this object.
   DataflowAnalysisContext *DACtx;
 
+  // In a properly initialized `Environment`, `ReturnLoc` should only be null if
+  // its `DeclContext` could not be cast to a `FunctionDecl`.
+  StorageLocation *ReturnLoc = nullptr;
+  // FIXME: Move `ThisPointeeLoc` here from `DataflowAnalysisContext`.
+
   // Maps from program declarations and statements to storage locations that are
   // assigned to them. Unlike the maps in `DataflowAnalysisContext`, these
   // include only storage locations that are in scope for a particular basic
   // block.
   llvm::DenseMap<const ValueDecl *, StorageLocation *> DeclToLoc;
   llvm::DenseMap<const Expr *, StorageLocation *> ExprToLoc;
 
   llvm::DenseMap<const StorageLocation *, Value *> LocToVal;

clang/lib/Analysis/FlowSensitive/DataflowEnvironment.cpp

if (auto *DS = dyn_cast<DeclStmt>(&S)) {
     initGlobalVar(*E->getMemberDecl(), Env);
   }
 }
 
 Environment::Environment(DataflowAnalysisContext &DACtx)
     : DACtx(&DACtx), FlowConditionToken(&DACtx.makeFlowConditionToken()) {}
 
 Environment::Environment(const Environment &Other)
-    : DACtx(Other.DACtx), DeclToLoc(Other.DeclToLoc),
-      ExprToLoc(Other.ExprToLoc), LocToVal(Other.LocToVal),
-      MemberLocToStruct(Other.MemberLocToStruct),
+    : DACtx(Other.DACtx), ReturnLoc(Other.ReturnLoc),
+      DeclToLoc(Other.DeclToLoc), ExprToLoc(Other.ExprToLoc),
+      LocToVal(Other.LocToVal), MemberLocToStruct(Other.MemberLocToStruct),
       FlowConditionToken(&DACtx->forkFlowCondition(*Other.FlowConditionToken)) {
 }
 
 Environment &Environment::operator=(const Environment &Other) {
   Environment Copy(Other);
   *this = std::move(Copy);
   return *this;
 }
 
 Environment::Environment(DataflowAnalysisContext &DACtx,
                          const DeclContext &DeclCtx)
     : Environment(DACtx) {
   if (const auto *FuncDecl = dyn_cast<FunctionDecl>(&DeclCtx)) {
     assert(FuncDecl->getBody() != nullptr);
     initGlobalVars(*FuncDecl->getBody(), *this);
     for (const auto *ParamDecl : FuncDecl->parameters()) {
       assert(ParamDecl != nullptr);
       auto &ParamLoc = createStorageLocation(*ParamDecl);
       setStorageLocation(*ParamDecl, ParamLoc);
       if (Value *ParamVal = createValue(ParamDecl->getType()))
         setValue(ParamLoc, *ParamVal);
     }
+
+    QualType ReturnType = FuncDecl->getReturnType();
+    ReturnLoc = &createStorageLocation(ReturnType);
   }
 
   if (const auto *MethodDecl = dyn_cast<CXXMethodDecl>(&DeclCtx)) {
     auto *Parent = MethodDecl->getParent();
     assert(Parent != nullptr);
     if (Parent->isLambda())
       MethodDecl = dyn_cast<CXXMethodDecl>(Parent->getDeclContext());
 
     if (MethodDecl && !MethodDecl->isStatic()) {
       QualType ThisPointeeType = MethodDecl->getThisObjectType();
       // FIXME: Add support for union types.
       if (!ThisPointeeType->isUnionType()) {
         auto &ThisPointeeLoc = createStorageLocation(ThisPointeeType);
         DACtx.setThisPointeeStorageLocation(ThisPointeeLoc);
         if (Value *ThisPointeeVal = createValue(ThisPointeeType))
           setValue(ThisPointeeLoc, *ThisPointeeVal);
       }
     }
   }
 }
 
 Environment Environment::pushCall(const CallExpr *Call) const {
   Environment Env(*this);
 
   const auto *FuncDecl = Call->getDirectCallee();
   assert(FuncDecl != nullptr);
   assert(FuncDecl->getBody() != nullptr);
   // FIXME: In order to allow the callee to reference globals, we probably need
   // to call `initGlobalVars` here in some way.
 
+  QualType ReturnType = FuncDecl->getReturnType();
+  Env.ReturnLoc = &Env.createStorageLocation(ReturnType);
+
   auto ParamIt = FuncDecl->param_begin();
   auto ArgIt = Call->arg_begin();
   auto ArgEnd = Call->arg_end();
 
   // FIXME: Parameters don't always map to arguments 1:1; examples include
   // overloaded operators implemented as member functions, and parameter packs.
   for (; ArgIt != ArgEnd; ++ParamIt, ++ArgIt) {
     assert(ParamIt != FuncDecl->param_end());
if (ParamType->isReferenceType()) {
       Env.setValue(Loc, *Val);
     }
   }
 
   return Env;
 }
 
 void Environment::popCall(const Environment &CalleeEnv) {
-  // We ignore `DACtx` because it's already the same in both. We don't bring
-  // back `DeclToLoc` and `ExprToLoc` because we want to be able to later
-  // analyze the same callee in a different context, and `setStorageLocation`
-  // requires there to not already be a storage location assigned. Conceptually,
-  // these maps capture information from the local scope, so when popping that
-  // scope, we do not propagate the maps.
+  // We ignore `DACtx` because it's already the same in both. We don't want the
+  // callee's `ReturnLoc`. We don't bring back `DeclToLoc` and `ExprToLoc`
+  // because we want to be able to later analyze the same callee in a different
+  // context, and `setStorageLocation` requires there to not already be a
+  // storage location assigned. Conceptually, these maps capture information
+  // from the local scope, so when popping that scope, we do not propagate the
+  // maps.
   this->LocToVal = std::move(CalleeEnv.LocToVal);
   this->MemberLocToStruct = std::move(CalleeEnv.MemberLocToStruct);
   this->FlowConditionToken = std::move(CalleeEnv.FlowConditionToken);
 }
 
 bool Environment::equivalentTo(const Environment &Other,
                                Environment::ValueModel &Model) const {
   assert(DACtx == Other.DACtx);
 
+  if (ReturnLoc != Other.ReturnLoc)
+    return false;
+
   if (DeclToLoc != Other.DeclToLoc)
     return false;
 
   if (ExprToLoc != Other.ExprToLoc)
     return false;
 
   // Compare the contents for the intersection of their domains.
   for (auto &Entry : LocToVal) {
bool Environment::equivalentTo(const Environment &Other,
   }
 
   return true;
 }
 
 LatticeJoinEffect Environment::join(const Environment &Other,
                                     Environment::ValueModel &Model) {
   assert(DACtx == Other.DACtx);
+  assert(ReturnLoc == Other.ReturnLoc);
 
   auto Effect = LatticeJoinEffect::Unchanged;
 
   Environment JoinedEnv(*DACtx);
 
+  JoinedEnv.ReturnLoc = ReturnLoc;
+
   JoinedEnv.DeclToLoc = intersectDenseMaps(DeclToLoc, Other.DeclToLoc);
   if (DeclToLoc.size() != JoinedEnv.DeclToLoc.size())
     Effect = LatticeJoinEffect::Changed;
 
   JoinedEnv.ExprToLoc = intersectDenseMaps(ExprToLoc, Other.ExprToLoc);
   if (ExprToLoc.size() != JoinedEnv.ExprToLoc.size())
     Effect = LatticeJoinEffect::Changed;
 
StorageLocation *Environment::getStorageLocation(const Expr &E,
   auto It = ExprToLoc.find(&ignoreCFGOmittedNodes(E));
   return It == ExprToLoc.end() ? nullptr : &skip(*It->second, SP);
 }
 
 StorageLocation *Environment::getThisPointeeStorageLocation() const {
   return DACtx->getThisPointeeStorageLocation();
 }
 
+StorageLocation *Environment::getReturnStorageLocation() const {
+  return ReturnLoc;
+}
+
 PointerValue &Environment::getOrCreateNullPointerValue(QualType PointeeType) {
   return DACtx->getOrCreateNullPointerValue(PointeeType);
 }
 
 void Environment::setValue(const StorageLocation &Loc, Value &Val) {
   LocToVal[&Loc] = &Val;
 
   if (auto *StructVal = dyn_cast<StructValue>(&Val)) {

clang/lib/Analysis/FlowSensitive/Transfer.cpp

if (ThisPointeeLoc == nullptr)
       return;
 
     auto &Loc = Env.createStorageLocation(*S);
     Env.setStorageLocation(*S, Loc);
     Env.setValue(Loc, Env.takeOwnership(
                           std::make_unique<PointerValue>(*ThisPointeeLoc)));
   }
 
+  void VisitReturnStmt(const ReturnStmt *S) {
+    auto *Ret = S->getRetValue();
+    if (Ret == nullptr)
+      return;
+
+    auto *Val = Env.getValue(*Ret, SkipPast::None);
+    if (Val == nullptr)
+      return;
+
+    // FIXME: Support reference-type returns.
+    assert(Val->getKind() != Value::Kind::Reference);
+
+    auto *Loc = Env.getReturnStorageLocation();
+    assert(Loc != nullptr);
+    // FIXME: Model NRVO.
+    Env.setValue(*Loc, *Val);
+  }
+
   void VisitMemberExpr(const MemberExpr *S) {
     ValueDecl *Member = S->getMemberDecl();
     assert(Member != nullptr);
 
     // FIXME: Consider assigning pointer values to function member expressions.
     if (Member->isFunctionOrFunctionTemplate())
       return;
 
if (S->isCallToStdMove()) {
         return;
 
       // FIXME: We don't support context-sensitive analysis of recursion, so
       // we should return early here if `F` is the same as the `FunctionDecl`
       // holding `S` itself.
 
       auto ExitBlock = CFCtx->getCFG().getExit().getBlockID();
 
       auto CalleeEnv = Env.pushCall(S);
 
       // FIXME: Use the same analysis as the caller for the callee. Note,
       // though, that doing so would require support for changing the analysis's
       // ASTContext.
       assert(
           CFCtx->getDecl() != nullptr &&
           "ControlFlowContexts in the environment should always carry a decl");
       auto Analysis = NoopAnalysis(CFCtx->getDecl()->getASTContext(),
                                    DataflowAnalysisOptions());
 
       auto BlockToOutputState =
           dataflow::runDataflowAnalysis(*CFCtx, Analysis, CalleeEnv);
       assert(BlockToOutputState);
       assert(ExitBlock < BlockToOutputState->size());
 
       auto ExitState = (*BlockToOutputState)[ExitBlock];
       assert(ExitState);
 
-      Env.popCall(ExitState->Env);
+      auto &ExitEnv = ExitState->Env;
+      auto *ReturnLoc = ExitEnv.getReturnStorageLocation();
+      assert(ReturnLoc != nullptr);
+      Env.setStorageLocation(*S, *ReturnLoc);
+      Env.popCall(ExitEnv);
     }
   }
 
   void VisitMaterializeTemporaryExpr(const MaterializeTemporaryExpr *S) {
     const Expr *SubExpr = S->getSubExpr();
     assert(SubExpr != nullptr);
 
     auto *SubExprLoc = Env.getStorageLocation(*SubExpr, SkipPast::None);

clang/unittests/Analysis/FlowSensitive/TransferTest.cpp

runDataflow(Code,
                /*.BuiltinTransferOptions=*/{/*.ContextSensitive=*/false}});
 }
 
 TEST(TransferTest, ContextSensitiveSetTrue) {
   std::string Code = R"(
     bool GiveBool();
     void SetBool(bool &Var) { Var = true; }
 
     void target() {
       bool Foo = GiveBool();
       SetBool(Foo);
       // [[p]]
     }
   )";
   runDataflow(Code,
               [](llvm::ArrayRef<
                      std::pair<std::string, DataflowAnalysisState<NoopLattice>>>
runDataflow(Code,
                     *cast<BoolValue>(Env.getValue(*BazDecl, SkipPast::None));
                 EXPECT_TRUE(Env.flowConditionImplies(BazVal));
                 EXPECT_FALSE(Env.flowConditionImplies(Env.makeNot(BazVal)));
               },
               {/*.ApplyBuiltinTransfer=*/true,
                /*.BuiltinTransferOptions=*/{/*.ContextSensitive=*/true}});
 }
 
+TEST(TransferTest, ContextSensitiveReturnVoid) {
+  std::string Code = R"(
+    void Noop() { return; }
+
+    void target() {
+      Noop();
+      // [[p]]
+    }
+  )";
+  runDataflow(Code,
+              [](llvm::ArrayRef<
+                     std::pair<std::string, DataflowAnalysisState<NoopLattice>>>
+                     Results,
+                 ASTContext &ASTCtx) {
+                ASSERT_THAT(Results, ElementsAre(Pair("p", _)));
+                // This just tests that the analysis doesn't crash.
+              },
+              {/*.ApplyBuiltinTransfer=*/true,
+               /*.BuiltinTransferOptions=*/{/*.ContextSensitive=*/true}});
+}
+
+TEST(TransferTest, ContextSensitiveReturnTrue) {
+  std::string Code = R"(
+    bool GiveBool() { return true; }
+
+    void target() {
+      bool Foo = GiveBool();
+      // [[p]]
+    }
+  )";
+  runDataflow(Code,
+              [](llvm::ArrayRef<
+                     std::pair<std::string, DataflowAnalysisState<NoopLattice>>>
+                     Results,
+                 ASTContext &ASTCtx) {
+                ASSERT_THAT(Results, ElementsAre(Pair("p", _)));
+                const Environment &Env = Results[0].second.Env;
+
+                const ValueDecl *FooDecl = findValueDecl(ASTCtx, "Foo");
+                ASSERT_THAT(FooDecl, NotNull());
+
+                auto &FooVal =
+                    *cast<BoolValue>(Env.getValue(*FooDecl, SkipPast::None));
+                EXPECT_TRUE(Env.flowConditionImplies(FooVal));
+              },
+              {/*.ApplyBuiltinTransfer=*/true,
+               /*.BuiltinTransferOptions=*/{/*.ContextSensitive=*/true}});
+}
+
+TEST(TransferTest, ContextSensitiveReturnFalse) {
+  std::string Code = R"(
+    bool GiveBool() { return false; }
+
+    void target() {
+      bool Foo = GiveBool();
+      // [[p]]
+    }
+  )";
+  runDataflow(Code,
+              [](llvm::ArrayRef<
+                     std::pair<std::string, DataflowAnalysisState<NoopLattice>>>
+                     Results,
+                 ASTContext &ASTCtx) {
+                ASSERT_THAT(Results, ElementsAre(Pair("p", _)));
+                const Environment &Env = Results[0].second.Env;
+
+                const ValueDecl *FooDecl = findValueDecl(ASTCtx, "Foo");
+                ASSERT_THAT(FooDecl, NotNull());
+
+                auto &FooVal =
+                    *cast<BoolValue>(Env.getValue(*FooDecl, SkipPast::None));
+                EXPECT_TRUE(Env.flowConditionImplies(Env.makeNot(FooVal)));
+              },
+              {/*.ApplyBuiltinTransfer=*/true,
+               /*.BuiltinTransferOptions=*/{/*.ContextSensitive=*/true}});
+}
+
+TEST(TransferTest, ContextSensitiveReturnArg) {
+  std::string Code = R"(
+    bool GiveBool();
+    bool GiveBack(bool Arg) { return Arg; }
+
+    void target() {
+      bool Foo = GiveBool();
+      bool Bar = GiveBack(Foo);
+      bool Baz = Foo == Bar;
+      // [[p]]
+    }
+  )";
+  runDataflow(Code,
+              [](llvm::ArrayRef<
+                     std::pair<std::string, DataflowAnalysisState<NoopLattice>>>
+                     Results,
+                 ASTContext &ASTCtx) {
+                ASSERT_THAT(Results, ElementsAre(Pair("p", _)));
+                const Environment &Env = Results[0].second.Env;
+
+                const ValueDecl *BazDecl = findValueDecl(ASTCtx, "Baz");
+                ASSERT_THAT(BazDecl, NotNull());
+
+                auto &BazVal =
+                    *cast<BoolValue>(Env.getValue(*BazDecl, SkipPast::None));
+                EXPECT_TRUE(Env.flowConditionImplies(BazVal));
+              },
+              {/*.ApplyBuiltinTransfer=*/true,
+               /*.BuiltinTransferOptions=*/{/*.ContextSensitive=*/true}});
+}
+
 } // namespace