[-Wunsafe-buffer-usage][NFC] Refactor checkUnsafeBufferUsage

[NoQ]

I like this, this is much easier to work with than my original speculative spaghetti!

clang/lib/Analysis/UnsafeBufferUsage.cpp

578–579

Maybe move unique_ptrs out of the original vector straight into the new location, and drop the original vector entirely, instead of dealing with raw pointers?

[jkorous]

• made WarningGadgetSets and FixableGadgetSets own the gadgets
• changed temporary function parameter name StrategyForDirectlyUnsafeVariables to S

[NoQ]

Aha looks great!

clang/lib/Analysis/UnsafeBufferUsage.cpp

 //===- UnsafeBufferUsage.cpp - Replace pointers with modern 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "clang/Analysis/Analyses/UnsafeBufferUsage.h"
-#include "clang/ASTMatchers/ASTMatchFinder.h"
 #include "clang/AST/RecursiveASTVisitor.h"
+#include "clang/ASTMatchers/ASTMatchFinder.h"
 #include "llvm/ADT/SmallVector.h"
+#include <memory>
 #include <optional>
 
 using namespace llvm;
 using namespace clang;
 using namespace ast_matchers;
 
 namespace clang::ast_matchers {
 // A `RecursiveASTVisitor` that traverses all descendants of a given node "n"
class DeclUseTracker {
   // Allocate on the heap for easier move.
   std::unique_ptr<UseSetTy> Uses{std::make_unique<UseSetTy>()};
   DefMapTy Defs{};
 
 public:
   DeclUseTracker() = default;
   DeclUseTracker(const DeclUseTracker &) = delete; // Let's avoid copies.
   DeclUseTracker(DeclUseTracker &&) = default;
+  DeclUseTracker &operator=(DeclUseTracker &&) = default;
 
   // Start tracking a freshly discovered DRE.
   void discoverUse(const DeclRefExpr *DRE) { Uses->insert(DRE); }
 
   // Stop tracking the DRE as it's been fully figured out.
   void claimUse(const DeclRefExpr *DRE) {
     assert(Uses->count(DRE) &&
            "DRE not found or claimed by multiple matchers!");
for (const auto &G : CB.FixableGadgets) {
     for (const auto *DRE : G->getClaimedVarUseSites()) {
       CB.Tracker.claimUse(DRE);
     }
   }
 
   return {std::move(CB.FixableGadgets), std::move(CB.WarningGadgets), std::move(CB.Tracker)};
 }
 
-void clang::checkUnsafeBufferUsage(const Decl *D,
-                                   UnsafeBufferUsageHandler &Handler) {
-  assert(D && D->getBody());
-
-  SmallSet<const VarDecl *, 8> WarnedDecls;
-
-  auto [FixableGadgets, WarningGadgets, Tracker] = findGadgets(D);
-
-  DenseMap<const VarDecl *, std::pair<std::vector<const FixableGadget *>,
-                                std::vector<const WarningGadget *>>> Map;
+struct WarningGadgetSets {
+  std::map<const VarDecl *, std::set<std::unique_ptr<WarningGadget>>> byVar;
+  // These Gadgets are not related to pointer variables (e. g. temporaries).
+  llvm::SmallVector<std::unique_ptr<WarningGadget>, 16> noVar;
+};
 
-  // First, let's sort gadgets by variables. If some gadgets cover more than one
+static WarningGadgetSets
+groupWarningGadgetsByVar(WarningGadgetList &&AllUnsafeOperations) {
+  WarningGadgetSets result;
+  // If some gadgets cover more than one
   // variable, they'll appear more than once in the map.
-  for (const auto &G : FixableGadgets) {
-    DeclUseList DREs = G->getClaimedVarUseSites();
+  for (auto &G : AllUnsafeOperations) {
+    DeclUseList ClaimedVarUseSites = G->getClaimedVarUseSites();
 
-    // Populate the map.
-    for (const DeclRefExpr *DRE : DREs) {
+    bool AssociatedWithVarDecl = false;
+    for (const DeclRefExpr *DRE : ClaimedVarUseSites) {
       if (const auto *VD = dyn_cast<VarDecl>(DRE->getDecl())) {
-        Map[VD].first.push_back(G.get());
+        result.byVar[VD].emplace(std::move(G));
+        AssociatedWithVarDecl = true;
       }
     }
-  }
+
+    if (!AssociatedWithVarDecl) {
+      result.noVar.emplace_back(std::move(G));
+      continue;
+    }
+  }
+  return result;
 
-  for (const auto &G : WarningGadgets) {
-    DeclUseList ClaimedVarUseSites = G->getClaimedVarUseSites();
+struct FixableGadgetSets {
+  std::map<const VarDecl *, std::set<std::unique_ptr<FixableGadget>>> byVar;
+};
 
-    // Populate the map.
-    bool Pushed = false;
-    for (const DeclRefExpr *DRE : ClaimedVarUseSites) {
+static FixableGadgetSets
+groupFixablesByVar(FixableGadgetList &&AllFixableOperations) {
+  FixableGadgetSets FixablesForUnsafeVars;
+  for (auto &F : AllFixableOperations) {
+    DeclUseList DREs = F->getClaimedVarUseSites();
+
+    for (const DeclRefExpr *DRE : DREs) {
       if (const auto *VD = dyn_cast<VarDecl>(DRE->getDecl())) {
-        Map[VD].second.push_back(G.get());
-        Pushed = true;
+        FixablesForUnsafeVars.byVar[VD].emplace(std::move(F));
       }
     }
+  }
+  return FixablesForUnsafeVars;
+}
 
-    if (!Pushed) {
-      // We won't return to this gadget later. Emit the warning right away.
-      Handler.handleUnsafeOperation(G->getBaseStmt());
-      continue;
+static std::map<const VarDecl *, FixItList>
+getFixIts(FixableGadgetSets &FixablesForUnsafeVars, const Strategy &S) {
+  std::map<const VarDecl *, FixItList> FixItsForVariable;
+  for (const auto &[VD, Fixables] : FixablesForUnsafeVars.byVar) {
+    // TODO fixVariable - fixit for the variable itself
+    bool ImpossibleToFix = false;
+    llvm::SmallVector<FixItHint, 16> FixItsForVD;
+    for (const auto &F : Fixables) {
+      llvm::Optional<FixItList> Fixits = F->getFixits(S);
+      if (!Fixits) {
+        ImpossibleToFix = true;
+        break;
+      } else {
+        const FixItList CorrectFixes = Fixits.value();
+        FixItsForVD.insert(FixItsForVD.end(), CorrectFixes.begin(),
+                           CorrectFixes.end());
       }
     }
+    if (ImpossibleToFix)
+      FixItsForVariable.erase(VD);
+    else
+      FixItsForVariable[VD].insert(FixItsForVariable[VD].end(),
+                                   FixItsForVD.begin(), FixItsForVD.end());
+  }
+  return FixItsForVariable;
+}
 
+static Strategy
+getNaiveStrategy(const llvm::SmallVectorImpl<const VarDecl *> &UnsafeVars) {
   Strategy S;
+  for (const VarDecl *VD : UnsafeVars) {
+    S.set(VD, Strategy::Kind::Span);
+  }
+  return S;
+}
+
+void clang::checkUnsafeBufferUsage(const Decl *D,
+                                   UnsafeBufferUsageHandler &Handler) {
+  assert(D && D->getBody());
 
-  for (const auto &Item : Map) {
+  WarningGadgetSets UnsafeOps;
+  FixableGadgetSets FixablesForUnsafeVars;
+  DeclUseTracker Tracker;
 
-    const VarDecl *VD = Item.first;
-    const std::vector<const FixableGadget *> &VDFixableGadgets = Item.second.first;
-    const std::vector<const WarningGadget *> &VDWarningGadgets = Item.second.second;
+  {
+    auto [FixableGadgets, WarningGadgets, TrackerRes] = findGadgets(D);
+    UnsafeOps = groupWarningGadgetsByVar(std::move(WarningGadgets));
+    FixablesForUnsafeVars = groupFixablesByVar(std::move(FixableGadgets));
+    Tracker = std::move(TrackerRes);
+  }
 
-    // If the variable has no unsafe gadgets, skip it entirely.
-    if (VDWarningGadgets.empty())
-      continue;
+  // Filter out non-local vars and vars with unclaimed DeclRefExpr-s.
+  for (auto it = FixablesForUnsafeVars.byVar.cbegin();
+       it != FixablesForUnsafeVars.byVar.cend();) {
+    // FIXME: Support ParmVarDecl as well.
+    if (it->first->isLocalVarDecl() || Tracker.hasUnclaimedUses(it->first)) {
+      it = FixablesForUnsafeVars.byVar.erase(it);
+    } else {
+      ++it;
+    }
+  }
 
-    std::optional<FixItList> Fixes;
+  llvm::SmallVector<const VarDecl *, 16> UnsafeVars;
+  for (const auto &[VD, ignore] : FixablesForUnsafeVars.byVar)
+    UnsafeVars.push_back(VD);
 
-    // Avoid suggesting fixes if not all uses of the variable are identified
-    // as known gadgets.
-    // FIXME: Support parameter variables as well.
-    if (!Tracker.hasUnclaimedUses(VD) && VD->isLocalVarDecl()) {
-      // Choose the appropriate strategy. FIXME: We should try different
-      // strategies.
-      S.set(VD, Strategy::Kind::Span);
+  Strategy NaiveStrategy = getNaiveStrategy(UnsafeVars);
+  std::map<const VarDecl *, FixItList> FixItsForVariable =
+      getFixIts(FixablesForUnsafeVars, NaiveStrategy);
 
-      // Check if it works.
-      // FIXME: This isn't sufficient (or even correct) when a gadget has
-      // already produced a fixit for a different variable i.e. it was mentioned
-      // in the map twice (or more). In such case the correct thing to do is
-      // to undo the previous fix first, and then if we can't produce the new
-      // fix for both variables, revert to the old one.
-      Fixes = FixItList{};
-      for (const FixableGadget *G : VDFixableGadgets) {
-        std::optional<FixItList> F = G->getFixits(S);
-        if (!F) {
-          Fixes = std::nullopt;
-          break;
-        }
+  // FIXME Detect overlapping FixIts.
 
-        for (auto &&Fixit: *F)
-          Fixes->push_back(std::move(Fixit));
-      }
-    }
+  for (const auto &G : UnsafeOps.noVar) {
+    Handler.handleUnsafeOperation(G->getBaseStmt());
 
-    if (Fixes) {
-      // If we reach this point, the strategy is applicable.
-      Handler.handleFixableVariable(VD, std::move(*Fixes));
+  for (const auto &[VD, WarningGadgets] : UnsafeOps.byVar) {
+    auto FixItsIt = FixItsForVariable.find(VD);
+    if (FixItsIt != FixItsForVariable.end()) {
+      Handler.handleFixableVariable(VD, std::move(FixItsIt->second));
     } else {
-      // The strategy has failed. Emit the warning without the fixit.
-      S.set(VD, Strategy::Kind::Wontfix);
-      for (const WarningGadget *G : VDWarningGadgets) {
+      for (const auto &G : WarningGadgets) {
         Handler.handleUnsafeOperation(G->getBaseStmt());
       }
     }
   }
 }