Fix compile time regression and memory leak

Author: usx95

clang/include/clang/Basic/Attr.td

@@ -1896,25 +1896,22 @@ def LifetimeCaptureBy : DeclOrTypeAttr {
   let Documentation = [LifetimeCaptureByDocs];
   let AdditionalMembers = [{
 private:
-  SmallVector<IdentifierInfo*, 1> ArgIdents;
-  SmallVector<SourceLocation, 1> ArgLocs;
+  IdentifierInfo** ArgIdents;
+  SourceLocation* ArgLocs;
 
 public:
   static constexpr int THIS = 0;
   static constexpr int INVALID = -1;
   static constexpr int UNKNOWN = -2;
   static constexpr int GLOBAL = -3;
 
-  void setArgs(SmallVector<IdentifierInfo*>&& Idents,
-               SmallVector<SourceLocation>&& Locs) { 
-    assert(Idents.size() == Locs.size());
-    assert(Idents.size() == params_Size);
-    ArgIdents = std::move(Idents);
-    ArgLocs = std::move(Locs);
+  void setArgs(IdentifierInfo **Idents, SourceLocation *Locs) {
+    ArgIdents = Idents;
+    ArgLocs = Locs;
   }
 
-  ArrayRef<IdentifierInfo*> getArgIdents() const { return ArgIdents; }
-  ArrayRef<SourceLocation> getArgLocs() const { return ArgLocs; }
+  auto getArgIdents() const { return ArrayRef<IdentifierInfo*>(ArgIdents, params_Size); }
+  auto getArgLocs() const { return ArrayRef<SourceLocation>(ArgLocs, params_Size); }
   void setParamIdx(size_t Idx, int Val) { 
     assert(Idx < params_Size);
     params_[Idx] = Val;

clang/lib/Sema/SemaDeclAttr.cpp

@@ -65,6 +65,7 @@
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/Demangle/Demangle.h"
 #include "llvm/IR/Assumptions.h"
+#include "llvm/IR/DerivedTypes.h"
 #include "llvm/MC/MCSectionMachO.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/MathExtras.h"
@@ -3876,33 +3877,38 @@ LifetimeCaptureByAttr *Sema::ParseLifetimeCaptureByAttr(const ParsedAttr &AL,
         << AL.getRange();
     return nullptr;
   }
-  SmallVector<IdentifierInfo *, 1> ParamIdents;
-  SmallVector<SourceLocation, 1> ParamLocs;
-  for (unsigned I = 0; I < AL.getNumArgs(); ++I) {
+  unsigned N = AL.getNumArgs();
+  IdentifierInfo **ParamIdents = new (Context) IdentifierInfo *[N];
+  SourceLocation *ParamLocs = new (Context) SourceLocation[N];
+  bool IsValid = true;
+  for (unsigned I = 0; I < N; ++I) {
     if (AL.isArgExpr(I)) {
       Expr *E = AL.getArgAsExpr(I);
       Diag(E->getExprLoc(), diag::err_capture_by_attribute_argument_unknown)
           << E << E->getExprLoc();
+      IsValid = false;
       continue;
     }
     assert(AL.isArgIdent(I));
     IdentifierLoc *IdLoc = AL.getArgAsIdent(I);
     if (IdLoc->Ident->getName() == ParamName) {
       Diag(IdLoc->Loc, diag::err_capture_by_references_itself) << IdLoc->Loc;
+      IsValid = false;
       continue;
     }
-    ParamIdents.push_back(IdLoc->Ident);
-    ParamLocs.push_back(IdLoc->Loc);
+    ParamIdents[I] = IdLoc->Ident;
+    ParamLocs[I] = IdLoc->Loc;
   }
-  SmallVector<int, 1> FakeParamIndices(ParamIdents.size(),
-                                       LifetimeCaptureByAttr::INVALID);
+  if (!IsValid)
+    return nullptr;
+  SmallVector<int> FakeParamIndices(N, LifetimeCaptureByAttr::INVALID);
   LifetimeCaptureByAttr *CapturedBy = ::new (Context) LifetimeCaptureByAttr(
       Context, AL, FakeParamIndices.data(), FakeParamIndices.size());
-  CapturedBy->setArgs(std::move(ParamIdents), std::move(ParamLocs));
+  CapturedBy->setArgs(ParamIdents, ParamLocs);
   return CapturedBy;
 }
 
-static void HandleLifetimeCaptureByAttr(Sema &S, Decl *D,
+static void handleLifetimeCaptureByAttr(Sema &S, Decl *D,
                                         const ParsedAttr &AL) {
   // Do not allow multiple attributes.
   if (D->hasAttr<LifetimeCaptureByAttr>()) {
@@ -3919,10 +3925,27 @@ static void HandleLifetimeCaptureByAttr(Sema &S, Decl *D,
 
 void Sema::LazyProcessLifetimeCaptureByParams(FunctionDecl *FD) {
   bool HasImplicitThisParam = isInstanceMethod(FD);
-
-  llvm::StringMap<int> NameIdxMapping;
-  NameIdxMapping["global"] = LifetimeCaptureByAttr::GLOBAL;
-  NameIdxMapping["unknown"] = LifetimeCaptureByAttr::UNKNOWN;
+  SmallVector<LifetimeCaptureByAttr *, 1> Attrs;
+  for (ParmVarDecl *PVD : FD->parameters())
+    if (auto *A = PVD->getAttr<LifetimeCaptureByAttr>())
+      Attrs.push_back(A);
+  if (HasImplicitThisParam) {
+    TypeSourceInfo *TSI = FD->getTypeSourceInfo();
+    if (!TSI)
+      return;
+    AttributedTypeLoc ATL;
+    for (TypeLoc TL = TSI->getTypeLoc();
+         (ATL = TL.getAsAdjusted<AttributedTypeLoc>());
+         TL = ATL.getModifiedLoc()) {
+      if (auto *A = ATL.getAttrAs<LifetimeCaptureByAttr>())
+        Attrs.push_back(const_cast<LifetimeCaptureByAttr *>(A));
+    }
+  }
+  if (Attrs.empty())
+    return;
+  llvm::StringMap<int> NameIdxMapping = {
+      {"global", LifetimeCaptureByAttr::GLOBAL},
+      {"unknown", LifetimeCaptureByAttr::UNKNOWN}};
   int Idx = 0;
   if (HasImplicitThisParam) {
     NameIdxMapping["this"] = 0;
@@ -3936,9 +3959,7 @@ void Sema::LazyProcessLifetimeCaptureByParams(FunctionDecl *FD) {
         Diag(PVD->getLocation(), diag::err_capture_by_param_uses_reserved_name)
             << (PVD->getName() == "unknown");
   };
-  auto HandleCaptureBy = [&](LifetimeCaptureByAttr *CapturedBy) {
-    if (!CapturedBy)
-      return;
+  for (auto *CapturedBy : Attrs) {
     const auto &Entities = CapturedBy->getArgIdents();
     for (size_t I = 0; I < Entities.size(); ++I) {
       StringRef Name = Entities[I]->getName();
@@ -3956,22 +3977,6 @@ void Sema::LazyProcessLifetimeCaptureByParams(FunctionDecl *FD) {
         DisallowReservedParams(Name);
       CapturedBy->setParamIdx(I, It->second);
     }
-  };
-  for (ParmVarDecl *PVD : FD->parameters())
-    HandleCaptureBy(PVD->getAttr<LifetimeCaptureByAttr>());
-  if (!HasImplicitThisParam)
-    return;
-  TypeSourceInfo *TSI = FD->getTypeSourceInfo();
-  if (!TSI)
-    return;
-  AttributedTypeLoc ATL;
-  for (TypeLoc TL = TSI->getTypeLoc();
-       (ATL = TL.getAsAdjusted<AttributedTypeLoc>());
-       TL = ATL.getModifiedLoc()) {
-    auto *A = ATL.getAttrAs<LifetimeCaptureByAttr>();
-    if (!A)
-      continue;
-    HandleCaptureBy(const_cast<LifetimeCaptureByAttr *>(A));
   }
 }
 
@@ -6753,7 +6758,7 @@ ProcessDeclAttribute(Sema &S, Scope *scope, Decl *D, const ParsedAttr &AL,
     handleCallbackAttr(S, D, AL);
     break;
   case ParsedAttr::AT_LifetimeCaptureBy:
-    HandleLifetimeCaptureByAttr(S, D, AL);
+    handleLifetimeCaptureByAttr(S, D, AL);
     break;
   case ParsedAttr::AT_CalledOnce:
     handleCalledOnceAttr(S, D, AL);

clang/test/SemaCXX/attr-lifetime-capture-by.cpp

@@ -22,6 +22,7 @@ void nonMember(
                                                x7)]], // expected-error {{'lifetime_capture_by' argument references itself}}
     const int &x8 [[clang::lifetime_capture_by(global)]],
     const int &x9 [[clang::lifetime_capture_by(unknown)]],
+    const int &test_memory_leak[[clang::lifetime_capture_by(x1,x2, x3, x4, x5, x6, x7, x8, x9)]],
     const S& u
   )
 {