[clang] [CodeGen] Fix cleanup attribute for C89 for-loop init variables (PR #156643)

[Jongmyeong Choi via cfe-commits]

https://github.com/jongmyeong-choi created https://github.com/llvm/llvm-project/pull/156643

In C89, for-init variables have function scope, so cleanup should occur at function exit, not loop exit. This implements deferred cleanup registration for C89 mode while preserving C99+ behavior.

Fixes #154624

>From 471273e66583dbec3f145c7f5723a4d5b524d15b Mon Sep 17 00:00:00 2001
From: Jongmyeong Choi <cheesechoi at gmail.com>
Date: Wed, 3 Sep 2025 18:41:46 +0900
Subject: [PATCH] [CodeGen] Fix cleanup attribute for C89 for-loop init
 variables

In C89, for-init variables have function scope, so cleanup should occur
at function exit, not loop exit. This implements deferred cleanup
registration for C89 mode while preserving C99+ behavior.

Fixes #154624
---
 clang/lib/CodeGen/CGDecl.cpp          | 32 ++++++++++++++++++++++++---
 clang/lib/CodeGen/CodeGenFunction.cpp |  6 ++++-
 clang/lib/CodeGen/CodeGenFunction.h   | 22 ++++++++++++++++++
 clang/test/CodeGen/attr-cleanup.c     | 25 +++++++++++++++++++++
 4 files changed, 81 insertions(+), 4 deletions(-)

diff --git a/clang/lib/CodeGen/CGDecl.cpp b/clang/lib/CodeGen/CGDecl.cpp
index 8a1675848e13c..e4a691a161026 100644
--- a/clang/lib/CodeGen/CGDecl.cpp
+++ b/clang/lib/CodeGen/CGDecl.cpp
@@ -2230,9 +2230,16 @@ void CodeGenFunction::EmitAutoVarCleanups(const AutoVarEmission &emission) {
     llvm::Constant *F = CGM.GetAddrOfFunction(FD);
     assert(F && "Could not find function!");
 
-    const CGFunctionInfo &Info = CGM.getTypes().arrangeFunctionDeclaration(FD);
-    EHStack.pushCleanup<CallCleanupFunction>(NormalAndEHCleanup, F, &Info, &D,
-                                             CA);
+    // Check if we're in C89 mode and should defer cleanup to function scope
+    bool isC89Mode = !getLangOpts().C99 && !getLangOpts().CPlusPlus;
+    if (isC89Mode) {
+      const CGFunctionInfo &Info = CGM.getTypes().arrangeFunctionDeclaration(FD);
+      addDeferredFunctionCleanup(F, &Info, &D, CA);
+    } else {
+      const CGFunctionInfo &Info = CGM.getTypes().arrangeFunctionDeclaration(FD);
+      EHStack.pushCleanup<CallCleanupFunction>(NormalAndEHCleanup, F, &Info, &D,
+                                               CA);
+    }
   }
 
   // If this is a block variable, call _Block_object_destroy
@@ -2963,3 +2970,22 @@ CodeGenModule::getOMPAllocateAlignment(const VarDecl *VD) {
   }
   return std::nullopt;
 }
+
+void CodeGenFunction::addDeferredFunctionCleanup(llvm::Constant *F,
+                                                 const CGFunctionInfo *Info,
+                                                 const VarDecl *Var,
+                                                 const CleanupAttr *Attribute) {
+  DeferredFunctionCleanups.emplace_back(F, Info, Var, Attribute);
+}
+
+void CodeGenFunction::processDeferredFunctionCleanups() {
+  // Process all deferred cleanups at function exit
+  for (const auto &cleanup : DeferredFunctionCleanups) {
+    EHStack.pushCleanup<CallCleanupFunction>(NormalAndEHCleanup, 
+                                             cleanup.CleanupFn, 
+                                             cleanup.FnInfo, 
+                                             cleanup.Var, 
+                                             cleanup.Attribute);
+  }
+  DeferredFunctionCleanups.clear();
+}
diff --git a/clang/lib/CodeGen/CodeGenFunction.cpp b/clang/lib/CodeGen/CodeGenFunction.cpp
index 652fe672f15e3..0ea33af47343a 100644
--- a/clang/lib/CodeGen/CodeGenFunction.cpp
+++ b/clang/lib/CodeGen/CodeGenFunction.cpp
@@ -405,6 +405,10 @@ void CodeGenFunction::FinishFunction(SourceLocation EndLoc) {
   // parameters.  Do this in whatever block we're currently in; it's
   // important to do this before we enter the return block or return
   // edges will be *really* confused.
+  
+  // Process deferred function cleanups before checking for regular cleanups
+  processDeferredFunctionCleanups();
+  
   bool HasCleanups = EHStack.stable_begin() != PrologueCleanupDepth;
   bool HasOnlyNoopCleanups =
       HasCleanups && EHStack.containsOnlyNoopCleanups(PrologueCleanupDepth);
@@ -422,7 +426,7 @@ void CodeGenFunction::FinishFunction(SourceLocation EndLoc) {
         // fall back to an artificial location if needed.
         OAL = ApplyDebugLocation::CreateDefaultArtificial(*this, EndLoc);
     }
-
+    
     PopCleanupBlocks(PrologueCleanupDepth);
   }
 
diff --git a/clang/lib/CodeGen/CodeGenFunction.h b/clang/lib/CodeGen/CodeGenFunction.h
index fc65199a0f154..19ea1e7e34052 100644
--- a/clang/lib/CodeGen/CodeGenFunction.h
+++ b/clang/lib/CodeGen/CodeGenFunction.h
@@ -392,6 +392,21 @@ class CodeGenFunction : public CodeGenTypeCache {
   /// cleanups associated with the parameters.
   EHScopeStack::stable_iterator PrologueCleanupDepth;
 
+  /// Structure for deferred function-level cleanups (e.g., C89 for-init cleanup variables)
+  struct DeferredCleanupInfo {
+    llvm::Constant *CleanupFn;
+    const CGFunctionInfo *FnInfo;
+    const VarDecl *Var;
+    const CleanupAttr *Attribute;
+    
+    DeferredCleanupInfo(llvm::Constant *F, const CGFunctionInfo *Info, 
+                        const VarDecl *V, const CleanupAttr *A)
+      : CleanupFn(F), FnInfo(Info), Var(V), Attribute(A) {}
+  };
+  
+  /// List of cleanups that should be registered at function exit instead of current scope
+  SmallVector<DeferredCleanupInfo, 4> DeferredFunctionCleanups;
+
   /// ReturnBlock - Unified return block.
   JumpDest ReturnBlock;
 
@@ -3470,6 +3485,13 @@ class CodeGenFunction : public CodeGenTypeCache {
   void EmitAutoVarCleanups(const AutoVarEmission &emission);
   void emitAutoVarTypeCleanup(const AutoVarEmission &emission,
                               QualType::DestructionKind dtorKind);
+  
+  /// Add a cleanup to be deferred until function exit (for C89 for-init variables)
+  void addDeferredFunctionCleanup(llvm::Constant *CleanupFn, const CGFunctionInfo *FnInfo,
+                                  const VarDecl *Var, const CleanupAttr *Attribute);
+  
+  /// Process all deferred function-level cleanups
+  void processDeferredFunctionCleanups();
 
   void MaybeEmitDeferredVarDeclInit(const VarDecl *var);
 
diff --git a/clang/test/CodeGen/attr-cleanup.c b/clang/test/CodeGen/attr-cleanup.c
index 755ede86c1382..c723014973332 100644
--- a/clang/test/CodeGen/attr-cleanup.c
+++ b/clang/test/CodeGen/attr-cleanup.c
@@ -5,3 +5,28 @@ void g(void) {
   __attribute__((cleanup(f))) void *g;
 }
 
+// Test for cleanup in for-loop initialization (PR #154624)
+// RUN: %clang_cc1 -std=c89 -emit-llvm %s -o - | FileCheck %s --check-prefix=C89
+// RUN: %clang_cc1 -std=c99 -emit-llvm %s -o - | FileCheck %s --check-prefix=C99
+
+void cleaner(int *p);
+
+// C89-LABEL: define{{.*}} void @test_for_loop_cleanup()  
+// C99-LABEL: define{{.*}} void @test_for_loop_cleanup()
+void test_for_loop_cleanup(void) {
+  for (__attribute__((cleanup(cleaner))) int i = 42; 0;)
+    ;
+  
+#ifndef __STDC_VERSION__
+  if (i > 40) {
+    // do something
+  }
+#endif
+}
+
+// C89: if.end:
+// C89-NEXT: call void @cleaner(ptr noundef %i)
+// C89-NEXT: ret void
+
+// C99: for.cond.cleanup:
+// C99-NEXT: call void @cleaner(ptr noundef %i)