[Clang] Extend lifetime of temporaries in mem-default-init for P2718R0

yronglin · yronglin · commit 8bd2742f6bcc · 2024-04-02T23:49:25.000+08:00
Signed-off-by: yronglin &lt;yronglin777@gmail.com&gt;
diff --git a/clang/lib/CodeGen/CGExpr.cpp b/clang/lib/CodeGen/CGExpr.cpp
@@ -268,9 +268,9 @@ void CodeGenFunction::EmitAnyExprToMem(const Expr *E,
   llvm_unreachable("bad evaluation kind");
 }
 
-static void
-pushTemporaryCleanup(CodeGenFunction &CGF, const MaterializeTemporaryExpr *M,
-                     const Expr *E, Address ReferenceTemporary) {
+void CodeGenFunction::pushTemporaryCleanup(const MaterializeTemporaryExpr *M,
+                                           const Expr *E,
+                                           Address ReferenceTemporary) {
   // Objective-C++ ARC:
   //   If we are binding a reference to a temporary that has ownership, we
   //   need to perform retain/release operations on the temporary.
@@ -305,9 +305,9 @@ pushTemporaryCleanup(CodeGenFunction &CGF, const MaterializeTemporaryExpr *M,
         CleanupKind CleanupKind;
         if (Lifetime == Qualifiers::OCL_Strong) {
           const ValueDecl *VD = M->getExtendingDecl();
-          bool Precise =
-              VD && isa<VarDecl>(VD) && VD->hasAttr<ObjCPreciseLifetimeAttr>();
-          CleanupKind = CGF.getARCCleanupKind();
+          bool Precise = isa_and_nonnull<VarDecl>(VD) &&
+                         VD->hasAttr<ObjCPreciseLifetimeAttr>();
+          CleanupKind = getARCCleanupKind();
           Destroy = Precise ? &CodeGenFunction::destroyARCStrongPrecise
                             : &CodeGenFunction::destroyARCStrongImprecise;
         } else {
@@ -317,13 +317,12 @@ pushTemporaryCleanup(CodeGenFunction &CGF, const MaterializeTemporaryExpr *M,
           Destroy = &CodeGenFunction::destroyARCWeak;
         }
         if (Duration == SD_FullExpression)
-          CGF.pushDestroy(CleanupKind, ReferenceTemporary,
-                          M->getType(), *Destroy,
-                          CleanupKind & EHCleanup);
+          pushDestroy(CleanupKind, ReferenceTemporary, M->getType(), *Destroy,
+                      CleanupKind & EHCleanup);
         else
-          CGF.pushLifetimeExtendedDestroy(CleanupKind, ReferenceTemporary,
-                                          M->getType(),
-                                          *Destroy, CleanupKind & EHCleanup);
+          pushLifetimeExtendedDestroy(CleanupKind, ReferenceTemporary,
+                                      M->getType(), *Destroy,
+                                      CleanupKind & EHCleanup);
         return;
 
       case SD_Dynamic:
@@ -352,32 +351,31 @@ pushTemporaryCleanup(CodeGenFunction &CGF, const MaterializeTemporaryExpr *M,
     llvm::FunctionCallee CleanupFn;
     llvm::Constant *CleanupArg;
     if (E->getType()->isArrayType()) {
-      CleanupFn = CodeGenFunction(CGF.CGM).generateDestroyHelper(
-          ReferenceTemporary, E->getType(),
-          CodeGenFunction::destroyCXXObject, CGF.getLangOpts().Exceptions,
+      CleanupFn = CodeGenFunction(CGM).generateDestroyHelper(
+          ReferenceTemporary, E->getType(), CodeGenFunction::destroyCXXObject,
+          getLangOpts().Exceptions,
           dyn_cast_or_null<VarDecl>(M->getExtendingDecl()));
-      CleanupArg = llvm::Constant::getNullValue(CGF.Int8PtrTy);
+      CleanupArg = llvm::Constant::getNullValue(Int8PtrTy);
     } else {
-      CleanupFn = CGF.CGM.getAddrAndTypeOfCXXStructor(
+      CleanupFn = CGM.getAddrAndTypeOfCXXStructor(
           GlobalDecl(ReferenceTemporaryDtor, Dtor_Complete));
-      CleanupArg = cast<llvm::Constant>(ReferenceTemporary.emitRawPointer(CGF));
+      CleanupArg =
+          cast<llvm::Constant>(ReferenceTemporary.emitRawPointer(*this));
     }
-    CGF.CGM.getCXXABI().registerGlobalDtor(
-        CGF, *cast<VarDecl>(M->getExtendingDecl()), CleanupFn, CleanupArg);
+    CGM.getCXXABI().registerGlobalDtor(
+        *this, *cast<VarDecl>(M->getExtendingDecl()), CleanupFn, CleanupArg);
     break;
   }
 
   case SD_FullExpression:
-    CGF.pushDestroy(NormalAndEHCleanup, ReferenceTemporary, E->getType(),
-                    CodeGenFunction::destroyCXXObject,
-                    CGF.getLangOpts().Exceptions);
+    pushDestroy(NormalAndEHCleanup, ReferenceTemporary, E->getType(),
+                CodeGenFunction::destroyCXXObject, getLangOpts().Exceptions);
     break;
 
   case SD_Automatic:
-    CGF.pushLifetimeExtendedDestroy(NormalAndEHCleanup,
-                                    ReferenceTemporary, E->getType(),
-                                    CodeGenFunction::destroyCXXObject,
-                                    CGF.getLangOpts().Exceptions);
+    pushLifetimeExtendedDestroy(NormalAndEHCleanup, ReferenceTemporary,
+                                E->getType(), CodeGenFunction::destroyCXXObject,
+                                getLangOpts().Exceptions);
     break;
 
   case SD_Dynamic:
@@ -484,7 +482,7 @@ EmitMaterializeTemporaryExpr(const MaterializeTemporaryExpr *M) {
     }
     }
 
-    pushTemporaryCleanup(*this, M, E, Object);
+    pushTemporaryCleanup(M, E, Object);
     return RefTempDst;
   }
 
@@ -573,7 +571,13 @@ EmitMaterializeTemporaryExpr(const MaterializeTemporaryExpr *M) {
     }
     EmitAnyExprToMem(E, Object, Qualifiers(), /*IsInit*/true);
   }
-  pushTemporaryCleanup(*this, M, E, Object);
+
+  // If this temporary extended by for-range variable, delay to emitting
+  // cleanup.
+  if (!CurLexicalScope->isExtendedByForRangeVar(M))
+    CurLexicalScope->addForRangeInitTemp(M, Object);
+  else
+    pushTemporaryCleanup(M, E, Object);
 
   // Perform derived-to-base casts and/or field accesses, to get from the
   // temporary object we created (and, potentially, for which we extended
diff --git a/clang/lib/CodeGen/CGStmt.cpp b/clang/lib/CodeGen/CGStmt.cpp
@@ -1230,11 +1230,23 @@ CodeGenFunction::EmitCXXForRangeStmt(const CXXForRangeStmt &S,
   JumpDest LoopExit = getJumpDestInCurrentScope("for.end");
 
   LexicalScope ForScope(*this, S.getSourceRange());
+  ForScope.setForRangeVar(S.getLoopVariable());
 
   // Evaluate the first pieces before the loop.
   if (S.getInit())
     EmitStmt(S.getInit());
   EmitStmt(S.getRangeStmt());
+
+  // Emit cleanup for tempories in for-range-init expression.
+  {
+    RunCleanupsScope Scope(*this);
+    auto LifetimeExtendTemps = ForScope.getForRangeInitTemps();
+    for (const auto &Temp : LifetimeExtendTemps) {
+      auto [M, Object] = Temp;
+      pushTemporaryCleanup(M, M->getSubExpr(), Object);
+    }
+  }
+
   EmitStmt(S.getBeginStmt());
   EmitStmt(S.getEndStmt());
 
diff --git a/clang/lib/CodeGen/CodeGenFunction.h b/clang/lib/CodeGen/CodeGenFunction.h
@@ -875,6 +875,9 @@ class CodeGenFunction : public CodeGenTypeCache {
       new (Buffer + sizeof(Header) + sizeof(T)) RawAddress(ActiveFlag);
   }
 
+  void pushTemporaryCleanup(const MaterializeTemporaryExpr *M, const Expr *E,
+                            Address ReferenceTemporary);
+
   /// Set up the last cleanup that was pushed as a conditional
   /// full-expression cleanup.
   void initFullExprCleanup() {
@@ -982,11 +985,24 @@ class CodeGenFunction : public CodeGenTypeCache {
   EHScopeStack::stable_iterator CurrentCleanupScopeDepth =
       EHScopeStack::stable_end();
 
+  struct ForRangeInitLifetimeExtendTemporary {
+    const MaterializeTemporaryExpr *M;
+    RawAddress Object;
+  };
+
   class LexicalScope : public RunCleanupsScope {
     SourceRange Range;
     SmallVector<const LabelDecl*, 4> Labels;
+    SmallVector<ForRangeInitLifetimeExtendTemporary, 4> ForRangeInitTemps;
     LexicalScope *ParentScope;
 
+    // This will be set to `__range` variable when we emitting a
+    // CXXForRangeStmt. It was used to check whether we are emitting a
+    // materialized temporary which in for-range-init and lifetime-extended by
+    // __range var. If so, the codegen of cleanup for that temporary object
+    // needs to be delayed.
+    const VarDecl *ForRangeVar = nullptr;
+
     LexicalScope(const LexicalScope &) = delete;
     void operator=(const LexicalScope &) = delete;
 
@@ -1004,6 +1020,21 @@ class CodeGenFunction : public CodeGenTypeCache {
       Labels.push_back(label);
     }
 
+    void addForRangeInitTemp(const MaterializeTemporaryExpr *M,
+                             RawAddress Object) {
+      assert(PerformCleanup && "adding temps to dead scope?");
+      ForRangeInitTemps.push_back({M, Object});
+    }
+
+    ArrayRef<ForRangeInitLifetimeExtendTemporary> getForRangeInitTemps() const {
+      return ForRangeInitTemps;
+    }
+
+    void setForRangeVar(const VarDecl *Var) { ForRangeVar = Var; }
+    bool isExtendedByForRangeVar(const MaterializeTemporaryExpr *M) const {
+      return M && ForRangeVar && M->getExtendingDecl() == ForRangeVar;
+    }
+
     /// Exit this cleanup scope, emitting any accumulated
     /// cleanups.
     ~LexicalScope() {
diff --git a/clang/lib/Sema/SemaExpr.cpp b/clang/lib/Sema/SemaExpr.cpp
@@ -6376,7 +6376,7 @@ ExprResult Sema::BuildCXXDefaultInitExpr(SourceLocation Loc, FieldDecl *Field) {
   Expr *Init = nullptr;
 
   bool NestedDefaultChecking = isCheckingDefaultArgumentOrInitializer();
-
+  bool InLifetimeExtendingContext = isInLifetimeExtendingContext();
   EnterExpressionEvaluationContext EvalContext(
       *this, ExpressionEvaluationContext::PotentiallyEvaluated, Field);
 
@@ -6411,12 +6411,15 @@ ExprResult Sema::BuildCXXDefaultInitExpr(SourceLocation Loc, FieldDecl *Field) {
   ImmediateCallVisitor V(getASTContext());
   if (!NestedDefaultChecking)
     V.TraverseDecl(Field);
-  if (V.HasImmediateCalls) {
+  if (V.HasImmediateCalls || InLifetimeExtendingContext) {
     ExprEvalContexts.back().DelayedDefaultInitializationContext = {Loc, Field,
                                                                    CurContext};
     ExprEvalContexts.back().IsCurrentlyCheckingDefaultArgumentOrInitializer =
         NestedDefaultChecking;
-
+    // Pass down lifetime extending flag, and collect temporaries in
+    // CreateMaterializeTemporaryExpr when we rewrite the call argument.
+    keepInLifetimeExtendingContext();
+    keepInMaterializeTemporaryObjectContext();
     EnsureImmediateInvocationInDefaultArgs Immediate(*this);
     ExprResult Res;
     runWithSufficientStackSpace(Loc, [&] {
diff --git a/clang/lib/Sema/SemaInit.cpp b/clang/lib/Sema/SemaInit.cpp
@@ -8208,9 +8208,10 @@ void Sema::checkInitializerLifetime(const InitializedEntity &Entity,
         // FIXME: Properly handle this situation. Perhaps the easiest approach
         // would be to clone the initializer expression on each use that would
         // lifetime extend its temporaries.
-        Diag(DiagLoc, diag::warn_unsupported_lifetime_extension)
-            << RK << DiagRange;
-        break;
+        // Diag(DiagLoc, diag::warn_unsupported_lifetime_extension)
+        //     << RK << DiagRange;
+        // break;
+        return true;
 
       case PathLifetimeKind::NoExtend:
         // If the path goes through the initialization of a variable or field,
