[clang] [clang] Infer lifetime_capture_by for map's subscript operator. (PR #118078)

[Haojian Wu via cfe-commits]

https://github.com/hokein created https://github.com/llvm/llvm-project/pull/118078

None

>From 18d56e847901653de9f84e256f30ec0e16ebca96 Mon Sep 17 00:00:00 2001
From: Haojian Wu <hokein.wu at gmail.com>
Date: Fri, 29 Nov 2024 11:07:56 +0100
Subject: [PATCH] [clang] Infer lifetime_capture_by for map's subscript
 operator.

---
 clang/lib/Sema/SemaAttr.cpp                 | 59 +++++++++++++--------
 clang/test/AST/attr-lifetime-capture-by.cpp | 16 ++++++
 2 files changed, 53 insertions(+), 22 deletions(-)

diff --git a/clang/lib/Sema/SemaAttr.cpp b/clang/lib/Sema/SemaAttr.cpp
index d3cf42251be2e7..3279f4109a9c02 100644
--- a/clang/lib/Sema/SemaAttr.cpp
+++ b/clang/lib/Sema/SemaAttr.cpp
@@ -270,34 +270,49 @@ void Sema::inferLifetimeBoundAttribute(FunctionDecl *FD) {
 }
 
 void Sema::inferLifetimeCaptureByAttribute(FunctionDecl *FD) {
-  if (!FD)
+  auto *MD = dyn_cast_if_present<CXXMethodDecl>(FD);
+  if (!MD || !MD->getParent()->isInStdNamespace())
     return;
-  auto *MD = dyn_cast<CXXMethodDecl>(FD);
-  if (!MD || !MD->getIdentifier() || !MD->getParent()->isInStdNamespace())
+  auto Annotate = [this](const FunctionDecl *MD) {
+    // Do not infer if any parameter is explicitly annotated.
+    for (ParmVarDecl *PVD : MD->parameters())
+      if (PVD->hasAttr<LifetimeCaptureByAttr>())
+        return;
+    for (ParmVarDecl *PVD : MD->parameters()) {
+      // Methods in standard containers that capture values typically accept
+      // reference-type parameters, e.g., `void push_back(const T& value)`.
+      // We only apply the lifetime_capture_by attribute to parameters of
+      // pointer-like reference types (`const T&`, `T&&`).
+      if (PVD->getType()->isReferenceType() &&
+          sema::isPointerLikeType(PVD->getType().getNonReferenceType())) {
+        int CaptureByThis[] = {LifetimeCaptureByAttr::THIS};
+        PVD->addAttr(
+            LifetimeCaptureByAttr::CreateImplicit(Context, CaptureByThis, 1));
+      }
+    }
+  };
+
+  if (!MD->getIdentifier()) {
+    static const llvm::StringSet<> MapLikeContainer{
+        "map",
+        "multimap",
+        "unordered_map",
+        "unordered_multimap",
+    };
+    // Infer for the map's operator []:
+    //    std::map<string_view, ...> m;
+    //    m[ReturnString(..)] = ...; // !dangling references in m.
+    if (MD->getOverloadedOperator() != OO_Subscript ||
+        !MapLikeContainer.contains(MD->getParent()->getName()))
+      return;
+    Annotate(MD);
     return;
-  // FIXME: Infer for operator[] for map-like containers. For example:
-  //    std::map<string_view, ...> m;
-  //    m[ReturnString(..)] = ...;
+  }
   static const llvm::StringSet<> CapturingMethods{"insert", "push",
                                                   "push_front", "push_back"};
   if (!CapturingMethods.contains(MD->getName()))
     return;
-  // Do not infer if any parameter is explicitly annotated.
-  for (ParmVarDecl *PVD : MD->parameters())
-    if (PVD->hasAttr<LifetimeCaptureByAttr>())
-      return;
-  for (ParmVarDecl *PVD : MD->parameters()) {
-    // Methods in standard containers that capture values typically accept
-    // reference-type parameters, e.g., `void push_back(const T& value)`.
-    // We only apply the lifetime_capture_by attribute to parameters of
-    // pointer-like reference types (`const T&`, `T&&`).
-    if (PVD->getType()->isReferenceType() &&
-        sema::isPointerLikeType(PVD->getType().getNonReferenceType())) {
-      int CaptureByThis[] = {LifetimeCaptureByAttr::THIS};
-      PVD->addAttr(
-          LifetimeCaptureByAttr::CreateImplicit(Context, CaptureByThis, 1));
-    }
-  }
+  Annotate(MD);
 }
 
 void Sema::inferNullableClassAttribute(CXXRecordDecl *CRD) {
diff --git a/clang/test/AST/attr-lifetime-capture-by.cpp b/clang/test/AST/attr-lifetime-capture-by.cpp
index debad9b7204d72..71b348dac764bc 100644
--- a/clang/test/AST/attr-lifetime-capture-by.cpp
+++ b/clang/test/AST/attr-lifetime-capture-by.cpp
@@ -30,6 +30,12 @@ struct vector {
 
   void insert(iterator, T&&);
 };
+
+template <typename Key, typename Value>
+struct map {
+  Value& operator[](Key&& p);
+  Value& operator[](const Key& p);
+};
 } // namespace std
 
 // CHECK-NOT:   LifetimeCaptureByAttr
@@ -99,3 +105,13 @@ std::vector<int> ints;
 // CHECK-NEXT:           ParmVarDecl {{.*}} 'iterator'
 // CHECK-NEXT:           ParmVarDecl {{.*}} 'int &&'
 // CHECK-NOT:   LifetimeCaptureByAttr
+
+std::map<View, int> map;
+// CHECK:   ClassTemplateSpecializationDecl {{.*}} struct map definition implicit_instantiation
+
+// CHECK:       CXXMethodDecl {{.*}} operator[] 'int &(View &&)' implicit_instantiation
+// CHECK-NEXT:           ParmVarDecl {{.*}} p 'View &&'
+// CHECK-NEXT:               LifetimeCaptureByAttr {{.*}} Implicit
+// CHECK:       CXXMethodDecl {{.*}} operator[] 'int &(const View &)' implicit_instantiation
+// CHECK-NEXT:           ParmVarDecl {{.*}} p 'const View &'
+// CHECK-NEXT:               LifetimeCaptureByAttr {{.*}} Implicit