[Frontend] Allow attaching an external sema source to compiler instance and extra diags to TypoCorrections

This can be used to append alternative typo corrections to an existing diag.
include-fixer can use it to suggest includes to be added.

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

git-svn-id: https://llvm.org/svn/llvm-project/cfe/trunk@287128 91177308-0d34-0410-b5e6-96231b3b80d8

5 files changed, 93 insertions, 0 deletions

diff --git a/include/clang/Frontend/CompilerInstance.h b/include/clang/Frontend/CompilerInstance.h
index c0960bdfb3..3f754d9998 100644
--- a/include/clang/Frontend/CompilerInstance.h
+++ b/include/clang/Frontend/CompilerInstance.h
@@ -96,6 +96,9 @@ class CompilerInstance : public ModuleLoader {
   /// The AST context.
   IntrusiveRefCntPtr<ASTContext> Context;
 
+  /// An optional sema source that will be attached to sema.
+  IntrusiveRefCntPtr<ExternalSemaSource> ExternalSemaSrc;
+
   /// The AST consumer.
   std::unique_ptr<ASTConsumer> Consumer;
 
@@ -774,6 +777,8 @@ public:
   void addDependencyCollector(std::shared_ptr<DependencyCollector> Listener) {
     DependencyCollectors.push_back(std::move(Listener));
   }
+
+  void setExternalSemaSource(IntrusiveRefCntPtr<ExternalSemaSource> ESS);
 };
 
 } // end namespace clang
diff --git a/include/clang/Sema/TypoCorrection.h b/include/clang/Sema/TypoCorrection.h
index 3b0385ef2a..5bebe52325 100644
--- a/include/clang/Sema/TypoCorrection.h
+++ b/include/clang/Sema/TypoCorrection.h
@@ -230,6 +230,15 @@ public:
   bool requiresImport() const { return RequiresImport; }
   void setRequiresImport(bool Req) { RequiresImport = Req; }
 
+  /// Extra diagnostics are printed after the first diagnostic for the typo.
+  /// This can be used to attach external notes to the diag.
+  void addExtraDiagnostic(PartialDiagnostic PD) {
+    ExtraDiagnostics.push_back(std::move(PD));
+  }
+  ArrayRef<PartialDiagnostic> getExtraDiagnostics() const {
+    return ExtraDiagnostics;
+  }
+
 private:
   bool hasCorrectionDecl() const {
     return (!isKeyword() && !CorrectionDecls.empty());
@@ -245,6 +254,8 @@ private:
   SourceRange CorrectionRange;
   bool ForceSpecifierReplacement;
   bool RequiresImport;
+
+  std::vector<PartialDiagnostic> ExtraDiagnostics;
 };
 
 /// @brief Base class for callback objects used by Sema::CorrectTypo to check
diff --git a/lib/Frontend/CompilerInstance.cpp b/lib/Frontend/CompilerInstance.cpp
index b2b7b7c67b..b96dc70cf5 100644
--- a/lib/Frontend/CompilerInstance.cpp
+++ b/lib/Frontend/CompilerInstance.cpp
@@ -538,6 +538,11 @@ void CompilerInstance::createSema(TranslationUnitKind TUKind,
                                   CodeCompleteConsumer *CompletionConsumer) {
   TheSema.reset(new Sema(getPreprocessor(), getASTContext(), getASTConsumer(),
                          TUKind, CompletionConsumer));
+  // Attach the external sema source if there is any.
+  if (ExternalSemaSrc) {
+    TheSema->addExternalSource(ExternalSemaSrc.get());
+    ExternalSemaSrc->InitializeSema(*TheSema);
+  }
 }
 
 // Output Files
@@ -1820,3 +1825,8 @@ CompilerInstance::lookupMissingImports(StringRef Name,
   return false;
 }
 void CompilerInstance::resetAndLeakSema() { BuryPointer(takeSema()); }
+
+void CompilerInstance::setExternalSemaSource(
+    IntrusiveRefCntPtr<ExternalSemaSource> ESS) {
+  ExternalSemaSrc = std::move(ESS);
+}
diff --git a/lib/Sema/SemaLookup.cpp b/lib/Sema/SemaLookup.cpp
index e5b7af8617..1fc6d4fd51 100644
--- a/lib/Sema/SemaLookup.cpp
+++ b/lib/Sema/SemaLookup.cpp
@@ -5082,6 +5082,10 @@ void Sema::diagnoseTypo(const TypoCorrection &Correction,
   if (PrevNote.getDiagID() && ChosenDecl)
     Diag(ChosenDecl->getLocation(), PrevNote)
       << CorrectedQuotedStr << (ErrorRecovery ? FixItHint() : FixTypo);
+
+  // Add any extra diagnostics.
+  for (const PartialDiagnostic &PD : Correction.getExtraDiagnostics())
+    Diag(Correction.getCorrectionRange().getBegin(), PD);
 }
 
 TypoExpr *Sema::createDelayedTypo(std::unique_ptr<TypoCorrectionConsumer> TCC,
diff --git a/unittests/Frontend/FrontendActionTest.cpp b/unittests/Frontend/FrontendActionTest.cpp
index 39a131f4a6..c3e6adb632 100644
--- a/unittests/Frontend/FrontendActionTest.cpp
+++ b/unittests/Frontend/FrontendActionTest.cpp
@@ -13,6 +13,7 @@
 #include "clang/Frontend/CompilerInstance.h"
 #include "clang/Frontend/CompilerInvocation.h"
 #include "clang/Frontend/FrontendAction.h"
+#include "clang/Frontend/FrontendActions.h"
 #include "clang/Lex/Preprocessor.h"
 #include "clang/Lex/PreprocessorOptions.h"
 #include "clang/Sema/Sema.h"
@@ -192,4 +193,66 @@ TEST(PreprocessorFrontendAction, EndSourceFile) {
   ASSERT_TRUE(TestAction.SeenEnd);
 }
 
+class TypoExternalSemaSource : public ExternalSemaSource {
+  CompilerInstance &CI;
+
+public:
+  TypoExternalSemaSource(CompilerInstance &CI) : CI(CI) {}
+
+  TypoCorrection CorrectTypo(const DeclarationNameInfo &Typo, int LookupKind,
+                             Scope *S, CXXScopeSpec *SS,
+                             CorrectionCandidateCallback &CCC,
+                             DeclContext *MemberContext, bool EnteringContext,
+                             const ObjCObjectPointerType *OPT) override {
+    // Generate a fake typo correction with one attached note.
+    ASTContext &Ctx = CI.getASTContext();
+    TypoCorrection TC(DeclarationName(&Ctx.Idents.get("moo")));
+    unsigned DiagID = Ctx.getDiagnostics().getCustomDiagID(
+        DiagnosticsEngine::Note, "This is a note");
+    TC.addExtraDiagnostic(PartialDiagnostic(DiagID, Ctx.getDiagAllocator()));
+    return TC;
+  }
+};
+
+struct TypoDiagnosticConsumer : public DiagnosticConsumer {
+  void HandleDiagnostic(DiagnosticsEngine::Level DiagLevel,
+                        const Diagnostic &Info) override {
+    // Capture errors and notes. There should be one of each.
+    if (DiagLevel == DiagnosticsEngine::Error) {
+      assert(Error.empty());
+      Info.FormatDiagnostic(Error);
+    } else {
+      assert(Note.empty());
+      Info.FormatDiagnostic(Note);
+    }
+  }
+  SmallString<32> Error;
+  SmallString<32> Note;
+};
+
+TEST(ASTFrontendAction, ExternalSemaSource) {
+  auto *Invocation = new CompilerInvocation;
+  Invocation->getLangOpts()->CPlusPlus = true;
+  Invocation->getPreprocessorOpts().addRemappedFile(
+      "test.cc", MemoryBuffer::getMemBuffer("void fooo();\n"
+                                            "int main() { foo(); }")
+                     .release());
+  Invocation->getFrontendOpts().Inputs.push_back(
+      FrontendInputFile("test.cc", IK_CXX));
+  Invocation->getFrontendOpts().ProgramAction = frontend::ParseSyntaxOnly;
+  Invocation->getTargetOpts().Triple = "i386-unknown-linux-gnu";
+  CompilerInstance Compiler;
+  Compiler.setInvocation(Invocation);
+  auto *TDC = new TypoDiagnosticConsumer;
+  Compiler.createDiagnostics(TDC, /*ShouldOwnClient=*/true);
+  Compiler.setExternalSemaSource(new TypoExternalSemaSource(Compiler));
+
+  SyntaxOnlyAction TestAction;
+  ASSERT_TRUE(Compiler.ExecuteAction(TestAction));
+  // There should be one error correcting to 'moo' and a note attached to it.
+  EXPECT_EQ("use of undeclared identifier 'foo'; did you mean 'moo'?",
+            TDC->Error.str().str());
+  EXPECT_EQ("This is a note", TDC->Note.str().str());
+}
+
 } // anonymous namespace