[clangd] Define out-of-line qualify return value

Summary:
Return type might need qualification if insertion context doesn't have
the same decls visible as the source context.

This patch adds qualification for return value to cover such cases.

Reviewers: hokein

Subscribers: ilya-biryukov, MaskRay, jkorous, arphaman, usaxena95, cfe-commits

Tags: #clang

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

2 files changed, 189 insertions, 24 deletions

diff --git a/clang-tools-extra/clangd/refactor/tweaks/DefineOutline.cpp b/clang-tools-extra/clangd/refactor/tweaks/DefineOutline.cpp
index 102c650eec1..1bd70b72aa1 100644
--- a/clang-tools-extra/clangd/refactor/tweaks/DefineOutline.cpp
+++ b/clang-tools-extra/clangd/refactor/tweaks/DefineOutline.cpp
@@ -6,13 +6,17 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "AST.h"
+#include "FindTarget.h"
 #include "HeaderSourceSwitch.h"
+#include "Logger.h"
 #include "Path.h"
 #include "Selection.h"
 #include "SourceCode.h"
 #include "refactor/Tweak.h"
 #include "clang/AST/ASTTypeTraits.h"
 #include "clang/AST/Decl.h"
+#include "clang/AST/DeclBase.h"
 #include "clang/AST/DeclTemplate.h"
 #include "clang/AST/Stmt.h"
 #include "clang/Basic/SourceLocation.h"
@@ -20,10 +24,13 @@
 #include "clang/Driver/Types.h"
 #include "clang/Format/Format.h"
 #include "clang/Tooling/Core/Replacement.h"
+#include "llvm/ADT/None.h"
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/Support/Casting.h"
 #include "llvm/Support/Error.h"
 #include <cstddef>
+#include <string>
 
 namespace clang {
 namespace clangd {
@@ -57,31 +64,136 @@ llvm::Optional<Path> getSourceFile(llvm::StringRef FileName,
   return getCorrespondingHeaderOrSource(FileName, Sel.AST, Sel.Index);
 }
 
-// Creates a modified version of function definition that can be inserted at a
-// different location. Contains both function signature and body.
-llvm::Optional<llvm::StringRef> getFunctionSourceCode(const FunctionDecl *FD) {
-  auto &SM = FD->getASTContext().getSourceManager();
-  auto CharRange = toHalfOpenFileRange(SM, FD->getASTContext().getLangOpts(),
-                                       FD->getSourceRange());
-  if (!CharRange)
+// Synthesize a DeclContext for TargetNS from CurContext. TargetNS must be empty
+// for global namespace, and endwith "::" otherwise.
+// Returns None if TargetNS is not a prefix of CurContext.
+llvm::Optional<const DeclContext *>
+findContextForNS(llvm::StringRef TargetNS, const DeclContext *CurContext) {
+  assert(TargetNS.empty() || TargetNS.endswith("::"));
+  // Skip any non-namespace contexts, e.g. TagDecls, functions/methods.
+  CurContext = CurContext->getEnclosingNamespaceContext();
+  // If TargetNS is empty, it means global ns, which is translation unit.
+  if (TargetNS.empty()) {
+    while (!CurContext->isTranslationUnit())
+      CurContext = CurContext->getParent();
+    return CurContext;
+  }
+  // Otherwise we need to drop any trailing namespaces from CurContext until
+  // we reach TargetNS.
+  std::string TargetContextNS =
+      CurContext->isNamespace()
+          ? llvm::cast<NamespaceDecl>(CurContext)->getQualifiedNameAsString()
+          : "";
+  TargetContextNS.append("::");
+
+  llvm::StringRef CurrentContextNS(TargetContextNS);
+  // If TargetNS is not a prefix of CurrentContext, there's no way to reach
+  // it.
+  if (!CurrentContextNS.startswith(TargetNS))
     return llvm::None;
+
+  while (CurrentContextNS != TargetNS) {
+    CurContext = CurContext->getParent();
+    // These colons always exists since TargetNS is a prefix of
+    // CurrentContextNS, it ends with "::" and they are not equal.
+    CurrentContextNS = CurrentContextNS.take_front(
+        CurrentContextNS.drop_back(2).rfind("::") + 2);
+  }
+  return CurContext;
+}
+
+// Returns source code for FD after applying Replacements.
+// FIXME: Make the function take a parameter to return only the function body,
+// afterwards it can be shared with define-inline code action.
+llvm::Expected<std::string>
+getFunctionSourceAfterReplacements(const FunctionDecl *FD,
+                                   const tooling::Replacements &Replacements) {
+  const auto &SM = FD->getASTContext().getSourceManager();
+  auto OrigFuncRange = toHalfOpenFileRange(
+      SM, FD->getASTContext().getLangOpts(), FD->getSourceRange());
+  if (!OrigFuncRange)
+    return llvm::createStringError(llvm::inconvertibleErrorCode(),
+                                   "Couldn't get range for function.");
   // Include template parameter list.
   if (auto *FTD = FD->getDescribedFunctionTemplate())
-    CharRange->setBegin(FTD->getBeginLoc());
+    OrigFuncRange->setBegin(FTD->getBeginLoc());
 
-  // FIXME: Qualify return type.
-  // FIXME: Qualify function name depending on the target context.
-  return toSourceCode(SM, *CharRange);
+  // Get new begin and end positions for the qualified function definition.
+  unsigned FuncBegin = SM.getFileOffset(OrigFuncRange->getBegin());
+  unsigned FuncEnd = Replacements.getShiftedCodePosition(
+      SM.getFileOffset(OrigFuncRange->getEnd()));
+
+  // Trim the result to function definition.
+  auto QualifiedFunc = tooling::applyAllReplacements(
+      SM.getBufferData(SM.getMainFileID()), Replacements);
+  if (!QualifiedFunc)
+    return QualifiedFunc.takeError();
+  return QualifiedFunc->substr(FuncBegin, FuncEnd - FuncBegin + 1);
 }
 
+// Creates a modified version of function definition that can be inserted at a
+// different location, qualifies return value and function name to achieve that.
+// Contains function signature, body and template parameters if applicable.
+// No need to qualify parameters, as they are looked up in the context
+// containing the function/method.
+// FIXME: Qualify function name depending on the target context.
+llvm::Expected<std::string>
+getFunctionSourceCode(const FunctionDecl *FD, llvm::StringRef TargetNamespace) {
+  auto &SM = FD->getASTContext().getSourceManager();
+  auto TargetContext = findContextForNS(TargetNamespace, FD->getDeclContext());
+  if (!TargetContext)
+    return llvm::createStringError(
+        llvm::inconvertibleErrorCode(),
+        "define outline: couldn't find a context for target");
+
+  llvm::Error Errors = llvm::Error::success();
+  tooling::Replacements QualifierInsertions;
+
+  // Finds the first unqualified name in function return type and qualifies it
+  // to be valid in TargetContext.
+  findExplicitReferences(FD, [&](ReferenceLoc Ref) {
+    // It is enough to qualify the first qualifier, so skip references with a
+    // qualifier. Also we can't do much if there are no targets or name is
+    // inside a macro body.
+    if (Ref.Qualifier || Ref.Targets.empty() || Ref.NameLoc.isMacroID())
+      return;
+    // Qualify return type
+    if (Ref.NameLoc != FD->getReturnTypeSourceRange().getBegin())
+      return;
+
+    for (const NamedDecl *ND : Ref.Targets) {
+      if (ND->getDeclContext() != Ref.Targets.front()->getDeclContext()) {
+        elog("Targets from multiple contexts: {0}, {1}",
+             printQualifiedName(*Ref.Targets.front()), printQualifiedName(*ND));
+        return;
+      }
+    }
+    const NamedDecl *ND = Ref.Targets.front();
+    const std::string Qualifier =
+        getQualification(FD->getASTContext(), *TargetContext,
+                         SM.getLocForStartOfFile(SM.getMainFileID()), ND);
+    if (auto Err = QualifierInsertions.add(
+            tooling::Replacement(SM, Ref.NameLoc, 0, Qualifier)))
+      Errors = llvm::joinErrors(std::move(Errors), std::move(Err));
+  });
+
+  if (Errors)
+    return std::move(Errors);
+  return getFunctionSourceAfterReplacements(FD, QualifierInsertions);
+}
+
+struct InsertionPoint {
+  std::string EnclosingNamespace;
+  size_t Offset;
+};
 // Returns the most natural insertion point for \p QualifiedName in \p Contents.
 // This currently cares about only the namespace proximity, but in feature it
 // should also try to follow ordering of declarations. For example, if decls
 // come in order `foo, bar, baz` then this function should return some point
 // between foo and baz for inserting bar.
-llvm::Expected<size_t> getInsertionOffset(llvm::StringRef Contents,
-                                          llvm::StringRef QualifiedName,
-                                          const format::FormatStyle &Style) {
+llvm::Expected<InsertionPoint>
+getInsertionPoint(llvm::StringRef Contents, llvm::StringRef QualifiedName,
+                  const format::FormatStyle &Style) {
   auto Region = getEligiblePoints(Contents, QualifiedName, Style);
 
   assert(!Region.EligiblePoints.empty());
@@ -89,8 +201,10 @@ llvm::Expected<size_t> getInsertionOffset(llvm::StringRef Contents,
   // locations for adjacent decls to Source. Unfortunately psudeo parsing in
   // getEligibleRegions only knows about namespace begin/end events so we
   // can't match function start/end positions yet.
-  auto InsertionPoint = Region.EligiblePoints.back();
-  return positionToOffset(Contents, InsertionPoint);
+  auto Offset = positionToOffset(Contents, Region.EligiblePoints.back());
+  if (!Offset)
+    return Offset.takeError();
+  return InsertionPoint{Region.EnclosingNamespace, *Offset};
 }
 
 /// Moves definition of a function/method to an appropriate implementation file.
@@ -170,21 +284,22 @@ public:
       return llvm::createStringError(Buffer.getError(),
                                      Buffer.getError().message());
     auto Contents = Buffer->get()->getBuffer();
-    auto InsertionOffset =
-        getInsertionOffset(Contents, Source->getQualifiedNameAsString(),
-                           getFormatStyleForFile(*CCFile, Contents, &FS));
-    if (!InsertionOffset)
-      return InsertionOffset.takeError();
+    auto InsertionPoint =
+        getInsertionPoint(Contents, Source->getQualifiedNameAsString(),
+                          getFormatStyleForFile(*CCFile, Contents, &FS));
+    if (!InsertionPoint)
+      return InsertionPoint.takeError();
 
-    auto FuncDef = getFunctionSourceCode(Source);
+    auto FuncDef =
+        getFunctionSourceCode(Source, InsertionPoint->EnclosingNamespace);
     if (!FuncDef)
       return llvm::createStringError(
           llvm::inconvertibleErrorCode(),
           "Couldn't get full source for function definition.");
 
     SourceManagerForFile SMFF(*CCFile, Contents);
-    const tooling::Replacement InsertFunctionDef(*CCFile, *InsertionOffset, 0,
-                                                 *FuncDef);
+    const tooling::Replacement InsertFunctionDef(
+        *CCFile, InsertionPoint->Offset, 0, *FuncDef);
     auto Effect = Effect::mainFileEdit(
         SMFF.get(), tooling::Replacements(InsertFunctionDef));
     if (!Effect)
diff --git a/clang-tools-extra/clangd/unittests/TweakTests.cpp b/clang-tools-extra/clangd/unittests/TweakTests.cpp
index 33665122ad6..9710465335d 100644
--- a/clang-tools-extra/clangd/unittests/TweakTests.cpp
+++ b/clang-tools-extra/clangd/unittests/TweakTests.cpp
@@ -1972,6 +1972,56 @@ TEST_F(DefineOutlineTest, HandleMacros) {
   }
 }
 
+TEST_F(DefineOutlineTest, QualifyReturnValue) {
+  FileName = "Test.hpp";
+  ExtraFiles["Test.cpp"] = "";
+
+  struct {
+    llvm::StringRef Test;
+    llvm::StringRef ExpectedHeader;
+    llvm::StringRef ExpectedSource;
+  } Cases[] = {
+      {R"cpp(
+        namespace a { class Foo; }
+        using namespace a;
+        Foo fo^o() { return; })cpp",
+       R"cpp(
+        namespace a { class Foo; }
+        using namespace a;
+        Foo foo() ;)cpp",
+       "a::Foo foo() { return; }"},
+      {R"cpp(
+        namespace a {
+          class Foo {
+            class Bar {};
+            Bar fo^o() { return {}; }
+          };
+        })cpp",
+       R"cpp(
+        namespace a {
+          class Foo {
+            class Bar {};
+            Bar foo() ;
+          };
+        })cpp",
+       "a::Foo::Bar foo() { return {}; }\n"},
+      {R"cpp(
+        class Foo;
+        Foo fo^o() { return; })cpp",
+       R"cpp(
+        class Foo;
+        Foo foo() ;)cpp",
+       "Foo foo() { return; }"},
+  };
+  llvm::StringMap<std::string> EditedFiles;
+  for (auto &Case : Cases) {
+    apply(Case.Test, &EditedFiles);
+    EXPECT_EQ(apply(Case.Test, &EditedFiles), Case.ExpectedHeader);
+    EXPECT_THAT(EditedFiles, testing::ElementsAre(FileWithContents(
+                                 testPath("Test.cpp"), Case.ExpectedSource)));
+  }
+}
+
 } // namespace
 } // namespace clangd
 } // namespace clang