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

[bkramer]

Summary

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

[klimek]

LG

include/clang/Frontend/CompilerInstance.h

class CompilerInstance : public ModuleLoader {
   IntrusiveRefCntPtr<SourceManager> SourceMgr;
 
   /// The preprocessor.
   IntrusiveRefCntPtr<Preprocessor> PP;
 
   /// 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;
 
   /// The code completion consumer.
   std::unique_ptr<CodeCompleteConsumer> CompletionConsumer;
 
   /// \brief The semantic analysis object.
   std::unique_ptr<Sema> TheSema;
public:
 
   GlobalModuleIndex *loadGlobalModuleIndex(SourceLocation TriggerLoc) override;
 
   bool lookupMissingImports(StringRef Name, SourceLocation TriggerLoc) override;
 
   void addDependencyCollector(std::shared_ptr<DependencyCollector> Listener) {
     DependencyCollectors.push_back(std::move(Listener));
   }
+
+  void setExternalSemaSource(IntrusiveRefCntPtr<ExternalSemaSource> ESS);
 };
 
 } // end namespace clang
 
 #endif

include/clang/Sema/TypoCorrection.h

public:
   }
   const_decl_iterator end() const { return CorrectionDecls.end(); }
 
   /// \brief Returns whether this typo correction is correcting to a
   /// declaration that was declared in a module that has not been imported.
   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());
   }
 
   // Results.
   DeclarationName CorrectionName;
   NestedNameSpecifier *CorrectionNameSpec;
   SmallVector<NamedDecl *, 1> CorrectionDecls;
   unsigned CharDistance;
   unsigned QualifierDistance;
   unsigned CallbackDistance;
   SourceRange CorrectionRange;
   bool ForceSpecifierReplacement;
   bool RequiresImport;
+
+  std::vector<PartialDiagnostic> ExtraDiagnostics;
 };
 
 /// @brief Base class for callback objects used by Sema::CorrectTypo to check
 /// the validity of a potential typo correction.
 class CorrectionCandidateCallback {
 public:
   static const unsigned InvalidDistance = TypoCorrection::InvalidDistance;
 

lib/Sema/SemaLookup.cpp

void Sema::diagnoseTypo(const TypoCorrection &Correction,
   Diag(Correction.getCorrectionRange().getBegin(), TypoDiag)
     << CorrectedQuotedStr << (ErrorRecovery ? FixTypo : FixItHint());
 
   NamedDecl *ChosenDecl =
       Correction.isKeyword() ? nullptr : Correction.getFoundDecl();
   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,
                                   TypoDiagnosticGenerator TDG,
                                   TypoRecoveryCallback TRC) {
   assert(TCC && "createDelayedTypo requires a valid TypoCorrectionConsumer");
   auto TE = new (Context) TypoExpr(Context.DependentTy);
   auto &State = DelayedTypos[TE];

unittests/Frontend/FrontendActionTest.cpp

 //===- unittests/Frontend/FrontendActionTest.cpp - FrontendAction tests ---===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 
 #include "clang/AST/ASTConsumer.h"
 #include "clang/AST/ASTContext.h"
 #include "clang/AST/RecursiveASTVisitor.h"
 #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"
 #include "llvm/ADT/Triple.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "gtest/gtest.h"
 
 using namespace llvm;
TEST(PreprocessorFrontendAction, EndSourceFile) {
   TestPPCallbacksFrontendAction TestAction(Callbacks);
   ASSERT_FALSE(Callbacks->SeenEnd);
   ASSERT_FALSE(TestAction.SeenEnd);
   ASSERT_TRUE(Compiler.ExecuteAction(TestAction));
   // Check that EndOfMainFile was called before EndSourceFileAction.
   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