[clangd] Find better insertion locations in DefineOutline tweak

If possible, put the definition next to the definition of an adjacent
declaration. For example:

struct S {
  void f^oo1() {}
  void foo2();
  void f^oo3() {}
};

// S::foo1() goes here
void S::foo2() {}
// S::foo3() goes here

Author: ckandeler

clang-tools-extra/clangd/SourceCode.cpp

@@ -1217,6 +1217,26 @@ EligibleRegion getEligiblePoints(llvm::StringRef Code,
   return ER;
 }
 
+std::string getNamespaceAtPosition(StringRef Code, const Position &Pos,
+                                   const LangOptions &LangOpts) {
+  std::vector<std::string> Enclosing = {""};
+  parseNamespaceEvents(Code, LangOpts, [&](NamespaceEvent Event) {
+    if (Pos < Event.Pos)
+      return;
+    if (Event.Trigger == NamespaceEvent::UsingDirective)
+      return;
+    if (!Event.Payload.empty())
+      Event.Payload.append("::");
+    if (Event.Trigger == NamespaceEvent::BeginNamespace) {
+      Enclosing.emplace_back(std::move(Event.Payload));
+    } else {
+      Enclosing.pop_back();
+      assert(Enclosing.back() == Event.Payload);
+    }
+  });
+  return Enclosing.back();
+}
+
 bool isHeaderFile(llvm::StringRef FileName,
                   std::optional<LangOptions> LangOpts) {
   // Respect the langOpts, for non-file-extension cases, e.g. standard library

clang-tools-extra/clangd/SourceCode.h

@@ -309,6 +309,9 @@ EligibleRegion getEligiblePoints(llvm::StringRef Code,
                                  llvm::StringRef FullyQualifiedName,
                                  const LangOptions &LangOpts);
 
+std::string getNamespaceAtPosition(llvm::StringRef Code, const Position &Pos,
+                                   const LangOptions &LangOpts);
+
 struct DefinedMacro {
   llvm::StringRef Name;
   const MacroInfo *Info;

clang-tools-extra/clangd/refactor/tweaks/DefineOutline.cpp

@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "AST.h"
+#include "FindSymbols.h"
 #include "FindTarget.h"
 #include "HeaderSourceSwitch.h"
 #include "ParsedAST.h"
@@ -34,6 +35,7 @@
 #include <cstddef>
 #include <optional>
 #include <string>
+#include <tuple>
 
 namespace clang {
 namespace clangd {
@@ -489,8 +491,18 @@ class DefineOutline : public Tweak {
 
   Expected<Effect> apply(const Selection &Sel) override {
     const SourceManager &SM = Sel.AST->getSourceManager();
-    auto CCFile = SameFile ? Sel.AST->tuPath().str()
-                           : getSourceFile(Sel.AST->tuPath(), Sel);
+    std::optional<Path> CCFile;
+    auto Anchor = getDefinitionOfAdjacentDecl(Sel);
+    if (Anchor) {
+      std::get<0>(*Anchor).FileURI = std::get<1>(*Anchor).c_str();
+      auto URI = URI::resolve(std::get<0>(*Anchor).FileURI);
+      if (!URI)
+        return URI.takeError();
+      CCFile = *URI;
+    } else {
+      CCFile = SameFile ? Sel.AST->tuPath().str()
+                        : getSourceFile(Sel.AST->tuPath(), Sel);
+    }
     if (!CCFile)
       return error("Couldn't find a suitable implementation file.");
     assert(Sel.FS && "FS Must be set in apply");
@@ -499,21 +511,62 @@ class DefineOutline : public Tweak {
     // doesn't exist?
     if (!Buffer)
       return llvm::errorCodeToError(Buffer.getError());
+
+    std::optional<Position> InsertionPos;
+    if (Anchor) {
+      if (auto P = getInsertionPointFromExistingDefinition(
+              **Buffer, std::get<0>(*Anchor), std::get<2>(*Anchor), Sel.AST)) {
+        InsertionPos = *P;
+      }
+    }
+
     auto Contents = Buffer->get()->getBuffer();
-    auto InsertionPoint = getInsertionPoint(Contents, Sel);
-    if (!InsertionPoint)
-      return InsertionPoint.takeError();
+
+    std::size_t Offset = -1;
+    const DeclContext *EnclosingNamespace = nullptr;
+    std::string EnclosingNamespaceName;
+
+    if (InsertionPos) {
+      EnclosingNamespaceName = getNamespaceAtPosition(Contents, *InsertionPos,
+                                                      Sel.AST->getLangOpts());
+    } else if (SameFile) {
+      auto P = getInsertionPointInMainFile(Sel.AST);
+      if (!P)
+        return P.takeError();
+      Offset = P->Offset;
+      EnclosingNamespace = P->EnclosingNamespace;
+    } else {
+      auto Region = getEligiblePoints(
+          Contents, Source->getQualifiedNameAsString(), Sel.AST->getLangOpts());
+      assert(!Region.EligiblePoints.empty());
+      EnclosingNamespaceName = Region.EnclosingNamespace;
+      InsertionPos = Region.EligiblePoints.back();
+    }
+
+    if (InsertionPos) {
+      auto O = positionToOffset(Contents, *InsertionPos);
+      if (!O)
+        return O.takeError();
+      Offset = *O;
+      auto TargetContext =
+          findContextForNS(EnclosingNamespaceName, Source->getDeclContext());
+      if (!TargetContext)
+        return error("define outline: couldn't find a context for target");
+      EnclosingNamespace = *TargetContext;
+    }
+
+    assert(Offset >= 0);
+    assert(EnclosingNamespace);
 
     auto FuncDef = getFunctionSourceCode(
-        Source, InsertionPoint->EnclosingNamespace, Sel.AST->getTokens(),
+        Source, EnclosingNamespace, Sel.AST->getTokens(),
         Sel.AST->getHeuristicResolver(),
         SameFile && isHeaderFile(Sel.AST->tuPath(), Sel.AST->getLangOpts()));
     if (!FuncDef)
       return FuncDef.takeError();
 
     SourceManagerForFile SMFF(*CCFile, Contents);
-    const tooling::Replacement InsertFunctionDef(
-        *CCFile, InsertionPoint->Offset, 0, *FuncDef);
+    const tooling::Replacement InsertFunctionDef(*CCFile, Offset, 0, *FuncDef);
     auto Effect = Effect::mainFileEdit(
         SMFF.get(), tooling::Replacements(InsertFunctionDef));
     if (!Effect)
@@ -548,59 +601,193 @@ class DefineOutline : public Tweak {
     return std::move(*Effect);
   }
 
-  // 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.
-  // FIXME: The selection can be made smarter by looking at the definition
-  // locations for adjacent decls to Source. Unfortunately pseudo parsing in
-  // getEligibleRegions only knows about namespace begin/end events so we
-  // can't match function start/end positions yet.
-  llvm::Expected<InsertionPoint> getInsertionPoint(llvm::StringRef Contents,
-                                                   const Selection &Sel) {
-    // If the definition goes to the same file and there is a namespace,
-    // we should (and, in the case of anonymous namespaces, need to)
-    // put the definition into the original namespace block.
-    if (SameFile) {
-      auto *Klass = Source->getDeclContext()->getOuterLexicalRecordContext();
-      if (!Klass)
-        return error("moving to same file not supported for free functions");
-      const SourceLocation EndLoc = Klass->getBraceRange().getEnd();
-      const auto &TokBuf = Sel.AST->getTokens();
-      auto Tokens = TokBuf.expandedTokens();
-      auto It = llvm::lower_bound(
-          Tokens, EndLoc, [](const syntax::Token &Tok, SourceLocation EndLoc) {
-            return Tok.location() < EndLoc;
-          });
-      while (It != Tokens.end()) {
-        if (It->kind() != tok::semi) {
-          ++It;
-          continue;
+  enum class RelativeInsertPos { Before, After };
+  std::optional<std::tuple<SymbolLocation, Path, RelativeInsertPos>>
+  getDefinitionOfAdjacentDecl(const Selection &Sel) {
+    if (!Sel.Index)
+      return {};
+    std::optional<std::pair<SymbolLocation, Path>> Anchor;
+    std::string TuURI = URI::createFile(Sel.AST->tuPath()).toString();
+    auto CheckCandidate = [&](Decl *Candidate) {
+      assert(Candidate != Source);
+      if (auto Func = llvm::dyn_cast_or_null<FunctionDecl>(Candidate);
+          !Func || Func->isThisDeclarationADefinition()) {
+        return;
+      }
+      std::optional<std::pair<SymbolLocation, Path>> CandidateLoc;
+      Sel.Index->lookup({{getSymbolID(Candidate)}}, [&](const Symbol &S) {
+        if (S.Definition) {
+          CandidateLoc = std::make_pair(S.Definition,
+                                        StringRef(S.Definition.FileURI).str());
         }
-        unsigned Offset = Sel.AST->getSourceManager()
-                              .getDecomposedLoc(It->endLocation())
-                              .second;
-        return InsertionPoint{Klass->getEnclosingNamespaceContext(), Offset};
+      });
+      if (!CandidateLoc)
+        return;
+
+      // If our definition is constrained to the same file, ignore
+      // definitions that are not located there.
+      // If our definition is not constrained to the same file, but
+      // our anchor definition is in the same file, then we also put our
+      // definition there, because that appears to be the user preference.
+      // Exception: If the existing definition is a template, then the
+      // location is likely due to technical necessity rather than preference,
+      // so ignore that definition.
+      bool CandidateSameFile = TuURI == CandidateLoc->second;
+      if (SameFile && !CandidateSameFile)
+        return;
+      if (!SameFile && CandidateSameFile) {
+        if (Candidate->isTemplateDecl())
+          return;
+        SameFile = true;
       }
-      return error(
-          "failed to determine insertion location: no end of class found");
+      Anchor = *CandidateLoc;
+    };
+
+    // Try to find adjacent function declarations.
+    // Determine the closest one by alternatingly going "up" and "down"
+    // from our function in increasing steps.
+    const DeclContext *ParentContext = Source->getParent();
+    const auto SourceIt = llvm::find_if(
+        ParentContext->decls(), [this](const Decl *D) { return D == Source; });
+    if (SourceIt == ParentContext->decls_end())
+      return {};
+    const int Preceding = std::distance(ParentContext->decls_begin(), SourceIt);
+    const int Following =
+        std::distance(SourceIt, ParentContext->decls_end()) - 1;
+    for (int Offset = 1; Offset <= Preceding || Offset <= Following; ++Offset) {
+      if (Offset <= Preceding)
+        CheckCandidate(
+            *std::next(ParentContext->decls_begin(), Preceding - Offset));
+      if (Anchor)
+        return std::make_tuple(Anchor->first, Anchor->second,
+                               RelativeInsertPos::After);
+      if (Offset <= Following)
+        CheckCandidate(*std::next(SourceIt, Offset));
+      if (Anchor)
+        return std::make_tuple(Anchor->first, Anchor->second,
+                               RelativeInsertPos::Before);
     }
+    return {};
+  }
 
-    auto Region = getEligiblePoints(
-        Contents, Source->getQualifiedNameAsString(), Sel.AST->getLangOpts());
+  // We don't know the actual start or end of the definition, only the position
+  // of the name. Therefore, we heuristically try to locate the last token
+  // before or in this function, respectively. Adapt as required by user code.
+  llvm::Expected<Position> getInsertionPointFromExistingDefinition(
+      const llvm::MemoryBuffer &Buffer, const SymbolLocation &Loc,
+      RelativeInsertPos RelInsertPos, ParsedAST *AST) {
+    auto LspPos = indexToLSPLocation(Loc, AST->tuPath());
+    if (!LspPos)
+      return LspPos.takeError();
+    auto StartOffset =
+        positionToOffset(Buffer.getBuffer(), LspPos->range.start);
+    if (!StartOffset)
+      return LspPos.takeError();
+    SourceLocation InsertionLoc;
+    SourceManager &SM = AST->getSourceManager();
+    FileID F = Buffer.getBufferIdentifier() == AST->tuPath()
+                   ? SM.getMainFileID()
+                   : SM.createFileID(Buffer);
+
+    auto InsertBefore = [&] {
+      // Go backwards until we encounter one of the following:
+      //   - An opening brace (of a namespace).
+      //   - A closing brace (of a function definition).
+      //   - A semicolon (of a declaration).
+      // If no such token was found, then the first token in the file starts the
+      // definition.
+      auto Tokens = syntax::tokenize(syntax::FileRange(F, 0, *StartOffset), SM,
+                                     AST->getLangOpts());
+      if (Tokens.empty())
+        return;
+      for (auto I = std::rbegin(Tokens);
+           InsertionLoc.isInvalid() && I != std::rend(Tokens); ++I) {
+        switch (I->kind()) {
+        case tok::l_brace:
+        case tok::r_brace:
+        case tok::semi:
+          if (I != std::rbegin(Tokens))
+            InsertionLoc = std::prev(I)->location();
+          else
+            InsertionLoc = I->endLocation();
+          break;
+        default:
+          break;
+        }
+      }
+      if (InsertionLoc.isInvalid())
+        InsertionLoc = Tokens.front().location();
+    };
 
-    assert(!Region.EligiblePoints.empty());
-    auto Offset = positionToOffset(Contents, Region.EligiblePoints.back());
-    if (!Offset)
-      return Offset.takeError();
+    if (RelInsertPos == RelativeInsertPos::Before) {
+      InsertBefore();
+    } else {
+      // Skip over one top-level pair of parentheses (for the parameter list)
+      // and one pair of curly braces (for the code block).
+      // If that fails, insert before the function instead.
+      auto Tokens = syntax::tokenize(
+          syntax::FileRange(F, *StartOffset, Buffer.getBuffer().size()), SM,
+          AST->getLangOpts());
+      bool SkippedParams = false;
+      int OpenParens = 0;
+      int OpenBraces = 0;
+      std::optional<syntax::Token> Tok;
+      for (const auto &T : Tokens) {
+        tok::TokenKind StartKind = SkippedParams ? tok::l_brace : tok::l_paren;
+        tok::TokenKind EndKind = SkippedParams ? tok::r_brace : tok::r_paren;
+        int &Count = SkippedParams ? OpenBraces : OpenParens;
+        if (T.kind() == StartKind) {
+          ++Count;
+        } else if (T.kind() == EndKind) {
+          if (--Count == 0) {
+            if (SkippedParams) {
+              Tok = T;
+              break;
+            }
+            SkippedParams = true;
+          } else if (Count < 0) {
+            break;
+          }
+        }
+      }
+      if (Tok)
+        InsertionLoc = Tok->endLocation();
+      else
+        InsertBefore();
+    }
 
-    auto TargetContext =
-        findContextForNS(Region.EnclosingNamespace, Source->getDeclContext());
-    if (!TargetContext)
-      return error("define outline: couldn't find a context for target");
+    return InsertionLoc.isValid() ? sourceLocToPosition(SM, InsertionLoc)
+                                  : Position();
+  }
 
-    return InsertionPoint{*TargetContext, *Offset};
+  // Returns the most natural insertion point in this file.
+  // This is a fallback for when we failed to find an existing definition to
+  // place the new one next to. It only considers namespace proximity.
+  llvm::Expected<InsertionPoint> getInsertionPointInMainFile(ParsedAST *AST) {
+    // If the definition goes to the same file and there is a namespace,
+    // we should (and, in the case of anonymous namespaces, need to)
+    // put the definition into the original namespace block.
+    auto *Klass = Source->getDeclContext()->getOuterLexicalRecordContext();
+    if (!Klass)
+      return error("moving to same file not supported for free functions");
+    const SourceLocation EndLoc = Klass->getBraceRange().getEnd();
+    const auto &TokBuf = AST->getTokens();
+    auto Tokens = TokBuf.expandedTokens();
+    auto It = llvm::lower_bound(
+        Tokens, EndLoc, [](const syntax::Token &Tok, SourceLocation EndLoc) {
+          return Tok.location() < EndLoc;
+        });
+    while (It != Tokens.end()) {
+      if (It->kind() != tok::semi) {
+        ++It;
+        continue;
+      }
+      unsigned Offset =
+          AST->getSourceManager().getDecomposedLoc(It->endLocation()).second;
+      return InsertionPoint{Klass->getEnclosingNamespaceContext(), Offset};
+    }
+    return error(
+        "failed to determine insertion location: no end of class found");
   }
 
 private: