[clangd] Simplify `insertDecl` and support insertion after the last declaration in a file

[avogelsgesang]

Summary

So far insertDecl, did not support insertion of new code after the
last top-level declaration in a file. This was, because a
TranslationUnitDecl does not have a valid getEndLoc().

Also, the logic of Anchor::Below was somewhat surprising: I would have
expected Below to insert below the first match. Instead it insert
after the first contigous run of matches. I found this highly
unintuitive, and replaced it by an additional optional First/Last
flag which can be specified for an anchor.

clang-tools-extra/clangd/refactor/InsertionPoint.h

 // It allows inserting above or below a block of decls matching some criterion.
 // For example, "insert after existing constructors".
 struct Anchor {
   // A predicate describing which decls are considered part of a block.
   // Match need not handle TemplateDecls, which are unwrapped before matching.
   std::function<bool(const Decl *)> Match;
   // Whether the insertion point should be before or after the matching block.
   enum Dir { Above, Below } Direction = Below;
+  // Whether to choose the first or the last match
+  enum SearchDir { First, Last } SearchDirection = First;
 };
 
 // Returns the point to insert a declaration according to Anchors.
 // Anchors are tried in order. For each, the first matching location is chosen.
 SourceLocation insertionPoint(const DeclContext &Ctx,
                               llvm::ArrayRef<Anchor> Anchors);
 
 // Returns an edit inserting Code inside Ctx.

clang-tools-extra/clangd/refactor/InsertionPoint.cpp

 #include "clang/AST/DeclObjC.h"
 #include "clang/AST/DeclTemplate.h"
 #include "clang/Basic/SourceManager.h"
 
 namespace clang {
 namespace clangd {
 namespace {
 
-// Choose the decl to insert before, according to an anchor.
-// Nullptr means insert at end of DC.
-// None means no valid place to insert.
-llvm::Optional<const Decl *> insertionDecl(const DeclContext &DC,
-                                           const Anchor &A) {
-  bool LastMatched = false;
-  bool ReturnNext = false;
-  for (const auto *D : DC.decls()) {
-    if (D->isImplicit())
-      continue;
-    if (ReturnNext)
-      return D;
-
-    const Decl *NonTemplate = D;
-    if (auto *TD = llvm::dyn_cast<TemplateDecl>(D))
-      NonTemplate = TD->getTemplatedDecl();
-    bool Matches = A.Match(NonTemplate);
-    dlog("    {0} {1} {2}", Matches, D->getDeclKindName(), D);
-
-    switch (A.Direction) {
-    case Anchor::Above:
-      if (Matches && !LastMatched) {
-        // Special case: if "above" matches an access specifier, we actually
-        // want to insert below it!
-        if (llvm::isa<AccessSpecDecl>(D)) {
-          ReturnNext = true;
-          continue;
-        }
-        return D;
-      }
-      break;
-    case Anchor::Below:
-      if (LastMatched && !Matches)
-        return D;
-      break;
-    }
-
-    LastMatched = Matches;
-  }
-  if (ReturnNext || (LastMatched && A.Direction == Anchor::Below))
-    return nullptr;
-  return llvm::None;
-}
-
 SourceLocation beginLoc(const Decl &D) {
   auto Loc = D.getBeginLoc();
   if (RawComment *Comment = D.getASTContext().getRawCommentForDeclNoCache(&D)) {
     auto CommentLoc = Comment->getBeginLoc();
     if (CommentLoc.isValid() && Loc.isValid() &&
         D.getASTContext().getSourceManager().isBeforeInTranslationUnit(
             CommentLoc, Loc))
       Loc = CommentLoc;

 AccessSpecifier getAccessAtEnd(const CXXRecordDecl &C) {
   AccessSpecifier Spec = (C.getTagKind() == TTK_Class ? AS_private : AS_public);
   for (const auto *D : C.decls())
     if (const auto *ASD = llvm::dyn_cast<AccessSpecDecl>(D))
       Spec = ASD->getAccess();
   return Spec;
 }
 
+SourceLocation insertionPoint(const DeclContext &DC, Anchor A) {
+  SourceLocation Result;
+  for (auto DIter = DC.decls_begin(); DIter != DC.decls_end(); ++DIter) {
+    const auto *D = *DIter;
+    if (D->isImplicit())
+      continue;
+
+    const Decl *NonTemplate = D;
+    if (auto *TD = llvm::dyn_cast<TemplateDecl>(D))
+      NonTemplate = TD->getTemplatedDecl();
+    bool Matches = A.Match(NonTemplate);
+    dlog("    {0} {1} {2}", Matches, D->getDeclKindName(), D);
+
+    if (Matches) {
+      switch (A.Direction) {
+      case Anchor::Above:
+        // Special case: if "above" matches an access specifier, we actually
+        // want to insert below it!
+        if (!llvm::isa<AccessSpecDecl>(D)) {
+          Result = beginLoc(*D);
+          break;
+        }
+        [[fallthrough]];
+      case Anchor::Below:
+        // We cannot use the `getEndLoc`, because the clang AST does not track,
+        // semicolons. "Below" hence actually means "Before the next one".
+        auto NextIter = DIter;
+        ++NextIter;
+        if (NextIter != DC.decls_end())
+          Result = beginLoc(**NextIter);
+        else if (!DC.isTranslationUnit())
+          // There was no next item. Use the endLoc instead.
+          return endLoc(DC);
+        else {
+          // TranslationUnits don't have a valid endLoc. We need to use the
+          // source manager instead to insert all the way at the bottom of the
+          // file.
+          auto EndLoc = D->getEndLoc();
+          const auto &SM = DC.getParentASTContext().getSourceManager();
+          return SM.getLocForEndOfFile(SM.getFileID(EndLoc));
+        }
+      }
+      if (A.SearchDirection == Anchor::First)
+        break;
+    }
+  }
+  return Result;
+}
+
 } // namespace
 
 SourceLocation insertionPoint(const DeclContext &DC,
                               llvm::ArrayRef<Anchor> Anchors) {
   dlog("Looking for insertion point in {0}", DC.getDeclKindName());
   for (const auto &A : Anchors) {
     dlog("  anchor ({0})", A.Direction == Anchor::Above ? "above" : "below");
-    if (auto D = insertionDecl(DC, A)) {
-      dlog("  anchor matched before {0}", *D);
-      return *D ? beginLoc(**D) : endLoc(DC);
+    auto L = insertionPoint(DC, A);
+    if (L.isValid()) {
+      return L;
     }
   }
   dlog("no anchor matched");
   return SourceLocation();
 }
 
 llvm::Expected<tooling::Replacement>
 insertDecl(llvm::StringRef Code, const DeclContext &DC,
SourceLocation insertionPoint(const CXXRecordDecl &InClass,
   return insertionPoint(InClass, Anchors);
 }
 
 llvm::Expected<tooling::Replacement> insertDecl(llvm::StringRef Code,
                                                 const CXXRecordDecl &InClass,
                                                 std::vector<Anchor> Anchors,
                                                 AccessSpecifier Protection) {
   // Fallback: insert at the bottom of the relevant access section.
-  Anchors.push_back({any, Anchor::Below});
+  Anchors.push_back({any, Anchor::Below, Anchor::Last});
   auto Loc = insertionPoint(InClass, std::move(Anchors), Protection);
   std::string CodeBuffer;
   auto &SM = InClass.getASTContext().getSourceManager();
   // Fallback: insert at the end of the class. Check if protection matches!
   if (Loc.isInvalid()) {
     Loc = InClass.getBraceRange().getEnd();
     if (Protection != getAccessAtEnd(InClass)) {
       CodeBuffer = (getAccessSpelling(Protection) + ":\n" + Code).str();

clang-tools-extra/clangd/unittests/InsertionPointTests.cpp

 
 TEST(InsertionPointTests, Generic) {
   Annotations Code(R"cpp(
   namespace ns {
     $a^int a1;
     $b^// leading comment
     int b;
     $c^int c1; // trailing comment
-    int c2;
+    $c2^int c2;
     $a2^int a2;
   $end^};
   )cpp");
 
   auto StartsWith =
       [&](llvm::StringLiteral S) -> std::function<bool(const Decl *)> {
     return [S](const Decl *D) {
       if (const auto *ND = llvm::dyn_cast<NamedDecl>(D))
         return llvm::StringRef(ND->getNameAsString()).startswith(S);
       return false;
     };
   };
 
   auto AST = TestTU::withCode(Code.code()).build();
   auto &NS = cast<NamespaceDecl>(findDecl(AST, "ns"));
 
   // Test single anchors.
-  auto Point = [&](llvm::StringLiteral Prefix, Anchor::Dir Direction) {
-    auto Loc = insertionPoint(NS, {Anchor{StartsWith(Prefix), Direction}});
+  auto Point = [&](llvm::StringLiteral Prefix, Anchor::Dir Direction,
+                   Anchor::SearchDir SearchDirection = Anchor::First) {
+    auto Loc = insertionPoint(
+        NS, {Anchor{StartsWith(Prefix), Direction, SearchDirection}});
     return sourceLocToPosition(AST.getSourceManager(), Loc);
   };
   EXPECT_EQ(Point("a", Anchor::Above), Code.point("a"));
   EXPECT_EQ(Point("a", Anchor::Below), Code.point("b"));
   EXPECT_EQ(Point("b", Anchor::Above), Code.point("b"));
   EXPECT_EQ(Point("b", Anchor::Below), Code.point("c"));
-  EXPECT_EQ(Point("c", Anchor::Above), Code.point("c"));
-  EXPECT_EQ(Point("c", Anchor::Below), Code.point("a2"));
+  EXPECT_EQ(Point("c", Anchor::Above, Anchor::First), Code.point("c"));
+  EXPECT_EQ(Point("c", Anchor::Below, Anchor::First), Code.point("c2"));
+  EXPECT_EQ(Point("c", Anchor::Above, Anchor::Last), Code.point("c2"));
+  EXPECT_EQ(Point("c", Anchor::Below, Anchor::Last), Code.point("a2"));
+  EXPECT_EQ(Point("a2", Anchor::Above), Code.point("a2"));
+  EXPECT_EQ(Point("a2", Anchor::Below), Code.point("end"));
   EXPECT_EQ(Point("", Anchor::Above), Code.point("a"));
-  EXPECT_EQ(Point("", Anchor::Below), Code.point("end"));
+  EXPECT_EQ(Point("", Anchor::Below), Code.point("b"));
   EXPECT_EQ(Point("no_match", Anchor::Below), Position{});
 
   // Test anchor chaining.
   auto Chain = [&](llvm::StringLiteral P1, llvm::StringLiteral P2) {
     auto Loc = insertionPoint(NS, {Anchor{StartsWith(P1), Anchor::Above},
                                    Anchor{StartsWith(P2), Anchor::Above}});
     return sourceLocToPosition(AST.getSourceManager(), Loc);
   };
   EXPECT_EQ(Chain("a", "b"), Code.point("a"));
   EXPECT_EQ(Chain("b", "a"), Code.point("b"));
   EXPECT_EQ(Chain("no_match", "a"), Code.point("a"));
 
   // Test edit generation.
   auto Edit = insertDecl("foo;", NS, {Anchor{StartsWith("a"), Anchor::Below}});
   ASSERT_THAT_EXPECTED(Edit, llvm::Succeeded());
   EXPECT_EQ(offsetToPosition(Code.code(), Edit->getOffset()), Code.point("b"));
   EXPECT_EQ(Edit->getReplacementText(), "foo;");
   // If no match, the edit is inserted at the end.
   Edit = insertDecl("x;", NS, {Anchor{StartsWith("no_match"), Anchor::Below}});
   ASSERT_THAT_EXPECTED(Edit, llvm::Succeeded());
   EXPECT_EQ(offsetToPosition(Code.code(), Edit->getOffset()),
             Code.point("end"));
 }
 
+TEST(InsertionPointTests, TopLevel) {
+  Annotations Code(R"cpp(
+  $a^int a;
+  $b^int b;
+  $end^)cpp");
+
+  auto HasName =
+      [&](llvm::StringLiteral S) -> std::function<bool(const Decl *)> {
+    return [S](const Decl *D) {
+      if (const auto *ND = llvm::dyn_cast<NamedDecl>(D))
+        return ND->getNameAsString() == S;
+      return false;
+    };
+  };
+
+  auto AST = TestTU::withCode(Code.code()).build();
+  auto &TUD = *AST.getASTContext().getTranslationUnitDecl();
+
+  // Test single anchors.
+  auto Point = [&](llvm::StringLiteral Prefix, Anchor::Dir Direction) {
+    auto Loc = insertionPoint(TUD, {Anchor{HasName(Prefix), Direction}});
+    return sourceLocToPosition(AST.getSourceManager(), Loc);
+  };
+  EXPECT_EQ(Point("a", Anchor::Above), Code.point("a"));
+  EXPECT_EQ(Point("a", Anchor::Below), Code.point("b"));
+  EXPECT_EQ(Point("b", Anchor::Above), Code.point("b"));
+  EXPECT_EQ(Point("b", Anchor::Below), Code.point("end"));
+}
+
 // For CXX, we should check:
 // - special handling for access specifiers
 // - unwrapping of template decls
 TEST(InsertionPointTests, CXX) {
   Annotations Code(R"cpp(
     class C {
     public:
       $Method^void pubMethod();
       $Field^int PubField;
 
     $private^private:
       $field^int PrivField;
       $method^void privMethod();
-      template <typename T> void privTemplateMethod();
+      $template^template <typename T> void privTemplateMethod();
     $end^};
   )cpp");
 
   auto AST = TestTU::withCode(Code.code()).build();
   const CXXRecordDecl &C = cast<CXXRecordDecl>(findDecl(AST, "C"));
 
   auto IsMethod = [](const Decl *D) { return llvm::isa<CXXMethodDecl>(D); };
   auto Any = [](const Decl *D) { return true; };
 
   // Test single anchors.
   auto Point = [&](Anchor A, AccessSpecifier Protection) {
     auto Loc = insertionPoint(C, {A}, Protection);
     return sourceLocToPosition(AST.getSourceManager(), Loc);
   };
   EXPECT_EQ(Point({IsMethod, Anchor::Above}, AS_public), Code.point("Method"));
   EXPECT_EQ(Point({IsMethod, Anchor::Below}, AS_public), Code.point("Field"));
   EXPECT_EQ(Point({Any, Anchor::Above}, AS_public), Code.point("Method"));
-  EXPECT_EQ(Point({Any, Anchor::Below}, AS_public), Code.point("private"));
+  EXPECT_EQ(Point({Any, Anchor::Below}, AS_public), Code.point("Method"));
   EXPECT_EQ(Point({IsMethod, Anchor::Above}, AS_private), Code.point("method"));
-  EXPECT_EQ(Point({IsMethod, Anchor::Below}, AS_private), Code.point("end"));
+  EXPECT_EQ(Point({IsMethod, Anchor::Below}, AS_private),
+            Code.point("template"));
   EXPECT_EQ(Point({Any, Anchor::Above}, AS_private), Code.point("field"));
-  EXPECT_EQ(Point({Any, Anchor::Below}, AS_private), Code.point("end"));
+  EXPECT_EQ(Point({Any, Anchor::Below}, AS_private), Code.point("field"));
   EXPECT_EQ(Point({IsMethod, Anchor::Above}, AS_protected), Position{});
   EXPECT_EQ(Point({IsMethod, Anchor::Below}, AS_protected), Position{});
   EXPECT_EQ(Point({Any, Anchor::Above}, AS_protected), Position{});
   EXPECT_EQ(Point({Any, Anchor::Below}, AS_protected), Position{});
 
   // Edits when there's no match --> end of matching access control section.
   auto Edit = insertDecl("x", C, {}, AS_public);
   ASSERT_THAT_EXPECTED(Edit, llvm::Succeeded());