[clang] [clang][AST][ASTMerge] prevent AST nodes from being deallocated early (PR #73096)

[Qizhi Hu via cfe-commits]

https://github.com/jcsxky created https://github.com/llvm/llvm-project/pull/73096

This patch aims to fix [error in ast-merge to new ast file](https://github.com/llvm/llvm-project/issues/72783).
`ASTUnit` is put in `for` body and AST nodes would be deallocated by allocator. Using these nodes later would lead to use after free bug.
Debug the [issue](https://github.com/llvm/llvm-project/issues/72783) can prove it. Address interval (local from 0x3a9a00 to 0x3aaa00) allocated by allocator contains a `IdentifierInfo` variable (local address:0x3aa190) whose address is freed early.
As system header like stdio.h or math.h can't be put into test, it's hard to add testcase. Could anyone give me some guidance? Thanks in advance!

>From 0c3b47b351eb90495d635fd37845d5c7da869b6e Mon Sep 17 00:00:00 2001
From: huqizhi <huqizhi at feysh.com>
Date: Wed, 22 Nov 2023 16:41:42 +0800
Subject: [PATCH] [clang][AST][ASTMerge] prevent AST nodes from being
 deallocated early

---
 clang/lib/Frontend/ASTMerge.cpp | 13 ++++++++-----
 1 file changed, 8 insertions(+), 5 deletions(-)

diff --git a/clang/lib/Frontend/ASTMerge.cpp b/clang/lib/Frontend/ASTMerge.cpp
index 057ea4fd5bb3518..94706ade4c876a2 100644
--- a/clang/lib/Frontend/ASTMerge.cpp
+++ b/clang/lib/Frontend/ASTMerge.cpp
@@ -5,14 +5,15 @@
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
-#include "clang/Frontend/ASTUnit.h"
 #include "clang/AST/ASTContext.h"
 #include "clang/AST/ASTDiagnostic.h"
 #include "clang/AST/ASTImporter.h"
 #include "clang/AST/ASTImporterSharedState.h"
 #include "clang/Basic/Diagnostic.h"
+#include "clang/Frontend/ASTUnit.h"
 #include "clang/Frontend/CompilerInstance.h"
 #include "clang/Frontend/FrontendActions.h"
+#include "llvm/ADT/SmallVector.h"
 
 using namespace clang;
 
@@ -40,24 +41,26 @@ void ASTMergeAction::ExecuteAction() {
       DiagIDs(CI.getDiagnostics().getDiagnosticIDs());
   auto SharedState = std::make_shared<ASTImporterSharedState>(
       *CI.getASTContext().getTranslationUnitDecl());
+  llvm::SmallVector<std::unique_ptr<ASTUnit>> Units(ASTFiles.size());
   for (unsigned I = 0, N = ASTFiles.size(); I != N; ++I) {
     IntrusiveRefCntPtr<DiagnosticsEngine>
         Diags(new DiagnosticsEngine(DiagIDs, &CI.getDiagnosticOpts(),
                                     new ForwardingDiagnosticConsumer(
                                           *CI.getDiagnostics().getClient()),
                                     /*ShouldOwnClient=*/true));
-    std::unique_ptr<ASTUnit> Unit = ASTUnit::LoadFromASTFile(
+    Units[I] = ASTUnit::LoadFromASTFile(
         ASTFiles[I], CI.getPCHContainerReader(), ASTUnit::LoadEverything, Diags,
         CI.getFileSystemOpts(), CI.getHeaderSearchOptsPtr(), false);
 
-    if (!Unit)
+    if (!Units[I])
       continue;
 
     ASTImporter Importer(CI.getASTContext(), CI.getFileManager(),
-                         Unit->getASTContext(), Unit->getFileManager(),
+                         Units[I]->getASTContext(), Units[I]->getFileManager(),
                          /*MinimalImport=*/false, SharedState);
 
-    TranslationUnitDecl *TU = Unit->getASTContext().getTranslationUnitDecl();
+    TranslationUnitDecl *TU =
+        Units[I]->getASTContext().getTranslationUnitDecl();
     for (auto *D : TU->decls()) {
       // Don't re-import __va_list_tag, __builtin_va_list.
       if (const auto *ND = dyn_cast<NamedDecl>(D))