Add support for remembering origins to ExternalASTMerger

ExternalASTMerger has hitherto relied on being able to look up 
any Decl through its named DeclContext chain. This works for 
many cases, but causes problems for function-local structs, 
which cannot be looked up in their containing FunctionDecl. An
example case is

void f() {
  { struct S { int a; }; }
  { struct S { bool b; }; }
}

It is not possible to lookup either of the two Ses individually 
(or even to provide enough information to disambiguate) after 
parsing is over; and there is typically no need to, since they 
are invisible to the outside world.

However, ExternalASTMerger needs to be able to complete either 
S on demand. This led to an XFAIL on test/Import/local-struct, 
which this patch removes. The way the patch works is:

It defines a new data structure, ExternalASTMerger::OriginMap,
which clients are expected to maintain (default-constructing 
if the origin does not have an ExternalASTMerger servicing it)
As DeclContexts are imported, if they cannot be looked up by 
name they are placed in the OriginMap. This allows 
ExternalASTMerger to complete them later if necessary.
As DeclContexts are imported from an origin that already has 
its own OriginMap, the origins are forwarded – but only for 
those DeclContexts that are actually used. This keeps the 
amount of stored data minimal.

The patch also applies several improvements from review:

- Thoroughly documents the interface to ExternalASTMerger;
- Adds optional logging to help track what's going on; and
- Cleans up a bunch of braces and dangling elses.

Differential Revision: https://reviews.llvm.org/D38208

llvm-svn: 314336

1 files changed, 309 insertions, 93 deletions

diff --git a/clang/lib/AST/ExternalASTMerger.cpp b/clang/lib/AST/ExternalASTMerger.cpp
index f69b5d62840..ec2e3034b15 100644
--- a/clang/lib/AST/ExternalASTMerger.cpp
+++ b/clang/lib/AST/ExternalASTMerger.cpp
@@ -14,6 +14,7 @@
 
 #include "clang/AST/ASTContext.h"
 #include "clang/AST/Decl.h"
+#include "clang/AST/DeclCXX.h"
 #include "clang/AST/DeclObjC.h"
 #include "clang/AST/ExternalASTMerger.h"
 
@@ -32,29 +33,18 @@ template <typename T> struct Source {
 
 typedef std::pair<Source<NamedDecl *>, ASTImporter *> Candidate;
 
-class LazyASTImporter : public ASTImporter {
-public:
-  LazyASTImporter(ASTContext &ToContext, FileManager &ToFileManager,
-                  ASTContext &FromContext, FileManager &FromFileManager)
-      : ASTImporter(ToContext, ToFileManager, FromContext, FromFileManager,
-                    /*MinimalImport=*/true) {}
-  Decl *Imported(Decl *From, Decl *To) override {
-    if (auto ToTag = dyn_cast<TagDecl>(To)) {
-      ToTag->setHasExternalLexicalStorage();
-      ToTag->setMustBuildLookupTable();
-    } else if (auto ToNamespace = dyn_cast<NamespaceDecl>(To)) {
-      ToNamespace->setHasExternalVisibleStorage();
-    } else if (auto ToContainer = dyn_cast<ObjCContainerDecl>(To)) {
-      ToContainer->setHasExternalLexicalStorage();
-      ToContainer->setMustBuildLookupTable();
-    }
-    return ASTImporter::Imported(From, To);
-  }
-};
+/// For the given DC, return the DC that is safe to perform lookups on.  This is
+/// the DC we actually want to work with most of the time.
+const DeclContext *CanonicalizeDC(const DeclContext *DC) {
+  if (isa<LinkageSpecDecl>(DC))
+    return DC->getRedeclContext();
+  return DC;
+}
 
 Source<const DeclContext *>
 LookupSameContext(Source<TranslationUnitDecl *> SourceTU, const DeclContext *DC,
                   ASTImporter &ReverseImporter) {
+  DC = CanonicalizeDC(DC);
   if (DC->isTranslationUnit()) {
     return SourceTU;
   }
@@ -64,101 +54,328 @@ LookupSameContext(Source<TranslationUnitDecl *> SourceTU, const DeclContext *DC,
     // If we couldn't find the parent DC in this TranslationUnit, give up.
     return nullptr;
   }
-  auto ND = cast<NamedDecl>(DC);
+  auto *ND = cast<NamedDecl>(DC);
   DeclarationName Name = ND->getDeclName();
   Source<DeclarationName> SourceName = ReverseImporter.Import(Name);
   DeclContext::lookup_result SearchResult =
       SourceParentDC.get()->lookup(SourceName.get());
   size_t SearchResultSize = SearchResult.size();
-  // Handle multiple candidates once we have a test for it.
-  // This may turn up when we import template specializations correctly.
-  assert(SearchResultSize < 2);
-  if (SearchResultSize == 0) {
-    // couldn't find the name, so we have to give up
+  if (SearchResultSize == 0 || SearchResultSize > 1) {
+    // There are two cases here.  First, we might not find the name.
+    // We might also find multiple copies, in which case we have no
+    // guarantee that the one we wanted is the one we pick.  (E.g.,
+    // if we have two specializations of the same template it is
+    // very hard to determine which is the one you want.)
+    //
+    // The Origins map fixes this problem by allowing the origin to be
+    // explicitly recorded, so we trigger that recording by returning
+    // nothing (rather than a possibly-inaccurate guess) here.
     return nullptr;
   } else {
     NamedDecl *SearchResultDecl = SearchResult[0];
-    return dyn_cast<DeclContext>(SearchResultDecl);
+    if (isa<DeclContext>(SearchResultDecl) &&
+        SearchResultDecl->getKind() == DC->getDeclKind())
+      return cast<DeclContext>(SearchResultDecl)->getPrimaryContext();
+    return nullptr; // This type of lookup is unsupported
   }
 }
 
-bool IsForwardDeclaration(Decl *D) {
-  if (auto TD = dyn_cast<TagDecl>(D)) {
-    return !TD->isThisDeclarationADefinition();
-  } else if (auto FD = dyn_cast<FunctionDecl>(D)) {
-    return !FD->isThisDeclarationADefinition();
-  } else if (auto OID = dyn_cast<ObjCInterfaceDecl>(D)) {
-     return OID->isThisDeclarationADefinition();
-  } else {
-    return false;
-  }
-}
+/// A custom implementation of ASTImporter, for ExternalASTMerger's purposes.
+///
+/// There are several modifications:
+///
+/// - It enables lazy lookup (via the HasExternalLexicalStorage flag and a few
+///   others), which instructs Clang to refer to ExternalASTMerger.  Also, it
+///   forces MinimalImport to true, which is necessary to make this work.
+/// - It maintains a reverse importer for use with names.  This allows lookup of
+///   arbitrary names in the source context.
+/// - It updates the ExternalASTMerger's origin map as needed whenever a
+///   it sees a DeclContext.
+class LazyASTImporter : public ASTImporter {
+private:
+  ExternalASTMerger &Parent;
+  ASTImporter Reverse;
+  const ExternalASTMerger::OriginMap &FromOrigins;
 
-template <typename CallbackType>
-void ForEachMatchingDC(
-    const DeclContext *DC,
-    llvm::ArrayRef<ExternalASTMerger::ImporterPair> Importers,
-    CallbackType Callback) {
-  for (const ExternalASTMerger::ImporterPair &IP : Importers) {
-    Source<TranslationUnitDecl *> SourceTU =
-        IP.Forward->getFromContext().getTranslationUnitDecl();
-    if (auto SourceDC = LookupSameContext(SourceTU, DC, *IP.Reverse))
-      Callback(IP, SourceDC);
+  llvm::raw_ostream &logs() { return Parent.logs(); }
+public:
+  LazyASTImporter(ExternalASTMerger &_Parent, ASTContext &ToContext,
+                  FileManager &ToFileManager, ASTContext &FromContext,
+                  FileManager &FromFileManager,
+                  const ExternalASTMerger::OriginMap &_FromOrigins)
+      : ASTImporter(ToContext, ToFileManager, FromContext, FromFileManager,
+                    /*MinimalImport=*/true),
+        Parent(_Parent), Reverse(FromContext, FromFileManager, ToContext,
+                                 ToFileManager, /*MinimalImport=*/true), FromOrigins(_FromOrigins) {}
+
+  /// Whenever a DeclContext is imported, ensure that ExternalASTSource's origin
+  /// map is kept up to date.  Also set the appropriate flags.
+  Decl *Imported(Decl *From, Decl *To) override {
+    if (auto *ToDC = dyn_cast<DeclContext>(To)) {
+      const bool LoggingEnabled = Parent.LoggingEnabled();
+      if (LoggingEnabled)
+        logs() << "(ExternalASTMerger*)" << (void*)&Parent
+               << " imported (DeclContext*)" << (void*)ToDC
+               << ", (ASTContext*)" << (void*)&getToContext()
+               << " from (DeclContext*)" << (void*)llvm::cast<DeclContext>(From)
+               << ", (ASTContext*)" << (void*)&getFromContext()
+               << "\n";
+      Source<DeclContext *> FromDC(
+          cast<DeclContext>(From)->getPrimaryContext());
+      if (FromOrigins.count(FromDC) &&
+          Parent.HasImporterForOrigin(*FromOrigins.at(FromDC).AST)) {
+        if (LoggingEnabled)
+          logs() << "(ExternalASTMerger*)" << (void*)&Parent
+                 << " forced origin (DeclContext*)"
+                 << (void*)FromOrigins.at(FromDC).DC
+                 << ", (ASTContext*)"
+                 << (void*)FromOrigins.at(FromDC).AST
+                 << "\n";
+        Parent.ForceRecordOrigin(ToDC, FromOrigins.at(FromDC));
+      } else {
+        if (LoggingEnabled)
+          logs() << "(ExternalASTMerger*)" << (void*)&Parent
+                 << " maybe recording origin (DeclContext*)" << (void*)FromDC
+                 << ", (ASTContext*)" << (void*)&getFromContext()
+                 << "\n";
+        Parent.MaybeRecordOrigin(ToDC, {FromDC, &getFromContext()});
+      }
+    }
+    if (auto *ToTag = dyn_cast<TagDecl>(To)) {
+      ToTag->setHasExternalLexicalStorage();
+      ToTag->setMustBuildLookupTable();
+      assert(Parent.CanComplete(ToTag));
+    } else if (auto *ToNamespace = dyn_cast<NamespaceDecl>(To)) {
+      ToNamespace->setHasExternalVisibleStorage();
+      assert(Parent.CanComplete(ToNamespace));
+    } else if (auto *ToContainer = dyn_cast<ObjCContainerDecl>(To)) {
+      ToContainer->setHasExternalLexicalStorage();
+      ToContainer->setMustBuildLookupTable();
+      assert(Parent.CanComplete(ToContainer));
+    }
+    return ASTImporter::Imported(From, To);
   }
-}
+  ASTImporter &GetReverse() { return Reverse; }
+};
 
 bool HasDeclOfSameType(llvm::ArrayRef<Candidate> Decls, const Candidate &C) {
+  if (isa<FunctionDecl>(C.first.get()))
+    return false;
   return llvm::any_of(Decls, [&](const Candidate &D) {
     return C.first.get()->getKind() == D.first.get()->getKind();
   });
 }
+
 } // end namespace
 
-ExternalASTMerger::ExternalASTMerger(const ImporterEndpoint &Target,
-                                     llvm::ArrayRef<ImporterEndpoint> Sources) {
-  for (const ImporterEndpoint &S : Sources) {
-    Importers.push_back(
-        {llvm::make_unique<LazyASTImporter>(Target.AST, Target.FM, S.AST, S.FM),
-         llvm::make_unique<ASTImporter>(S.AST, S.FM, Target.AST, Target.FM,
-                                        /*MinimalImport=*/true)});
-  }
+ASTImporter &ExternalASTMerger::ImporterForOrigin(ASTContext &OriginContext) {
+  for (const std::unique_ptr<ASTImporter> &I : Importers)
+    if (&I->getFromContext() == &OriginContext)
+      return *I;
+  llvm_unreachable("We should have an importer for this origin!");
 }
 
-bool ExternalASTMerger::FindExternalVisibleDeclsByName(const DeclContext *DC,
-                                                       DeclarationName Name) {
-  llvm::SmallVector<NamedDecl *, 1> Decls;
-  llvm::SmallVector<Candidate, 4> CompleteDecls;
-  llvm::SmallVector<Candidate, 4> ForwardDecls;
+namespace {
+LazyASTImporter &LazyImporterForOrigin(ExternalASTMerger &Merger,
+                                   ASTContext &OriginContext) {
+  return static_cast<LazyASTImporter &>(
+      Merger.ImporterForOrigin(OriginContext));
+}
+}
 
-  auto FilterFoundDecl = [&CompleteDecls, &ForwardDecls](const Candidate &C) {
-    if (IsForwardDeclaration(C.first.get())) {
-      if (!HasDeclOfSameType(ForwardDecls, C)) {
-        ForwardDecls.push_back(C);
+bool ExternalASTMerger::HasImporterForOrigin(ASTContext &OriginContext) {
+  for (const std::unique_ptr<ASTImporter> &I : Importers)
+    if (&I->getFromContext() == &OriginContext)
+      return true;
+  return false;
+}
+
+template <typename CallbackType>
+void ExternalASTMerger::ForEachMatchingDC(const DeclContext *DC,
+                                          CallbackType Callback) {
+  if (Origins.count(DC)) {
+    ExternalASTMerger::DCOrigin Origin = Origins[DC];
+    LazyASTImporter &Importer = LazyImporterForOrigin(*this, *Origin.AST);
+    Callback(Importer, Importer.GetReverse(), Origin.DC);
+  } else {
+    bool DidCallback = false;
+    for (const std::unique_ptr<ASTImporter> &Importer : Importers) {
+      Source<TranslationUnitDecl *> SourceTU =
+          Importer->getFromContext().getTranslationUnitDecl();
+      ASTImporter &Reverse =
+          static_cast<LazyASTImporter *>(Importer.get())->GetReverse();
+      if (auto SourceDC = LookupSameContext(SourceTU, DC, Reverse)) {
+        DidCallback = true;
+        if (Callback(*Importer, Reverse, SourceDC))
+          break;
       }
-    } else {
-      CompleteDecls.push_back(C);
     }
-  };
+    if (!DidCallback && LoggingEnabled())
+      logs() << "(ExternalASTMerger*)" << (void*)this
+             << " asserting for (DeclContext*)" << (void*)DC
+             << ", (ASTContext*)" << (void*)&Target.AST
+             << "\n";
+    assert(DidCallback && "Couldn't find a source context matching our DC");
+  }
+}
+
+void ExternalASTMerger::CompleteType(TagDecl *Tag) {
+  assert(Tag->hasExternalLexicalStorage());
+  ForEachMatchingDC(Tag, [&](ASTImporter &Forward, ASTImporter &Reverse,
+                             Source<const DeclContext *> SourceDC) -> bool {
+    auto *SourceTag = const_cast<TagDecl *>(cast<TagDecl>(SourceDC.get()));
+    if (SourceTag->hasExternalLexicalStorage())
+      SourceTag->getASTContext().getExternalSource()->CompleteType(SourceTag);
+    if (!SourceTag->getDefinition())
+      return false;
+    Forward.Imported(SourceTag, Tag);
+    Forward.ImportDefinition(SourceTag);
+    Tag->setCompleteDefinition(SourceTag->isCompleteDefinition());
+    return true;
+  });
+}
 
+void ExternalASTMerger::CompleteType(ObjCInterfaceDecl *Interface) {
+  assert(Interface->hasExternalLexicalStorage());
   ForEachMatchingDC(
-      DC, Importers,
-      [&](const ImporterPair &IP, Source<const DeclContext *> SourceDC) {
-        DeclarationName FromName = IP.Reverse->Import(Name);
-        DeclContextLookupResult Result = SourceDC.get()->lookup(FromName);
-        for (NamedDecl *FromD : Result) {
-          FilterFoundDecl(std::make_pair(FromD, IP.Forward.get()));
-        }
+      Interface, [&](ASTImporter &Forward, ASTImporter &Reverse,
+                     Source<const DeclContext *> SourceDC) -> bool {
+        auto *SourceInterface = const_cast<ObjCInterfaceDecl *>(
+            cast<ObjCInterfaceDecl>(SourceDC.get()));
+        if (SourceInterface->hasExternalLexicalStorage())
+          SourceInterface->getASTContext().getExternalSource()->CompleteType(
+              SourceInterface);
+        if (!SourceInterface->getDefinition())
+          return false;
+        Forward.Imported(SourceInterface, Interface);
+        Forward.ImportDefinition(SourceInterface);
+        return true;
       });
+}
 
-  llvm::ArrayRef<Candidate> DeclsToReport =
-      CompleteDecls.empty() ? ForwardDecls : CompleteDecls;
+bool ExternalASTMerger::CanComplete(DeclContext *Interface) {
+  assert(Interface->hasExternalLexicalStorage() ||
+         Interface->hasExternalVisibleStorage());
+  bool FoundMatchingDC = false;
+  ForEachMatchingDC(Interface,
+                    [&](ASTImporter &Forward, ASTImporter &Reverse,
+                        Source<const DeclContext *> SourceDC) -> bool {
+                      FoundMatchingDC = true;
+                      return true;
+                    });
+  return FoundMatchingDC;
+}
 
-  if (DeclsToReport.empty()) {
-    return false;
+namespace {
+bool IsSameDC(const DeclContext *D1, const DeclContext *D2) {
+  if (isa<ObjCContainerDecl>(D1) && isa<ObjCContainerDecl>(D2))
+    return true; // There are many cases where Objective-C is ambiguous.
+  if (auto *T1 = dyn_cast<TagDecl>(D1))
+    if (auto *T2 = dyn_cast<TagDecl>(D2))
+      if (T1->getFirstDecl() == T2->getFirstDecl())
+        return true;
+  return D1 == D2 || D1 == CanonicalizeDC(D2);
+}
+}
+
+void ExternalASTMerger::MaybeRecordOrigin(const DeclContext *ToDC,
+                                          DCOrigin Origin) {
+  LazyASTImporter &Importer = LazyImporterForOrigin(*this, *Origin.AST);
+  ASTImporter &Reverse = Importer.GetReverse();
+  Source<const DeclContext *> FoundFromDC =
+      LookupSameContext(Origin.AST->getTranslationUnitDecl(), ToDC, Reverse);
+  const bool DoRecord = !FoundFromDC || !IsSameDC(FoundFromDC.get(), Origin.DC);
+  if (DoRecord)
+    RecordOriginImpl(ToDC, Origin, Importer);
+  if (LoggingEnabled())
+    logs() << "(ExternalASTMerger*)" << (void*)this
+             << (DoRecord ? " decided " : " decided NOT")
+             << " to record origin (DeclContext*)" << (void*)Origin.DC
+             << ", (ASTContext*)" << (void*)&Origin.AST
+             << "\n";
+}
+
+void ExternalASTMerger::ForceRecordOrigin(const DeclContext *ToDC,
+                                          DCOrigin Origin) {
+  RecordOriginImpl(ToDC, Origin, ImporterForOrigin(*Origin.AST));
+}
+
+void ExternalASTMerger::RecordOriginImpl(const DeclContext *ToDC, DCOrigin Origin,
+                                         ASTImporter &Importer) {
+  Origins[ToDC] = Origin;
+  Importer.ASTImporter::Imported(cast<Decl>(Origin.DC), const_cast<Decl*>(cast<Decl>(ToDC)));
+}
+
+ExternalASTMerger::ExternalASTMerger(const ImporterTarget &Target,
+                                     llvm::ArrayRef<ImporterSource> Sources) : LogStream(&llvm::nulls()), Target(Target) {
+  AddSources(Sources);
+}
+
+void ExternalASTMerger::AddSources(llvm::ArrayRef<ImporterSource> Sources) {
+  for (const ImporterSource &S : Sources) {
+    assert(&S.AST != &Target.AST);
+    Importers.push_back(llvm::make_unique<LazyASTImporter>(
+        *this, Target.AST, Target.FM, S.AST, S.FM, S.OM));
+  }
+}
+
+void ExternalASTMerger::RemoveSources(llvm::ArrayRef<ImporterSource> Sources) {
+  if (LoggingEnabled())
+    for (const ImporterSource &S : Sources)
+      logs() << "(ExternalASTMerger*)" << (void*)this
+             << " removing source (ASTContext*)" << (void*)&S.AST
+             << "\n";
+  Importers.erase(
+      std::remove_if(Importers.begin(), Importers.end(),
+                     [&Sources](std::unique_ptr<ASTImporter> &Importer) -> bool {
+                       for (const ImporterSource &S : Sources) {
+                         if (&Importer->getFromContext() == &S.AST)
+                           return true;
+                       }
+                       return false;
+                     }),
+      Importers.end());
+  for (OriginMap::iterator OI = Origins.begin(), OE = Origins.end(); OI != OE; ) {
+    std::pair<const DeclContext *, DCOrigin> Origin = *OI;
+    bool Erase = false;
+    for (const ImporterSource &S : Sources) {
+      if (&S.AST == Origin.second.AST) {
+        Erase = true;
+        break;
+      }
+    }
+    if (Erase)
+      OI = Origins.erase(OI);
+    else
+      ++OI;
   }
+}
+
+bool ExternalASTMerger::FindExternalVisibleDeclsByName(const DeclContext *DC,
+                                                       DeclarationName Name) {
+  llvm::SmallVector<NamedDecl *, 1> Decls;
+  llvm::SmallVector<Candidate, 4> Candidates;
 
-  Decls.reserve(DeclsToReport.size());
-  for (const Candidate &C : DeclsToReport) {
+  auto FilterFoundDecl = [&Candidates](const Candidate &C) {
+   if (!HasDeclOfSameType(Candidates, C))
+     Candidates.push_back(C);
+  };
+
+  ForEachMatchingDC(DC, [&](ASTImporter &Forward, ASTImporter &Reverse,
+                            Source<const DeclContext *> SourceDC) -> bool {
+    DeclarationName FromName = Reverse.Import(Name);
+    DeclContextLookupResult Result = SourceDC.get()->lookup(FromName);
+    for (NamedDecl *FromD : Result) {
+      FilterFoundDecl(std::make_pair(FromD, &Forward));
+    }
+    return false;
+  });
+
+  if (Candidates.empty())
+    return false;
+
+  Decls.reserve(Candidates.size());
+  for (const Candidate &C : Candidates) {
     NamedDecl *d = cast<NamedDecl>(C.second->Import(C.first.get()));
     assert(d);
     Decls.push_back(d);
@@ -170,17 +387,16 @@ bool ExternalASTMerger::FindExternalVisibleDeclsByName(const DeclContext *DC,
 void ExternalASTMerger::FindExternalLexicalDecls(
     const DeclContext *DC, llvm::function_ref<bool(Decl::Kind)> IsKindWeWant,
     SmallVectorImpl<Decl *> &Result) {
-  ForEachMatchingDC(
-      DC, Importers,
-      [&](const ImporterPair &IP, Source<const DeclContext *> SourceDC) {
-        for (const Decl *SourceDecl : SourceDC.get()->decls()) {
-          if (IsKindWeWant(SourceDecl->getKind())) {
-            Decl *ImportedDecl =
-                IP.Forward->Import(const_cast<Decl *>(SourceDecl));
-            assert(ImportedDecl->getDeclContext() == DC);
-            (void)ImportedDecl;
-          }
-        }
-      });
+  ForEachMatchingDC(DC, [&](ASTImporter &Forward, ASTImporter &Reverse,
+                            Source<const DeclContext *> SourceDC) -> bool {
+    for (const Decl *SourceDecl : SourceDC.get()->decls()) {
+      if (IsKindWeWant(SourceDecl->getKind())) {
+        Decl *ImportedDecl = Forward.Import(const_cast<Decl *>(SourceDecl));
+        assert(!ImportedDecl || IsSameDC(ImportedDecl->getDeclContext(), DC));
+        (void)ImportedDecl;
+      }
+    }
+    return false;
+  });
 }