[clang][dataflow] Parameterize analysis by explicit map of analyzable functions.

[ymandel]

Summary

This patch evolves the cache of ControlFlowContexts into a parameter that the
user can supply. Specifically, the system now limits inline analysis of function
calls specifically to those functions contained in a user-supplied map from
"fully qualified name" to ControlFlowContext.

To avoid the cost of generating the fully qualified function name for every
CallExpr, we cache the results of the map lookup in an intermediary map over
FunctionDecl*s.

Additionally, this patch adds a utility function to construct such a map from an ASTUnit. This allows either admitting all the functions in the main AST (for example, a test code snippet) or from a separate AST. This latter functionality is a step towards supporting models from independent ASTs.

Part of issue #56879

[ymandel]

add comment

[sgatev]

clang/include/clang/Analysis/FlowSensitive/ControlFlowContext.h

84

#include "llvm/ADT/StringMap.h"

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

27

This should be StringMap.

60

How about InlineableFunctions?

60

clang/lib/Analysis/FlowSensitive/Transfer.cpp

549–555

How is that different from F? Why not let Environment::pushCall get this from the CallExpr argument?

clang/unittests/Analysis/FlowSensitive/TestingSupport.h

74

[ymandel]

Address reviewer comments.

[xazax.hun]

The concept of fully qualified name is somewhat ambiguous for me. Do we expect the user to specify inline namespaces as well? I'd love to see some test cases and comments that clarify this.

clang/lib/Analysis/FlowSensitive/ControlFlowContext.cpp

96

Do we exclude non-toplevel declarations on purpuse? Or would this work for methods of inline classes, methods of classes defined within a function?

clang/lib/Analysis/FlowSensitive/DataflowAnalysisContext.cpp

341

Out of curiosity, I'd expect getCanonicalDecl to always succeed. Do you know cases where it would not?

[ymandel]

In D131280#3706836, @xazax.hun wrote:

The concept of fully qualified name is somewhat ambiguous for me. Do we expect the user to specify inline namespaces as well? I'd love to see some test cases and comments that clarify this.

Sure. This is probably worth some discussion. Fully qualified names, however we define them, will not be enough, since they don't cover overload sets. So, regardless, we need some internally-consistent (in clang) mechanism of uniquely identifying function declarations. Since we're also supporting methods and constructors, that implies IDs for classes as well.

I'd like some mechanism that matches how identifiers are used. So, for example, inline namespaces should *not* be necessary, since they are an implementation detail from this perspective.

clang/lib/Analysis/FlowSensitive/ControlFlowContext.cpp

96

Do we exclude non-toplevel declarations on purpuse? Or would this work for methods of inline classes, methods of classes defined within a function?

I think that for the current use case -- models of library types and their methods/functions -- we don't have a good usecase for this. But, I can see this becoming an issue if we want to expand to inlining other declarations. So, I'm inclined to hold off on this for the time being, since that's a larger design discussion.

Yet, I also think this should be generalized to take any decl and extract the functions/methods. I limited to ASTUnit for convenience, since that was the immediate need. I'm happy to either:

1. Add a FIXME, and/or,
2. Split this function into two: one that takes two decl iterators (begin, end) and does this work here and another which is just a convenience function for ASTUnit to apply the above to the top-level decls.

WDYT?

clang/lib/Analysis/FlowSensitive/DataflowAnalysisContext.cpp

341

Hah. No, I don't but I also didn't find any mention in the headers guaranteeing otherwise, so I added this. I'm fine removing but I feel a little uncomfortable with the fact that it's not documented.

[xazax.hun]

In D131280#3706915, @ymandel wrote:

Sure. This is probably worth some discussion. Fully qualified names, however we define them, will not be enough, since they don't cover overload sets.

I see. This is not a unique problem. I think there were multiple discussions about API Notes and those need to solve the same problem. @gribozavr2 probably has more context on the current status of API Notes. An alternative to fully qualified names is Clang's USR that is often used for cross-referencing functions across translation units. Less user friendly, but will support overloads.

I'd like some mechanism that matches how identifiers are used. So, for example, inline namespaces should *not* be necessary, since they are an implementation detail from this perspective.

A similar matching is already implemented for the Clang Static Analyzer. See https://clang.llvm.org/doxygen/classclang_1_1ento_1_1CallDescription.html and https://clang.llvm.org/doxygen/classclang_1_1ento_1_1CallDescriptionMap.html

One example use is in the CStringChecker: https://github.com/llvm/llvm-project/blob/main/clang/lib/StaticAnalyzer/Checkers/CStringChecker.cpp#L136

clang/lib/Analysis/FlowSensitive/ControlFlowContext.cpp

96

I am fine with the current behavior, but I think the docstring "Builds a map of all declared functions in the given AST" is misleading in this case. Specifying in the docstring that this function will only map the top-level functions and methods of top-level classes sounds good to me.

clang/lib/Analysis/FlowSensitive/DataflowAnalysisContext.cpp

341

I see other places where we don't check for nullness, e.g.: https://github.com/llvm/llvm-project/blob/bb297024fad2f6c3ccaaa6a5c3a270f73f15f3ac/clang-tools-extra/clang-tidy/misc/UnusedParametersCheck.cpp#L83

I think the CanonicalDecl is defined as the first declaration for redeclarable entities. I think getting the first element of a non-empty list should always succeed. But I definitely agree, this fact should be documented, or even better, this should probably return a reference.

[ymandel]

respond to comments.

[ymandel]

fix typo

[ymandel]

In D131280#3706988, @xazax.hun wrote:

In D131280#3706915, @ymandel wrote:

Sure. This is probably worth some discussion. Fully qualified names, however we define them, will not be enough, since they don't cover overload sets.

I see. This is not a unique problem. I think there were multiple discussions about API Notes and those need to solve the same problem. @gribozavr2 probably has more context on the current status of API Notes. An alternative to fully qualified names is Clang's USR that is often used for cross-referencing functions across translation units. Less user friendly, but will support overloads.

I'd like some mechanism that matches how identifiers are used. So, for example, inline namespaces should *not* be necessary, since they are an implementation detail from this perspective.

A similar matching is already implemented for the Clang Static Analyzer. See https://clang.llvm.org/doxygen/classclang_1_1ento_1_1CallDescription.html and https://clang.llvm.org/doxygen/classclang_1_1ento_1_1CallDescriptionMap.html

One example use is in the CStringChecker: https://github.com/llvm/llvm-project/blob/main/clang/lib/StaticAnalyzer/Checkers/CStringChecker.cpp#L136

Thanks. That looks good, but I'm concerned that it only counts the arguments and doesn't look at their types. I'd imagine this will be a limitation down the line when we want to deal with overload sets w/ the same number of arguments, but different types.

Aside: why the const char * interface? Do you think owners would be open to a llvm::StringRef overload for the constructor?

clang/lib/Analysis/FlowSensitive/ControlFlowContext.cpp

96

Thanks, done.

clang/lib/Analysis/FlowSensitive/Transfer.cpp

549–555

Added comment to explain.

[xazax.hun]

In D131280#3709781, @ymandel wrote:

Thanks. That looks good, but I'm concerned that it only counts the arguments and doesn't look at their types. I'd imagine this will be a limitation down the line when we want to deal with overload sets w/ the same number of arguments, but different types.

Yeah, it is not a fully baked solution at this point, but it does implement some of the features that you plan to add (like skipping inline namespaces), and some more (argument count, checking if a function is from a system header).

Aside: why the const char * interface? Do you think owners would be open to a llvm::StringRef overload for the constructor?

I am sure that the analyzer community is open to any improvements. The main reason I'd be glad if that facility could be shared across the two static analysis solution because improvements from one community could be benefited by the other, also the code would be more similar which could be great for cross pollination of ideas.

If you think it is feasible to reuse some of those facilities for your purposes, I am happy to review all of those patches. If it is not a good fit for some reason, I am ok with having a new custom solution here.

[ymandel]

In D131280#3709964, @xazax.hun wrote:

In D131280#3709781, @ymandel wrote:

Thanks. That looks good, but I'm concerned that it only counts the arguments and doesn't look at their types. I'd imagine this will be a limitation down the line when we want to deal with overload sets w/ the same number of arguments, but different types.

Yeah, it is not a fully baked solution at this point, but it does implement some of the features that you plan to add (like skipping inline namespaces), and some more (argument count, checking if a function is from a system header).

Aside: why the const char * interface? Do you think owners would be open to a llvm::StringRef overload for the constructor?

I am sure that the analyzer community is open to any improvements. The main reason I'd be glad if that facility could be shared across the two static analysis solution because improvements from one community could be benefited by the other, also the code would be more similar which could be great for cross pollination of ideas.

If you think it is feasible to reuse some of those facilities for your purposes, I am happy to review all of those patches. If it is not a good fit for some reason, I am ok with having a new custom solution here.

I think we should try to reuse for the reasons you mention. We can evolve it with more features as needed. I'm kind of allergic to const char * so I propose that I first send a patch for that and can then follow up here.

[xazax.hun]

In D131280#3710080, @ymandel wrote:

I think we should try to reuse for the reasons you mention. We can evolve it with more features as needed. I'm kind of allergic to const char * so I propose that I first send a patch for that and can then follow up here.

This sounds wonderful, thanks!

[samestep]

clang/include/clang/Analysis/FlowSensitive/ControlFlowContext.h

71–74

clang/include/clang/Analysis/FlowSensitive/ControlFlowContext.h

 
 #ifndef LLVM_CLANG_ANALYSIS_FLOWSENSITIVE_CONTROLFLOWCONTEXT_H
 #define LLVM_CLANG_ANALYSIS_FLOWSENSITIVE_CONTROLFLOWCONTEXT_H
 
 #include "clang/AST/ASTContext.h"
 #include "clang/AST/Decl.h"
 #include "clang/AST/Stmt.h"
 #include "clang/Analysis/CFG.h"
+#include "clang/Frontend/ASTUnit.h"
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/StringMap.h"
 #include "llvm/Support/Error.h"
 #include <memory>
 #include <utility>
 
 namespace clang {
 namespace dataflow {
 
 /// Holds CFG and other derived context that is needed to perform dataflow
 /// analysis.
 class ControlFlowContext {
 public:
   /// Builds a ControlFlowContext from an AST node. `D` is the function in which
   /// `S` resides and must not be null.
   static llvm::Expected<ControlFlowContext> build(const Decl *D, Stmt &S,
                                                   ASTContext &C);
 
-  // DEPRECATED. Use overload above.
+  // FIXME: fix known callers and uncomment the annotation.
+  // [[deprecated("Use overload above.")]]
   static llvm::Expected<ControlFlowContext> build(const Decl *D, Stmt *S,
                                                   ASTContext *C);
 
   /// Returns the `Decl` containing the statement used to construct the CFG, if
   /// available.
   const Decl *getDecl() const { return ContainingDecl; }
 
   /// Returns the CFG that is stored in this context.
ControlFlowContext(const Decl *D, std::unique_ptr<CFG> Cfg,
         StmtToBlock(std::move(StmtToBlock)) {}
 
   /// The `Decl` containing the statement used to construct the CFG.
   const Decl *ContainingDecl;
   std::unique_ptr<CFG> Cfg;
   llvm::DenseMap<const Stmt *, const CFGBlock *> StmtToBlock;
 };
 
+/// Builds a map of top-level functions and methods of top-level classes sounds
+/// in the given AST, mapping from fully qualified function name to the
+/// `ControlFlowContext` that corresponds to that function. Fully qualified
+/// names are used so that we can related same-named functions from two
+/// different ASTs. This is critical for providing implementations of functions
+/// that are suitable for inlining during analysis of their callsites.
+///
+/// N.B. The return value, if not empty, will point into `Unit`, so `Unit` must
+/// outlive the return value.
+///
+/// FIXME: add support for constructors and declarations nested in namespaces.
+/// FIXME: fully qualified names are not enough to uniquely identify members of
+/// overload sets. change the key to support uniqueness.
+llvm::Expected<llvm::StringMap<ControlFlowContext>>
+buildFunctionMapFromAST(ASTUnit &Unit);
+
 } // namespace dataflow
 } // namespace clang
 
 #endif // LLVM_CLANG_ANALYSIS_FLOWSENSITIVE_CONTROLFLOWCONTEXT_H

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

 #include "clang/AST/Expr.h"
 #include "clang/AST/TypeOrdering.h"
 #include "clang/Analysis/FlowSensitive/ControlFlowContext.h"
 #include "clang/Analysis/FlowSensitive/Solver.h"
 #include "clang/Analysis/FlowSensitive/StorageLocation.h"
 #include "clang/Analysis/FlowSensitive/Value.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/DenseSet.h"
+#include "llvm/ADT/StringMap.h"
 #include "llvm/Support/Compiler.h"
 #include <cassert>
 #include <memory>
 #include <type_traits>
 #include <utility>
 #include <vector>
 
 namespace clang {

 llvm::DenseSet<const FieldDecl *> getObjectFields(QualType Type);
 
 /// Owns objects that encompass the state of a program and stores context that
 /// is used during dataflow analysis.
 class DataflowAnalysisContext {
 public:
   /// Constructs a dataflow analysis context.
   ///
+  /// `InlineableFunctions` is a map of inlineable function bodies (represented
+  /// as `ControlFlowContext`s), keyed by the fully-qualified function name. The
+  /// analysis will inline calls -- analyzed the function body as if it were
+  /// inlined at the call site -- *only* to the listed functions and will use
+  /// the provided implementation for that purpose. The actual function
+  /// definition, even if available in the current TU, is ignored.
+  ///
   /// Requirements:
   ///
   ///  `S` must not be null.
-  DataflowAnalysisContext(std::unique_ptr<Solver> S)
+  DataflowAnalysisContext(
+      std::unique_ptr<Solver> S,
+      llvm::StringMap<ControlFlowContext> InlineableFunctions = {})
       : S(std::move(S)), TrueVal(createAtomicBoolValue()),
-        FalseVal(createAtomicBoolValue()) {
+        FalseVal(createAtomicBoolValue()),
+        InlineableFunctions(std::move(InlineableFunctions)) {
     assert(this->S != nullptr);
   }
 
   /// Takes ownership of `Loc` and returns a reference to it.
   ///
   /// Requirements:
   ///
   ///  `Loc` must not be null.
private:
   //
   // Flow conditions depend on other flow conditions if they are created using
   // `forkFlowCondition` or `joinFlowConditions`. The graph of flow condition
   // dependencies is stored in the `FlowConditionDeps` map.
   llvm::DenseMap<AtomicBoolValue *, llvm::DenseSet<AtomicBoolValue *>>
       FlowConditionDeps;
   llvm::DenseMap<AtomicBoolValue *, BoolValue *> FlowConditionConstraints;
 
-  llvm::DenseMap<const FunctionDecl *, ControlFlowContext> FunctionContexts;
+  // Analyzable function definitions, keyed by fully qualified names. See
+  // constructor for detailed explanation.
+  llvm::StringMap<ControlFlowContext> InlineableFunctions;
+
+  // Cache mapping function decls directly to `ControlFlowContext`s. If present
+  // in the map, the function decl has already been checked against
+  // `InlineableFunctions` (null if not a member of `InlineableFunctions`). If
+  // absent, the decl has not yet been checked.
+  llvm::DenseMap<const FunctionDecl *, const ControlFlowContext*>
+      FunctionContexts;
 };
 
 } // namespace dataflow
 } // namespace clang
 
 #endif // LLVM_CLANG_ANALYSIS_FLOWSENSITIVE_DATAFLOWANALYSISCONTEXT_H

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

public:
   ///
   /// If `DeclCtx` is a function, initializes the environment with symbolic
   /// representations of the function parameters.
   ///
   /// If `DeclCtx` is a non-static member function, initializes the environment
   /// with a symbolic representation of the `this` pointee.
   Environment(DataflowAnalysisContext &DACtx, const DeclContext &DeclCtx);
 
-  /// Creates and returns an environment to use for an inline analysis  of the
-  /// callee. Uses the storage location from each argument in the `Call` as the
-  /// storage location for the corresponding parameter in the callee.
+  /// Creates and returns an environment to use for an inline analysis of the
+  /// callee. `FuncDecl` provides the definition of the function to be analyzed.
+  /// Uses the storage location from each argument in the `Call` as the storage
+  /// location for the corresponding parameter in the callee.
   ///
   /// Requirements:
   ///
-  ///  The callee of `Call` must be a `FunctionDecl` with a body.
+  ///  The body of `FuncDecl` must not reference globals.
   ///
-  ///  The body of the callee must not reference globals.
-  ///
-  ///  The arguments of `Call` must map 1:1 to the callee's parameters.
+  ///  The arguments of `Call` must map 1:1 to `FuncDecl`'s parameters.
   ///
   ///  Each argument of `Call` must already have a `StorageLocation`.
-  Environment pushCall(const CallExpr *Call) const;
+  Environment pushCall(const CallExpr *Call, const FunctionDecl &FuncDecl) const;
 
   /// Moves gathered information back into `this` from a `CalleeEnv` created via
   /// `pushCall`.
   void popCall(const Environment &CalleeEnv);
 
   /// Returns true if and only if the environment is equivalent to `Other`, i.e
   /// the two environments:
   ///  - have the same mappings from declarations to storage locations,

clang/lib/Analysis/FlowSensitive/ControlFlowContext.cpp

 //
 //===----------------------------------------------------------------------===//
 
 #include "clang/Analysis/FlowSensitive/ControlFlowContext.h"
 #include "clang/AST/ASTContext.h"
 #include "clang/AST/Decl.h"
 #include "clang/AST/Stmt.h"
 #include "clang/Analysis/CFG.h"
+#include "clang/Frontend/ASTUnit.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/Support/Error.h"
 #include <utility>
 
 namespace clang {
 namespace dataflow {
 
 /// Returns a map from statements to basic blocks that contain them.

 
 llvm::Expected<ControlFlowContext>
 ControlFlowContext::build(const Decl *D, Stmt *S, ASTContext *C) {
   assert(S != nullptr);
   assert(C != nullptr);
   return build(D, *S, *C);
 }
 
+static void InsertFunction(FunctionDecl *F,
+                           llvm::StringMap<ControlFlowContext> &Functions) {
+  auto *Body = F->getBody();
+  if (Body == nullptr)
+    return;
+  auto CFCtx = ControlFlowContext::build(F, *Body, F->getASTContext());
+  assert(CFCtx &&
+         "All functions in model should translate to CFG successfully");
+  std::string QualName;
+  llvm::raw_string_ostream OS(QualName);
+  F->printQualifiedName(OS);
+  Functions.insert({std::move(QualName), std::move(*CFCtx)});
+}
+
+llvm::Expected<llvm::StringMap<ControlFlowContext>>
+buildFunctionMapFromAST(ASTUnit &Unit) {
+  llvm::StringMap<ControlFlowContext> Functions;
+  for (auto It = Unit.top_level_begin(); It != Unit.top_level_end(); ++It) {
+    if (auto *C = dyn_cast<CXXRecordDecl>(*It)) {
+      for (auto *M : C->methods())
+        InsertFunction(M, Functions);
+    } else if (auto *F = dyn_cast<FunctionDecl>(*It)) {
+      InsertFunction(F, Functions);
+    }
+  }
+  return Functions;
+}
+
 } // namespace dataflow
 } // namespace clang

clang/lib/Analysis/FlowSensitive/DataflowAnalysisContext.cpp

 //===-- DataflowAnalysisContext.cpp -----------------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 //  This file defines a DataflowAnalysisContext class that owns objects that
 //  encompass the state of a program and stores context that is used during
 //  dataflow analysis.
 //
 //===----------------------------------------------------------------------===//
 
 #include "clang/Analysis/FlowSensitive/DataflowAnalysisContext.h"
 #include "clang/AST/ExprCXX.h"
+#include "clang/Analysis/FlowSensitive/ControlFlowContext.h"
 #include "clang/Analysis/FlowSensitive/DebugSupport.h"
 #include "clang/Analysis/FlowSensitive/Value.h"
 #include "llvm/Support/Debug.h"
 #include <cassert>
 #include <memory>
 #include <utility>
 
 namespace clang {
void DataflowAnalysisContext::dumpFlowCondition(AtomicBoolValue &Token) {
   llvm::DenseMap<const AtomicBoolValue *, std::string> AtomNames = {
       {&getBoolLiteralValue(false), "False"},
       {&getBoolLiteralValue(true), "True"}};
   llvm::dbgs() << debugString(Constraints, AtomNames);
 }
 
 const ControlFlowContext *
 DataflowAnalysisContext::getControlFlowContext(const FunctionDecl *F) {
-  // Canonicalize the key:
-  F = F->getDefinition();
-  if (F == nullptr)
-    return nullptr;
+  // Canonicalize the key.
+  F = F->getCanonicalDecl();
   auto It = FunctionContexts.find(F);
   if (It != FunctionContexts.end())
-    return &It->second;
+    return It->second;
 
-  if (Stmt *Body = F->getBody()) {
-    auto CFCtx = ControlFlowContext::build(F, *Body, F->getASTContext());
-    // FIXME: Handle errors.
-    assert(CFCtx);
-    auto Result = FunctionContexts.insert({F, std::move(*CFCtx)});
-    return &Result.first->second;
-  }
-
-  return nullptr;
+  const ControlFlowContext *CFCtx = nullptr;
+  std::string QualName;
+  llvm::raw_string_ostream OS(QualName);
+  F->printQualifiedName(OS);
+  auto It2 = InlineableFunctions.find(QualName);
+  if (It2 != InlineableFunctions.end())
+    CFCtx = &It2->second;
+  FunctionContexts.insert({F, CFCtx});
+  return CFCtx;
 }
 
 } // namespace dataflow
 } // namespace clang
 
 using namespace clang;
 
 const Expr &clang::dataflow::ignoreCFGOmittedNodes(const Expr &E) {

clang/lib/Analysis/FlowSensitive/DataflowEnvironment.cpp

if (MethodDecl && !MethodDecl->isStatic()) {
         ThisPointeeLoc = &createStorageLocation(ThisPointeeType);
         if (Value *ThisPointeeVal = createValue(ThisPointeeType))
           setValue(*ThisPointeeLoc, *ThisPointeeVal);
       }
     }
   }
 }
 
-Environment Environment::pushCall(const CallExpr *Call) const {
+Environment Environment::pushCall(const CallExpr *Call,
+                                  const FunctionDecl &FuncDecl) const {
   Environment Env(*this);
   // FIXME: Support references here.
   Env.ReturnLoc = Env.getStorageLocation(*Call, SkipPast::Reference);
 
-  const auto *FuncDecl = Call->getDirectCallee();
-  assert(FuncDecl != nullptr);
   // FIXME: In order to allow the callee to reference globals, we probably need
   // to call `initGlobalVars` here in some way.
 
   if (const auto *MethodCall = dyn_cast<CXXMemberCallExpr>(Call)) {
     if (const Expr *Arg = MethodCall->getImplicitObjectArgument()) {
       Env.ThisPointeeLoc = Env.getStorageLocation(*Arg, SkipPast::Reference);
     }
   }
 
-  auto ParamIt = FuncDecl->param_begin();
+  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());
+    assert(ParamIt != FuncDecl.param_end());
 
     const Expr *Arg = *ArgIt;
     auto *ArgLoc = Env.getStorageLocation(*Arg, SkipPast::Reference);
     assert(ArgLoc != nullptr);
 
     const VarDecl *Param = *ParamIt;
     auto &Loc = Env.createStorageLocation(*Param);
     Env.setStorageLocation(*Param, Loc);

clang/lib/Analysis/FlowSensitive/Transfer.cpp

if (S->isCallToStdMove()) {
       if (ArgLoc == nullptr)
         return;
       Env.setStorageLocation(*S, *ArgLoc);
     } else if (const FunctionDecl *F = S->getDirectCallee()) {
       // This case is for context-sensitive analysis.
       if (!Options.ContextSensitive)
         return;
 
+      // Note that it is important for the storage location of `S` to be set
+      // before `pushCall`, because the latter uses it to set the storage
+      // location for `return`.
+      auto &ReturnLoc = Env.createStorageLocation(*S);
+      Env.setStorageLocation(*S, ReturnLoc);
+
       const ControlFlowContext *CFCtx = Env.getControlFlowContext(F);
       if (!CFCtx)
         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();
 
-      // Note that it is important for the storage location of `S` to be set
-      // before `pushCall`, because the latter uses it to set the storage
-      // location for `return`.
-      auto &ReturnLoc = Env.createStorageLocation(*S);
-      Env.setStorageLocation(*S, ReturnLoc);
-      auto CalleeEnv = Env.pushCall(S);
+      // Note that the decl may be different than that associated with the call
+      // expression, because the environment can map the actual callee to any
+      // (compatible) implementation.
+      const FunctionDecl *FuncDecl = CFCtx->getDecl()->getAsFunction();
+      assert(FuncDecl != nullptr && "ControlFlowContexts in the environment "
+                                    "should always carry a FunctionDecl");
+      auto CalleeEnv = Env.pushCall(S, *FuncDecl);
 
       // 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());
 

clang/unittests/Analysis/FlowSensitive/TestingSupport.h

struct AnalysisData {
   ASTContext &ASTCtx;
   const ControlFlowContext &CFCtx;
   const Environment &Env;
   TypeErasedDataflowAnalysis &Analysis;
   llvm::DenseMap<const clang::Stmt *, std::string> &Annotations;
   std::vector<llvm::Optional<TypeErasedDataflowAnalysisState>> &BlockStates;
 };
 
+/// Requirements:
+///
+///  `ASTUnit` must be built from `Code`.
 template <typename AnalysisT>
 llvm::Error checkDataflow(
-    llvm::StringRef Code,
+    llvm::StringRef Code, std::unique_ptr<ASTUnit> ASTUnit,
     ast_matchers::internal::Matcher<FunctionDecl> TargetFuncMatcher,
     std::function<AnalysisT(ASTContext &, Environment &)> MakeAnalysis,
     std::function<void(ASTContext &, const CFGStmt &,
                        const TypeErasedDataflowAnalysisState &)>
         PostVisitStmt,
-    std::function<void(AnalysisData)> VerifyResults, ArrayRef<std::string> Args,
-    const tooling::FileContentMappings &VirtualMappedFiles = {}) {
-  llvm::Annotations AnnotatedCode(Code);
-  auto Unit = tooling::buildASTFromCodeWithArgs(
-      AnnotatedCode.code(), Args, "input.cc", "clang-dataflow-test",
-      std::make_shared<PCHContainerOperations>(),
-      tooling::getClangStripDependencyFileAdjuster(), VirtualMappedFiles);
-  auto &Context = Unit->getASTContext();
+    std::function<void(AnalysisData)> VerifyResults,
+    llvm::StringMap<ControlFlowContext> AnalyzableFunctions = {}) {
+  const llvm::Annotations AnnotatedCode(Code);
+  auto &Context = ASTUnit->getASTContext();
 
   if (Context.getDiagnostics().getClient()->getNumErrors() != 0) {
     return llvm::make_error<llvm::StringError>(
         llvm::errc::invalid_argument, "Source file has syntax or type errors, "
                                       "they were printed to the test log");
   }
 
   const FunctionDecl *F = ast_matchers::selectFirst<FunctionDecl>(
       "target",
       ast_matchers::match(ast_matchers::functionDecl(
                               ast_matchers::isDefinition(), TargetFuncMatcher)
                               .bind("target"),
                           Context));
   if (F == nullptr)
     return llvm::make_error<llvm::StringError>(
         llvm::errc::invalid_argument, "Could not find target function.");
 
-  auto CFCtx = ControlFlowContext::build(F, F->getBody(), &F->getASTContext());
+  auto CFCtx = ControlFlowContext::build(F, *F->getBody(), F->getASTContext());
   if (!CFCtx)
     return CFCtx.takeError();
 
-  DataflowAnalysisContext DACtx(std::make_unique<WatchedLiteralsSolver>());
+  DataflowAnalysisContext DACtx(std::make_unique<WatchedLiteralsSolver>(),
+                                std::move(AnalyzableFunctions));
   Environment Env(DACtx, *F);
   auto Analysis = MakeAnalysis(Context, Env);
 
   std::function<void(const CFGStmt &, const TypeErasedDataflowAnalysisState &)>
       PostVisitStmtClosure = nullptr;
   if (PostVisitStmt != nullptr) {
     PostVisitStmtClosure = [&PostVisitStmt, &Context](
                                const CFGStmt &Stmt,
llvm::Error checkDataflow(
   auto &BlockStates = *MaybeBlockStates;
 
   AnalysisData AnalysisData{Context,  *CFCtx,      Env,
                             Analysis, Annotations, BlockStates};
   VerifyResults(AnalysisData);
   return llvm::Error::success();
 }
 
-// Runs dataflow on the body of the function that matches `TargetFuncMatcher` in
-// code snippet `Code`. Requires: `AnalysisT` contains a type `Lattice`.
+/// `AnalyzableFunctions` is passed to the `DataflowAnalysisContext`
+/// constructor.
+// FIXME: This now called in one place. Inline and remove.
 template <typename AnalysisT>
 llvm::Error checkDataflow(
     llvm::StringRef Code,
     ast_matchers::internal::Matcher<FunctionDecl> TargetFuncMatcher,
     std::function<AnalysisT(ASTContext &, Environment &)> MakeAnalysis,
+    std::function<void(ASTContext &, const CFGStmt &,
+                       const TypeErasedDataflowAnalysisState &)>
+        PostVisitStmt,
+    std::function<void(AnalysisData)> VerifyResults, ArrayRef<std::string> Args,
+    const tooling::FileContentMappings &VirtualMappedFiles = {},
+    llvm::StringMap<ControlFlowContext> AnalyzableFunctions = {}) {
+  auto Unit = tooling::buildASTFromCodeWithArgs(
+      Code, Args, "input.cc", "clang-dataflow-test",
+      std::make_shared<PCHContainerOperations>(),
+      tooling::getClangStripDependencyFileAdjuster(), VirtualMappedFiles);
+  return checkDataflow(Code, std::move(Unit), std::move(TargetFuncMatcher),
+                       std::move(MakeAnalysis), std::move(PostVisitStmt),
+                       std::move(VerifyResults),
+                       std::move(AnalyzableFunctions));
+}
+
+/// Runs dataflow on the body of the function that matches `TargetFuncMatcher`
+/// in code snippet `Code`. `AnalyzableFunctions` is passed to the
+/// `DataflowAnalysisContext` constructor. Requires: `AnalysisT` contains a type
+/// `Lattice`.
+template <typename AnalysisT>
+llvm::Error checkDataflow(
+    llvm::StringRef Code, std::unique_ptr<ASTUnit> ASTUnit,
+    ast_matchers::internal::Matcher<FunctionDecl> TargetFuncMatcher,
+    std::function<AnalysisT(ASTContext &, Environment &)> MakeAnalysis,
     std::function<void(
         llvm::ArrayRef<std::pair<
             std::string, DataflowAnalysisState<typename AnalysisT::Lattice>>>,
         ASTContext &)>
         VerifyResults,
-    ArrayRef<std::string> Args,
-    const tooling::FileContentMappings &VirtualMappedFiles = {}) {
+    llvm::StringMap<ControlFlowContext> AnalyzableFunctions = {}) {
   using StateT = DataflowAnalysisState<typename AnalysisT::Lattice>;
 
   return checkDataflow(
-      Code, std::move(TargetFuncMatcher), std::move(MakeAnalysis),
+      Code, std::move(ASTUnit), std::move(TargetFuncMatcher),
+      std::move(MakeAnalysis),
       /*PostVisitStmt=*/nullptr,
       [&VerifyResults](AnalysisData AnalysisData) {
         if (AnalysisData.BlockStates.empty()) {
           VerifyResults({}, AnalysisData.ASTCtx);
           return;
         }
 
         auto &Annotations = AnalysisData.Annotations;
return checkDataflow(
                   return;
                 auto *Lattice = llvm::any_cast<typename AnalysisT::Lattice>(
                     &State.Lattice.Value);
                 Results.emplace_back(It->second, StateT{*Lattice, State.Env});
               });
         }
         VerifyResults(Results, AnalysisData.ASTCtx);
       },
-      Args, VirtualMappedFiles);
+      std::move(AnalyzableFunctions));
 }
 
-// Runs dataflow on the body of the function named `target_fun` in code snippet
-// `code`.
+/// Runs dataflow on the body of the function named `target_fun` in code snippet
+/// `code`. `AnalyzableFunctions` is passed to the `DataflowAnalysisContext`
+/// constructor.
 template <typename AnalysisT>
 llvm::Error checkDataflow(
     llvm::StringRef Code, llvm::StringRef TargetFun,
     std::function<AnalysisT(ASTContext &, Environment &)> MakeAnalysis,
     std::function<void(
         llvm::ArrayRef<std::pair<
             std::string, DataflowAnalysisState<typename AnalysisT::Lattice>>>,
         ASTContext &)>
         VerifyResults,
     ArrayRef<std::string> Args,
-    const tooling::FileContentMappings &VirtualMappedFiles = {}) {
-  return checkDataflow(Code, ast_matchers::hasName(TargetFun),
-                       std::move(MakeAnalysis), std::move(VerifyResults), Args,
-                       VirtualMappedFiles);
+    const tooling::FileContentMappings &VirtualMappedFiles = {},
+    llvm::StringMap<ControlFlowContext> AnalyzableFunctions = {}) {
+  auto ASTUnit = tooling::buildASTFromCodeWithArgs(
+      Code, Args, "input.cc", "clang-dataflow-test",
+      std::make_shared<PCHContainerOperations>(),
+      tooling::getClangStripDependencyFileAdjuster(), VirtualMappedFiles);
+
+  return checkDataflow(Code, std::move(ASTUnit),
+                       ast_matchers::hasName(TargetFun),
+                       std::move(MakeAnalysis), std::move(VerifyResults),
+                       std::move(AnalyzableFunctions));
 }
 
 /// Returns the `ValueDecl` for the given identifier.
 ///
 /// Requirements:
 ///
 ///  `Name` must be unique in `ASTCtx`.
 const ValueDecl *findValueDecl(ASTContext &ASTCtx, llvm::StringRef Name);

clang/unittests/Analysis/FlowSensitive/TransferTest.cpp

 //===- unittests/Analysis/FlowSensitive/TransferTest.cpp ------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "TestingSupport.h"
 #include "clang/AST/ASTContext.h"
 #include "clang/AST/Decl.h"
 #include "clang/ASTMatchers/ASTMatchFinder.h"
 #include "clang/ASTMatchers/ASTMatchers.h"
+#include "clang/Analysis/FlowSensitive/ControlFlowContext.h"
 #include "clang/Analysis/FlowSensitive/DataflowEnvironment.h"
 #include "clang/Analysis/FlowSensitive/NoopAnalysis.h"
 #include "clang/Analysis/FlowSensitive/StorageLocation.h"
 #include "clang/Analysis/FlowSensitive/Value.h"
 #include "clang/Basic/LangStandard.h"
+#include "clang/Tooling/Tooling.h"
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Testing/Support/Error.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 #include <string>
 #include <utility>
 
 namespace {
 
 using namespace clang;
 using namespace dataflow;
 using namespace test;
 using ::testing::_;
 using ::testing::ElementsAre;
 using ::testing::IsNull;
 using ::testing::NotNull;
 using ::testing::Pair;
 using ::testing::SizeIs;
 
 template <typename Matcher>
 void runDataflow(llvm::StringRef Code, Matcher Match,
                  DataflowAnalysisOptions Options,
+                 llvm::StringRef ModeledFunctionsCode = {},
                  LangStandard::Kind Std = LangStandard::lang_cxx17,
                  llvm::StringRef TargetFun = "target") {
+  const std::vector<std::string> Args = {
+      "-fsyntax-only", "-fno-delayed-template-parsing",
+      "-std=" +
+          std::string(LangStandard::getLangStandardForKind(Std).getName())};
+
+  auto MainASTUnit = tooling::buildASTFromCodeWithArgs(Code, Args);
+  ASSERT_NE(MainASTUnit, nullptr);
+
+  // The unique_ptr has to live outside the if-body since the lifetime of the
+  // function map is linked to the lifetime of the ASTUnit. The raw pointer
+  // allows us to select between the main TU and the (optional) model TU.
+  std::unique_ptr<ASTUnit> ModelASTUnit = nullptr;
+  ASTUnit *UnitPtr = nullptr;
+  if (ModeledFunctionsCode.empty()) {
+    UnitPtr = MainASTUnit.get();
+  } else {
+    ModelASTUnit =
+        tooling::buildASTFromCodeWithArgs(ModeledFunctionsCode, Args);
+    ASSERT_NE(ModelASTUnit, nullptr);
+    UnitPtr = ModelASTUnit.get();
+  }
+
+  llvm::Expected<llvm::StringMap<ControlFlowContext>> Result =
+      buildFunctionMapFromAST(*UnitPtr);
+  ASSERT_TRUE((bool)Result);
+  llvm::StringMap<ControlFlowContext> AnalyzableFunctions = std::move(*Result);
+
   ASSERT_THAT_ERROR(
       test::checkDataflow<NoopAnalysis>(
-          Code, TargetFun,
+          Code, std::move(MainASTUnit), ast_matchers::hasName(TargetFun),
           [Options](ASTContext &C, Environment &) {
             return NoopAnalysis(C, Options);
           },
           [&Match](
               llvm::ArrayRef<
                   std::pair<std::string, DataflowAnalysisState<NoopLattice>>>
                   Results,
               ASTContext &ASTCtx) { Match(Results, ASTCtx); },
-          {"-fsyntax-only", "-fno-delayed-template-parsing",
-           "-std=" + std::string(
-                         LangStandard::getLangStandardForKind(Std).getName())}),
+          std::move(AnalyzableFunctions)),
       llvm::Succeeded());
 }
 
+// FIXME: inline this overload and remove it.
 template <typename Matcher>
 void runDataflow(llvm::StringRef Code, Matcher Match,
                  LangStandard::Kind Std = LangStandard::lang_cxx17,
                  bool ApplyBuiltinTransfer = true,
                  llvm::StringRef TargetFun = "target") {
-  runDataflow(Code, Match, {ApplyBuiltinTransfer, {}}, Std, TargetFun);
+  runDataflow(Code, Match, /*Options=*/{ApplyBuiltinTransfer, {}},
+              /*AnalyzableFunctions=*/{}, Std, TargetFun);
 }
 
 TEST(TransferTest, IntVarDeclNotTrackedWhenTransferDisabled) {
   std::string Code = R"(
     void target() {
       int Foo;
       // [[p]]
     }
std::string Code = R"(
     class B : public A {
     public:
       B(int I) : A(I) {
         (void)0;
         // [[p]]
       }
     };
   )";
+
+  auto ASTUnit = tooling::buildASTFromCodeWithArgs(
+      Code,
+      /*Args=*/
+      {"-fsyntax-only", "-fno-delayed-template-parsing",
+       "-std=" + std::string(LangStandard::getLangStandardForKind(
+                                 LangStandard::lang_cxx17)
+                                 .getName())},
+      "input.cc", "clang-dataflow-test");
+
   ASSERT_THAT_ERROR(
       test::checkDataflow<NoopAnalysis>(
-          Code, cxxConstructorDecl(ofClass(hasName("B"))),
+          Code, std::move(ASTUnit), cxxConstructorDecl(ofClass(hasName("B"))),
           [](ASTContext &C, Environment &) {
             return NoopAnalysis(C, /*ApplyBuiltinTransfer=*/true);
           },
           [](llvm::ArrayRef<
                  std::pair<std::string, DataflowAnalysisState<NoopLattice>>>
                  Results,
              ASTContext &ASTCtx) {
             // Regression test to verify that base-class initializers do not
             // trigger an assertion. If we add support for such initializers in
             // the future, we can expand this test to check more specific
             // properties.
             EXPECT_THAT(Results, ElementsAre(Pair("p", _)));
-          },
-          {"-fsyntax-only", "-fno-delayed-template-parsing",
-           "-std=" + std::string(LangStandard::getLangStandardForKind(
-                                     LangStandard::lang_cxx17)
-                                     .getName())}),
+          }),
       llvm::Succeeded());
 }
 
 TEST(TransferTest, ReferenceMember) {
   std::string Code = R"(
     struct A {
       int &Bar;
     };
std::string Code = R"(
     void SetBool(bool &Var) { Var = true; }
 
     void target() {
       bool Foo = GiveBool();
       SetBool(Foo);
       // [[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;
+  runDataflow(
+      Code,
 
-                const ValueDecl *FooDecl = findValueDecl(ASTCtx, "Foo");
-                ASSERT_THAT(FooDecl, NotNull());
 
-                auto &FooVal =
-                    *cast<BoolValue>(Env.getValue(*FooDecl, SkipPast::None));
-                EXPECT_FALSE(Env.flowConditionImplies(FooVal));
-                EXPECT_FALSE(Env.flowConditionImplies(Env.makeNot(FooVal)));
-              },
-              {/*.ApplyBuiltinTransfer=*/true,
-               /*.BuiltinTransferOptions=*/{/*.ContextSensitive=*/false}});
+        auto &FooVal = *cast<BoolValue>(Env.getValue(*FooDecl, SkipPast::None));
+       /*.BuiltinTransferOptions=*/{/*.ContextSensitive=*/false}},
+      /*ModeledFunctionsCode=*/"");
 }
 
 TEST(TransferTest, ContextSensitiveSetTrue) {
+  std::string Functions = R"(
+    void SetBool(bool &Var) { Var = true; }
+  )";
   std::string Code = R"(
     bool GiveBool();
-    void SetBool(bool &Var) { Var = true; }
+    void SetBool(bool &Var);
 
     void target() {
       bool Foo = GiveBool();
       SetBool(Foo);
       // [[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}});
+               /*.BuiltinTransferOptions=*/{/*.ContextSensitive=*/true}},
+              Functions);
 }
 
 TEST(TransferTest, ContextSensitiveSetFalse) {
+  std::string Functions = R"(void SetBool(bool &Var) { Var = false; })";
   std::string Code = R"(
     bool GiveBool();
-    void SetBool(bool &Var) { Var = false; }
+    void SetBool(bool &Var);
 
     void target() {
       bool Foo = GiveBool();
       SetBool(Foo);
       // [[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;
+  runDataflow(
+      Code,
 
-                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}});
+        auto &FooVal = *cast<BoolValue>(Env.getValue(*FooDecl, SkipPast::None));
+       /*.BuiltinTransferOptions=*/{/*.ContextSensitive=*/true}},
+      Functions);
 }
 
 TEST(TransferTest, ContextSensitiveSetBothTrueAndFalse) {
+  std::string Functions = R"(
+    void SetBool(bool &Var, bool Val) { Var = Val; }
+  )";
   std::string Code = R"(
     bool GiveBool();
-    void SetBool(bool &Var, bool Val) { Var = Val; }
+    void SetBool(bool &Var, bool Val);
 
     void target() {
       bool Foo = GiveBool();
       bool Bar = GiveBool();
       SetBool(Foo, true);
       SetBool(Bar, false);
       // [[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;
+  runDataflow(
+      Code,
 
-                const ValueDecl *FooDecl = findValueDecl(ASTCtx, "Foo");
-                ASSERT_THAT(FooDecl, NotNull());
 
-                const ValueDecl *BarDecl = findValueDecl(ASTCtx, "Bar");
-                ASSERT_THAT(BarDecl, NotNull());
 
-                auto &FooVal =
-                    *cast<BoolValue>(Env.getValue(*FooDecl, SkipPast::None));
-                EXPECT_TRUE(Env.flowConditionImplies(FooVal));
-                EXPECT_FALSE(Env.flowConditionImplies(Env.makeNot(FooVal)));
+        auto &FooVal = *cast<BoolValue>(Env.getValue(*FooDecl, SkipPast::None));
 
-                auto &BarVal =
-                    *cast<BoolValue>(Env.getValue(*BarDecl, SkipPast::None));
-                EXPECT_FALSE(Env.flowConditionImplies(BarVal));
-                EXPECT_TRUE(Env.flowConditionImplies(Env.makeNot(BarVal)));
-              },
-              {/*.ApplyBuiltinTransfer=*/true,
-               /*.BuiltinTransferOptions=*/{/*.ContextSensitive=*/true}});
+        auto &BarVal = *cast<BoolValue>(Env.getValue(*BarDecl, SkipPast::None));
+       /*.BuiltinTransferOptions=*/{/*.ContextSensitive=*/true}},
+      Functions);
 }
 
 TEST(TransferTest, ContextSensitiveSetTwoLayers) {
-  std::string Code = R"(
-    bool GiveBool();
+  std::string Functions = R"(
     void SetBool1(bool &Var) { Var = true; }
     void SetBool2(bool &Var) { SetBool1(Var); }
+  )";
+  std::string Code = R"(
+    bool GiveBool();
+    void SetBool2(bool &Var);
 
     void target() {
       bool Foo = GiveBool();
       SetBool2(Foo);
       // [[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;
+  runDataflow(
+      Code,
 
-                const ValueDecl *FooDecl = findValueDecl(ASTCtx, "Foo");
-                ASSERT_THAT(FooDecl, NotNull());
 
-                auto &FooVal =
-                    *cast<BoolValue>(Env.getValue(*FooDecl, SkipPast::None));
-                EXPECT_FALSE(Env.flowConditionImplies(FooVal));
-                EXPECT_FALSE(Env.flowConditionImplies(Env.makeNot(FooVal)));
-              },
-              {/*.ApplyBuiltinTransfer=*/true,
-               /*.BuiltinTransferOptions=*/{/*.ContextSensitive=*/true}});
+        auto &FooVal = *cast<BoolValue>(Env.getValue(*FooDecl, SkipPast::None));
+       /*.BuiltinTransferOptions=*/{/*.ContextSensitive=*/true}},
+      Functions);
 }
 
 TEST(TransferTest, ContextSensitiveSetMultipleLines) {
-  std::string Code = R"(
+  std::string Functions = R"(
     void SetBools(bool &Var1, bool &Var2) {
       Var1 = true;
       Var2 = false;
     }
+  )";
+  std::string Code = R"(
+    void SetBools(bool &Var1, bool &Var2);
 
     void target() {
       bool Foo = false;
       bool Bar = true;
       SetBools(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;
+  runDataflow(
+      Code,
 
-                const ValueDecl *FooDecl = findValueDecl(ASTCtx, "Foo");
-                ASSERT_THAT(FooDecl, NotNull());
 
-                const ValueDecl *BarDecl = findValueDecl(ASTCtx, "Bar");
-                ASSERT_THAT(BarDecl, NotNull());
 
-                auto &FooVal =
-                    *cast<BoolValue>(Env.getValue(*FooDecl, SkipPast::None));
-                EXPECT_TRUE(Env.flowConditionImplies(FooVal));
-                EXPECT_FALSE(Env.flowConditionImplies(Env.makeNot(FooVal)));
+        auto &FooVal = *cast<BoolValue>(Env.getValue(*FooDecl, SkipPast::None));
 
-                auto &BarVal =
-                    *cast<BoolValue>(Env.getValue(*BarDecl, SkipPast::None));
-                EXPECT_FALSE(Env.flowConditionImplies(BarVal));
-                EXPECT_TRUE(Env.flowConditionImplies(Env.makeNot(BarVal)));
-              },
-              {/*.ApplyBuiltinTransfer=*/true,
-               /*.BuiltinTransferOptions=*/{/*.ContextSensitive=*/true}});
+        auto &BarVal = *cast<BoolValue>(Env.getValue(*BarDecl, SkipPast::None));
+       /*.BuiltinTransferOptions=*/{/*.ContextSensitive=*/true}},
+      Functions);
 }
 
 TEST(TransferTest, ContextSensitiveSetMultipleBlocks) {
-  std::string Code = R"(
+  std::string Functions = R"(
     void IfCond(bool Cond, bool &Then, bool &Else) {
       if (Cond) {
         Then = true;
       } else {
         Else = true;
       }
     }
+  )";
+  std::string Code = R"(
+    void IfCond(bool Cond, bool &Then, bool &Else);
 
     void target() {
       bool Foo = false;
       bool Bar = false;
       bool Baz = false;
       IfCond(Foo, Bar, Baz);
       // [[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;
+  runDataflow(
+      Code,
 
-                const ValueDecl *BarDecl = findValueDecl(ASTCtx, "Bar");
-                ASSERT_THAT(BarDecl, NotNull());
 
-                const ValueDecl *BazDecl = findValueDecl(ASTCtx, "Baz");
-                ASSERT_THAT(BazDecl, NotNull());
 
-                auto &BarVal =
-                    *cast<BoolValue>(Env.getValue(*BarDecl, SkipPast::None));
-                EXPECT_FALSE(Env.flowConditionImplies(BarVal));
-                EXPECT_TRUE(Env.flowConditionImplies(Env.makeNot(BarVal)));
+        auto &BarVal = *cast<BoolValue>(Env.getValue(*BarDecl, SkipPast::None));
 
-                auto &BazVal =
-                    *cast<BoolValue>(Env.getValue(*BazDecl, SkipPast::None));
-                EXPECT_TRUE(Env.flowConditionImplies(BazVal));
-                EXPECT_FALSE(Env.flowConditionImplies(Env.makeNot(BazVal)));
-              },
-              {/*.ApplyBuiltinTransfer=*/true,
-               /*.BuiltinTransferOptions=*/{/*.ContextSensitive=*/true}});
+        auto &BazVal = *cast<BoolValue>(Env.getValue(*BazDecl, SkipPast::None));
+       /*.BuiltinTransferOptions=*/{/*.ContextSensitive=*/true}},
+      Functions);
 }
 
+llvm::StringLiteral Model = R"(
+    void Noop() { return; }
+    bool GiveTrue() { return true; }
+    bool GiveFalse() { return false; }
+    bool GiveBack(bool Arg) { return Arg; }
+)";
+
 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}});
+  runDataflow(
+      Code,
+       /*.BuiltinTransferOptions=*/{/*.ContextSensitive=*/true}},
+      Model);
 }
 
 TEST(TransferTest, ContextSensitiveReturnTrue) {
   std::string Code = R"(
-    bool GiveBool() { return true; }
+    bool GiveTrue();
 
     void target() {
-      bool Foo = GiveBool();
+      bool Foo = GiveTrue();
       // [[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}});
+               /*.BuiltinTransferOptions=*/{/*.ContextSensitive=*/true}}, Model);
 }
 
 TEST(TransferTest, ContextSensitiveReturnFalse) {
   std::string Code = R"(
-    bool GiveBool() { return false; }
+    bool GiveFalse();
 
     void target() {
-      bool Foo = GiveBool();
+      bool Foo = GiveFalse();
       // [[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}});
+               /*.BuiltinTransferOptions=*/{/*.ContextSensitive=*/true}}, Model);
 }
 
 TEST(TransferTest, ContextSensitiveReturnArg) {
   std::string Code = R"(
     bool GiveBool();
-    bool GiveBack(bool Arg) { return Arg; }
+    bool GiveBack(bool 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}});
+               /*.BuiltinTransferOptions=*/{/*.ContextSensitive=*/true}}, Model);
 }
 
 TEST(TransferTest, ContextSensitiveMethodLiteral) {
   std::string Code = R"(
     class MyClass {
     public:
       bool giveBool() { return true; }
     };
runDataflow(Code,
               {/*.ApplyBuiltinTransfer=*/true,
                /*.BuiltinTransferOptions=*/{/*.ContextSensitive=*/true}});
 }
 
 TEST(TransferTest, ContextSensitiveMethodSetter) {
   std::string Code = R"(
     class MyClass {
     public:
-      bool setField(bool Val) { Field = Val; }
+      void setField(bool Val) { Field = Val; }
 
       bool Field;
     };
 
     void target() {
       MyClass MyObj;
       MyObj.setField(true);
       bool Foo = MyObj.Field;
runDataflow(Code,
                /*.BuiltinTransferOptions=*/{/*.ContextSensitive=*/true}});
 }
 
 TEST(TransferTest, ContextSensitiveMethodGetterAndSetter) {
   std::string Code = R"(
     class MyClass {
     public:
       bool getField() { return Field; }
-      bool setField(bool Val) { Field = Val; }
+      void setField(bool Val) { Field = Val; }
 
     private:
       bool Field;
     };
 
     void target() {
       MyClass MyObj;
       MyObj.setField(true);