4 files changed, 42 insertions, 4 deletions

diff --git a/include/clang/AST/ASTContext.h b/include/clang/AST/ASTContext.h
index a6b4254db3..87a12cde2a 100644
--- a/include/clang/AST/ASTContext.h
+++ b/include/clang/AST/ASTContext.h
@@ -1272,6 +1272,15 @@ public:
   TypeSourceInfo *
   getTrivialTypeSourceInfo(QualType T, SourceLocation Loc = SourceLocation());
 
+  /// \brief Add a deallocation callback that will be invoked when the 
+  /// ASTContext is destroyed.
+  ///
+  /// \brief Callback A callback function that will be invoked on destruction.
+  ///
+  /// \brief Data Pointer data that will be provided to the callback function
+  /// when it is called.
+  void AddDeallocation(void (*Callback)(void*), void *Data);
+  
 private:
   ASTContext(const ASTContext&); // DO NOT IMPLEMENT
   void operator=(const ASTContext&); // DO NOT IMPLEMENT
@@ -1286,11 +1295,15 @@ private:
                                   const FieldDecl *Field,
                                   bool OutermostType = false,
                                   bool EncodingProperty = false);
-
+ 
   const ASTRecordLayout &getObjCLayout(const ObjCInterfaceDecl *D,
                                        const ObjCImplementationDecl *Impl);
   
 private:
+  /// \brief A set of deallocations that should be performed when the 
+  /// ASTContext is destroyed.
+  llvm::SmallVector<std::pair<void (*)(void*), void *>, 16> Deallocations;
+                                       
   // FIXME: This currently contains the set of StoredDeclMaps used
   // by DeclContext objects.  This probably should not be in ASTContext,
   // but we include it here so that ASTContext can quickly deallocate them.
diff --git a/include/clang/AST/DeclTemplate.h b/include/clang/AST/DeclTemplate.h
index dc4aecce03..89b43a243c 100644
--- a/include/clang/AST/DeclTemplate.h
+++ b/include/clang/AST/DeclTemplate.h
@@ -460,6 +460,8 @@ public:
   
 /// Declaration of a template function.
 class FunctionTemplateDecl : public TemplateDecl {
+  static void DeallocateCommon(void *Ptr);
+                        
 protected:
   /// \brief Data that is common to all of the declarations of a given
   /// function template.
@@ -1164,6 +1166,8 @@ public:
 
 /// Declaration of a class template.
 class ClassTemplateDecl : public TemplateDecl {
+  static void DeallocateCommon(void *Ptr);
+  
 protected:
   /// \brief Data that is common to all of the declarations of a given
   /// class template.
diff --git a/lib/AST/ASTContext.cpp b/lib/AST/ASTContext.cpp
index 801a1f6391..d6e094e5b4 100644
--- a/lib/AST/ASTContext.cpp
+++ b/lib/AST/ASTContext.cpp
@@ -64,6 +64,12 @@ ASTContext::~ASTContext() {
   // FIXME: Is this the ideal solution?
   ReleaseDeclContextMaps();
 
+  if (!FreeMemory) {
+    // Call all of the deallocation functions.
+    for (unsigned I = 0, N = Deallocations.size(); I != N; ++I)
+      Deallocations[I].first(Deallocations[I].second);
+  }
+  
   // Release all of the memory associated with overridden C++ methods.
   for (llvm::DenseMap<const CXXMethodDecl *, CXXMethodVector>::iterator 
          OM = OverriddenMethods.begin(), OMEnd = OverriddenMethods.end();
@@ -114,6 +120,10 @@ ASTContext::~ASTContext() {
   TUDecl->Destroy(*this);
 }
 
+void ASTContext::AddDeallocation(void (*Callback)(void*), void *Data) {
+  Deallocations.push_back(std::make_pair(Callback, Data));
+}
+
 void
 ASTContext::setExternalSource(llvm::OwningPtr<ExternalASTSource> &Source) {
   ExternalSource.reset(Source.take());
diff --git a/lib/AST/DeclTemplate.cpp b/lib/AST/DeclTemplate.cpp
index 98a724aeee..17bd2217fc 100644
--- a/lib/AST/DeclTemplate.cpp
+++ b/lib/AST/DeclTemplate.cpp
@@ -94,6 +94,10 @@ TemplateDecl::~TemplateDecl() {
 // FunctionTemplateDecl Implementation
 //===----------------------------------------------------------------------===//
 
+void FunctionTemplateDecl::DeallocateCommon(void *Ptr) {
+  static_cast<Common *>(Ptr)->~Common();
+}
+
 FunctionTemplateDecl *FunctionTemplateDecl::Create(ASTContext &C,
                                                    DeclContext *DC,
                                                    SourceLocation L,
@@ -129,8 +133,9 @@ FunctionTemplateDecl::Common *FunctionTemplateDecl::getCommonPtr() {
     First = First->getPreviousDeclaration();
 
   if (First->CommonOrPrev.isNull()) {
-    // FIXME: Allocate with the ASTContext
-    First->CommonOrPrev = new Common;
+    Common *CommonPtr = new (getASTContext()) Common;
+    getASTContext().AddDeallocation(DeallocateCommon, CommonPtr);
+    First->CommonOrPrev = CommonPtr;
   }
   return First->CommonOrPrev.get<Common*>();
 }
@@ -139,6 +144,10 @@ FunctionTemplateDecl::Common *FunctionTemplateDecl::getCommonPtr() {
 // ClassTemplateDecl Implementation
 //===----------------------------------------------------------------------===//
 
+void ClassTemplateDecl::DeallocateCommon(void *Ptr) {
+  static_cast<Common *>(Ptr)->~Common();
+}
+
 ClassTemplateDecl *ClassTemplateDecl::getCanonicalDecl() {
   ClassTemplateDecl *Template = this;
   while (Template->getPreviousDeclaration())
@@ -156,8 +165,10 @@ ClassTemplateDecl *ClassTemplateDecl::Create(ASTContext &C,
   Common *CommonPtr;
   if (PrevDecl)
     CommonPtr = PrevDecl->CommonPtr;
-  else
+  else {
     CommonPtr = new (C) Common;
+    C.AddDeallocation(DeallocateCommon, CommonPtr);
+  }
 
   return new (C) ClassTemplateDecl(DC, L, Name, Params, Decl, PrevDecl,
                                    CommonPtr);