[clang][dataflow] Fully support Environment construction for Stmt analysis. (llvm#91616)

Assume in fewer places that the analysis is of a `FunctionDecl`, and
initialize the `Environment` properly for `Stmt`s.

Moves constructors for `Environment` to header to make it more obvious
that there are only minor differences between them and very little
initialization in the constructors.

Tested with check-clang-tooling.

Author: bazuzi

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

@@ -19,6 +19,7 @@
 #include "clang/AST/DeclBase.h"
 #include "clang/AST/Expr.h"
 #include "clang/AST/Type.h"
+#include "clang/Analysis/FlowSensitive/ASTOps.h"
 #include "clang/Analysis/FlowSensitive/DataflowAnalysisContext.h"
 #include "clang/Analysis/FlowSensitive/DataflowLattice.h"
 #include "clang/Analysis/FlowSensitive/Formula.h"
@@ -30,9 +31,11 @@
 #include "llvm/ADT/MapVector.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/ErrorHandling.h"
+#include <cassert>
 #include <memory>
 #include <type_traits>
 #include <utility>
+#include <vector>
 
 namespace clang {
 namespace dataflow {
@@ -155,32 +158,40 @@ class Environment {
 
   /// Creates an environment that uses `DACtx` to store objects that encompass
   /// the state of a program.
-  explicit Environment(DataflowAnalysisContext &DACtx);
+  explicit Environment(DataflowAnalysisContext &DACtx)
+      : DACtx(&DACtx),
+        FlowConditionToken(DACtx.arena().makeFlowConditionToken()) {}
+
+  /// Creates an environment that uses `DACtx` to store objects that encompass
+  /// the state of a program, with `S` as the statement to analyze.
+  Environment(DataflowAnalysisContext &DACtx, Stmt &S) : Environment(DACtx) {
+    InitialTargetStmt = &S;
+  }
+
+  /// Creates an environment that uses `DACtx` to store objects that encompass
+  /// the state of a program, with `FD` as the function to analyze.
+  ///
+  /// Requirements:
+  ///
+  ///  The function must have a body, i.e.
+  ///  `FunctionDecl::doesThisDecalarationHaveABody()` must be true.
+  Environment(DataflowAnalysisContext &DACtx, const FunctionDecl &FD)
+      : Environment(DACtx, *FD.getBody()) {
+    assert(FD.doesThisDeclarationHaveABody());
+    InitialTargetFunc = &FD;
+  }
 
   // Copy-constructor is private, Environments should not be copied. See fork().
   Environment &operator=(const Environment &Other) = delete;
 
   Environment(Environment &&Other) = default;
   Environment &operator=(Environment &&Other) = default;
 
-  /// Creates an environment that uses `DACtx` to store objects that encompass
-  /// the state of a program.
-  ///
-  /// 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);
-
   /// Assigns storage locations and values to all parameters, captures, global
-  /// variables, fields and functions referenced in the function currently being
-  /// analyzed.
-  ///
-  /// Requirements:
+  /// variables, fields and functions referenced in the `Stmt` or `FunctionDecl`
+  /// passed to the constructor.
   ///
-  ///  The function must have a body, i.e.
-  ///  `FunctionDecl::doesThisDecalarationHaveABody()` must be true.
+  /// If no `Stmt` or `FunctionDecl` was supplied, this function does nothing.
   void initialize();
 
   /// Returns a new environment that is a copy of this one.
@@ -193,7 +204,7 @@ class Environment {
   /// forked flow condition references the original).
   Environment fork() const;
 
-  /// Creates and returns an environment to use for an inline analysis  of the
+  /// 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.
   ///
@@ -365,46 +376,51 @@ class Environment {
   RecordStorageLocation &
   getResultObjectLocation(const Expr &RecordPRValue) const;
 
-  /// Returns the return value of the current function. This can be null if:
+  /// Returns the return value of the function currently being analyzed.
+  /// This can be null if:
   /// - The function has a void return type
   /// - No return value could be determined for the function, for example
   ///   because it calls a function without a body.
   ///
   /// Requirements:
-  ///  The current function must have a non-reference return type.
+  ///  The current analysis target must be a function and must have a
+  ///  non-reference return type.
   Value *getReturnValue() const {
     assert(getCurrentFunc() != nullptr &&
            !getCurrentFunc()->getReturnType()->isReferenceType());
     return ReturnVal;
   }
 
-  /// Returns the storage location for the reference returned by the current
-  /// function. This can be null if function doesn't return a single consistent
-  /// reference.
+  /// Returns the storage location for the reference returned by the function
+  /// currently being analyzed. This can be null if the function doesn't return
+  /// a single consistent reference.
   ///
   /// Requirements:
-  ///  The current function must have a reference return type.
+  ///  The current analysis target must be a function and must have a reference
+  ///  return type.
   StorageLocation *getReturnStorageLocation() const {
     assert(getCurrentFunc() != nullptr &&
            getCurrentFunc()->getReturnType()->isReferenceType());
     return ReturnLoc;
   }
 
-  /// Sets the return value of the current function.
+  /// Sets the return value of the function currently being analyzed.
   ///
   /// Requirements:
-  ///  The current function must have a non-reference return type.
+  ///  The current analysis target must be a function and must have a
+  ///  non-reference return type.
   void setReturnValue(Value *Val) {
     assert(getCurrentFunc() != nullptr &&
            !getCurrentFunc()->getReturnType()->isReferenceType());
     ReturnVal = Val;
   }
 
-  /// Sets the storage location for the reference returned by the current
-  /// function.
+  /// Sets the storage location for the reference returned by the function
+  /// currently being analyzed.
   ///
   /// Requirements:
-  ///  The current function must have a reference return type.
+  ///  The current analysis target must be a function and must have a reference
+  ///  return type.
   void setReturnStorageLocation(StorageLocation *Loc) {
     assert(getCurrentFunc() != nullptr &&
            getCurrentFunc()->getReturnType()->isReferenceType());
@@ -641,23 +657,21 @@ class Environment {
   /// (or the flow condition is overly constraining) or if the solver times out.
   bool allows(const Formula &) const;
 
-  /// Returns the `DeclContext` of the block being analysed, if any. Otherwise,
-  /// returns null.
-  const DeclContext *getDeclCtx() const { return CallStack.back(); }
-
   /// Returns the function currently being analyzed, or null if the code being
   /// analyzed isn't part of a function.
   const FunctionDecl *getCurrentFunc() const {
-    return dyn_cast<FunctionDecl>(getDeclCtx());
+    return CallStack.empty() ? InitialTargetFunc : CallStack.back();
   }
 
-  /// Returns the size of the call stack.
+  /// Returns the size of the call stack, not counting the initial analysis
+  /// target.
   size_t callStackSize() const { return CallStack.size(); }
 
   /// Returns whether this `Environment` can be extended to analyze the given
-  /// `Callee` (i.e. if `pushCall` can be used), with recursion disallowed and a
-  /// given `MaxDepth`.
-  bool canDescend(unsigned MaxDepth, const DeclContext *Callee) const;
+  /// `Callee` (i.e. if `pushCall` can be used).
+  /// Recursion is not allowed. `MaxDepth` is the maximum size of the call stack
+  /// (i.e. the maximum value that `callStackSize()` may assume after the call).
+  bool canDescend(unsigned MaxDepth, const FunctionDecl *Callee) const;
 
   /// Returns the `DataflowAnalysisContext` used by the environment.
   DataflowAnalysisContext &getDataflowAnalysisContext() const { return *DACtx; }
@@ -719,15 +733,20 @@ class Environment {
                         ArrayRef<const Expr *> Args);
 
   /// Assigns storage locations and values to all global variables, fields
-  /// and functions referenced in `FuncDecl`. `FuncDecl` must have a body.
-  void initFieldsGlobalsAndFuncs(const FunctionDecl *FuncDecl);
+  /// and functions in `Referenced`.
+  void initFieldsGlobalsAndFuncs(const ReferencedDecls &Referenced);
 
   static PrValueToResultObject
   buildResultObjectMap(DataflowAnalysisContext *DACtx,
                        const FunctionDecl *FuncDecl,
                        RecordStorageLocation *ThisPointeeLoc,
                        RecordStorageLocation *LocForRecordReturnVal);
 
+  static PrValueToResultObject
+  buildResultObjectMap(DataflowAnalysisContext *DACtx, Stmt *S,
+                       RecordStorageLocation *ThisPointeeLoc,
+                       RecordStorageLocation *LocForRecordReturnVal);
+
   // `DACtx` is not null and not owned by this object.
   DataflowAnalysisContext *DACtx;
 
@@ -736,11 +755,20 @@ class Environment {
   // shared between environments in the same call.
   // https://github.com/llvm/llvm-project/issues/59005
 
-  // `DeclContext` of the block being analysed if provided.
-  std::vector<const DeclContext *> CallStack;
+  // The stack of functions called from the initial analysis target.
+  std::vector<const FunctionDecl *> CallStack;
+
+  // Initial function to analyze, if a function was passed to the constructor.
+  // Null otherwise.
+  const FunctionDecl *InitialTargetFunc = nullptr;
+  // Top-level statement of the initial analysis target.
+  // If a function was passed to the constructor, this is its body.
+  // If a statement was passed to the constructor, this is that statement.
+  // Null if no analysis target was passed to the constructor.
+  Stmt *InitialTargetStmt = nullptr;
 
   // Maps from prvalues of record type to their result objects. Shared between
-  // all environments for the same function.
+  // all environments for the same analysis target.
   // FIXME: It's somewhat unsatisfactory that we have to use a `shared_ptr`
   // here, though the cost is acceptable: The overhead of a `shared_ptr` is
   // incurred when it is copied, and this happens only relatively rarely (when
@@ -749,7 +777,7 @@ class Environment {
   std::shared_ptr<PrValueToResultObject> ResultObjectMap;
 
   // The following three member variables handle various different types of
-  // return values.
+  // return values when the current analysis target is a function.
   // - If the return type is not a reference and not a record: Value returned
   //   by the function.
   Value *ReturnVal = nullptr;
@@ -762,7 +790,7 @@ class Environment {
   RecordStorageLocation *LocForRecordReturnVal = nullptr;
 
   // The storage location of the `this` pointee. Should only be null if the
-  // function being analyzed is only a function and not a method.
+  // analysis target is not a method.
   RecordStorageLocation *ThisPointeeLoc = nullptr;
 
   // Maps from declarations and glvalue expression to storage locations that are