4 files changed, 36 insertions, 56 deletions

diff --git a/lib/Sema/Sema.h b/lib/Sema/Sema.h
index df6f18f95b..87ebea72f2 100644
--- a/lib/Sema/Sema.h
+++ b/lib/Sema/Sema.h
@@ -2172,11 +2172,11 @@ public:
   void MarkBaseAndMemberDestructorsReferenced(CXXDestructorDecl *Destructor);
 
   /// ClassesWithUnmarkedVirtualMembers - Contains record decls whose virtual
-  /// members might need to be marked as referenced. This is either done when
-  /// the key function definition is emitted (this is handled by by 
-  /// MaybeMarkVirtualMembersReferenced), or at the end of the translation unit
-  /// (done by ProcessPendingClassesWithUnmarkedVirtualMembers).
-  std::map<CXXRecordDecl *, SourceLocation> ClassesWithUnmarkedVirtualMembers;
+  /// members need to be marked as referenced at the end of the translation
+  /// unit. It will contain polymorphic classes that do not have a key
+  /// function or have a key function that has been defined.
+  llvm::SmallVector<std::pair<CXXRecordDecl *, SourceLocation>, 4>
+    ClassesWithUnmarkedVirtualMembers;
 
   /// MaybeMarkVirtualMembersReferenced - If the passed in method is the
   /// key function of the record decl, will mark virtual member functions as 
diff --git a/lib/Sema/SemaDecl.cpp b/lib/Sema/SemaDecl.cpp
index 0fc61e3b58..383893cc2f 100644
--- a/lib/Sema/SemaDecl.cpp
+++ b/lib/Sema/SemaDecl.cpp
@@ -5098,6 +5098,16 @@ void Sema::ActOnTagFinishDefinition(Scope *S, DeclPtrTy TagD,
   // Exit this scope of this tag's definition.
   PopDeclContext();
 
+  // If this is a polymorphic C++ class without a key function, we'll
+  // have to mark all of the virtual members to allow emission of a vtable
+  // in this translation unit.
+  if (CXXRecordDecl *Record = dyn_cast<CXXRecordDecl>(Tag)) {
+    if (!Record->isDependentContext() && Record->isDynamicClass() &&
+        !Context.getKeyFunction(Record))
+      ClassesWithUnmarkedVirtualMembers.push_back(std::make_pair(Record, 
+                                                                 RBraceLoc));
+  }
+
   // Notify the consumer that we've defined a tag.
   Consumer.HandleTagDeclDefinition(Tag);
 }
diff --git a/lib/Sema/SemaDeclCXX.cpp b/lib/Sema/SemaDeclCXX.cpp
index ecf95d7eb6..623bf8ae7e 100644
--- a/lib/Sema/SemaDeclCXX.cpp
+++ b/lib/Sema/SemaDeclCXX.cpp
@@ -2054,7 +2054,7 @@ void Sema::CheckCompletedCXXClass(CXXRecordDecl *Record) {
 
   if (!Record->isDependentType())
     AddImplicitlyDeclaredMembersToClass(Record);
-
+  
   if (Record->isInvalidDecl())
     return;
 
@@ -5693,67 +5693,30 @@ void Sema::MaybeMarkVirtualMembersReferenced(SourceLocation Loc,
   if (!RD->isDynamicClass())
     return;
 
-  if (!MD->isOutOfLine()) {
-    // The only inline functions we care about are constructors. We also defer
-    // marking the virtual members as referenced until we've reached the end
-    // of the translation unit. We do this because we need to know the key
-    // function of the class in order to determine the key function.
-    if (isa<CXXConstructorDecl>(MD))
-      ClassesWithUnmarkedVirtualMembers.insert(std::make_pair(RD, Loc));
+  // Only out-of-line definitions matter.
+  if (!MD->isOutOfLine())
+    return;
+  
+  const CXXMethodDecl *KeyFunction = Context.getKeyFunction(RD);
+  if (!KeyFunction || KeyFunction->getCanonicalDecl() != MD->getCanonicalDecl())
     return;
-  }
-
-  switch (RD->getTemplateSpecializationKind()) {
-  case TSK_Undeclared:
-  case TSK_ExplicitSpecialization: {
-    const CXXMethodDecl *KeyFunction = Context.getKeyFunction(RD);
-
-    if (!KeyFunction) {
-      // This record does not have a key function, so we assume that the vtable
-      // will be emitted when it's used by the constructor.
-      if (!isa<CXXConstructorDecl>(MD))
-        return;
-    } else if (KeyFunction->getCanonicalDecl() != MD->getCanonicalDecl()) {
-      // We don't have the right key function.
-      return;
-    }
-    break;
-  }
-
-  case TSK_ImplicitInstantiation:
-  case TSK_ExplicitInstantiationDeclaration:
-  case TSK_ExplicitInstantiationDefinition:
-    // Always mark the virtual members of an instantiated template.
-    break;
-  }
 
-  // Mark the members as referenced.
-  MarkVirtualMembersReferenced(Loc, RD);
-  ClassesWithUnmarkedVirtualMembers.erase(RD);
+  // We will need to mark all of the virtual members as referenced to build the
+  // vtable.
+  ClassesWithUnmarkedVirtualMembers.push_back(std::make_pair(RD, Loc));
 }
 
 bool Sema::ProcessPendingClassesWithUnmarkedVirtualMembers() {
   if (ClassesWithUnmarkedVirtualMembers.empty())
     return false;
   
-  for (std::map<CXXRecordDecl *, SourceLocation>::iterator i = 
-       ClassesWithUnmarkedVirtualMembers.begin(), 
-       e = ClassesWithUnmarkedVirtualMembers.end(); i != e; ++i) {
-    CXXRecordDecl *RD = i->first;
-    
-    const CXXMethodDecl *KeyFunction = Context.getKeyFunction(RD);
-    if (KeyFunction) {
-      // We know that the class has a key function. If the key function was
-      // declared in this translation unit, then it the class decl would not 
-      // have been in the ClassesWithUnmarkedVirtualMembers map.
-      continue;
-    }
-    
-    SourceLocation Loc = i->second;
+  while (!ClassesWithUnmarkedVirtualMembers.empty()) {
+    CXXRecordDecl *RD = ClassesWithUnmarkedVirtualMembers.back().first;
+    SourceLocation Loc = ClassesWithUnmarkedVirtualMembers.back().second;
+    ClassesWithUnmarkedVirtualMembers.pop_back();
     MarkVirtualMembersReferenced(Loc, RD);
   }
   
-  ClassesWithUnmarkedVirtualMembers.clear();
   return true;
 }
 
diff --git a/lib/Sema/SemaTemplateInstantiate.cpp b/lib/Sema/SemaTemplateInstantiate.cpp
index d974f89bc6..2db0deb509 100644
--- a/lib/Sema/SemaTemplateInstantiate.cpp
+++ b/lib/Sema/SemaTemplateInstantiate.cpp
@@ -1060,6 +1060,13 @@ Sema::InstantiateClass(SourceLocation PointOfInstantiation,
   // Exit the scope of this instantiation.
   CurContext = PreviousContext;
 
+  // If this is a polymorphic C++ class without a key function, we'll
+  // have to mark all of the virtual members to allow emission of a vtable
+  // in this translation unit.
+  if (Instantiation->isDynamicClass() && !Context.getKeyFunction(Instantiation))
+      ClassesWithUnmarkedVirtualMembers.push_back(std::make_pair(Instantiation,
+                                                       PointOfInstantiation));
+
   if (!Invalid)
     Consumer.HandleTagDeclDefinition(Instantiation);