[clangd] Helper for determining member insertion point.

To be used in D116490 and D116385, and an upcoming patch to generate C++
constructors.

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

clang-tools-extra/clangd/CMakeLists.txt

@@ -131,6 +131,7 @@ add_clang_library(clangDaemon
   index/dex/PostingList.cpp
   index/dex/Trigram.cpp
 
+  refactor/InsertionPoint.cpp
   refactor/Rename.cpp
   refactor/Tweak.cpp
 

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

@@ -0,0 +1,157 @@
+//===--- InsertionPoint.cpp - Where should we add new code? ---------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "refactor/InsertionPoint.h"
+#include "support/Logger.h"
+#include "clang/AST/ASTContext.h"
+#include "clang/AST/DeclCXX.h"
+#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;
+  }
+  return Loc;
+}
+
+bool any(const Decl *D) { return true; }
+
+SourceLocation endLoc(const DeclContext &DC) {
+  const Decl *D = llvm::cast<Decl>(&DC);
+  if (auto *OCD = llvm::dyn_cast<ObjCContainerDecl>(D))
+    return OCD->getAtEndRange().getBegin();
+  return D->getEndLoc();
+}
+
+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;
+}
+
+} // 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);
+    }
+  }
+  dlog("no anchor matched");
+  return SourceLocation();
+}
+
+llvm::Expected<tooling::Replacement>
+insertDecl(llvm::StringRef Code, const DeclContext &DC,
+           llvm::ArrayRef<Anchor> Anchors) {
+  auto Loc = insertionPoint(DC, Anchors);
+  // Fallback: insert at the end.
+  if (Loc.isInvalid())
+    Loc = endLoc(DC);
+  const auto &SM = DC.getParentASTContext().getSourceManager();
+  if (!SM.isWrittenInSameFile(Loc, cast<Decl>(DC).getLocation()))
+    return error("{0} body in wrong file: {1}", DC.getDeclKindName(),
+                 Loc.printToString(SM));
+  return tooling::Replacement(SM, Loc, 0, Code);
+}
+
+SourceLocation insertionPoint(const CXXRecordDecl &InClass,
+                              std::vector<Anchor> Anchors,
+                              AccessSpecifier Protection) {
+  for (auto &A : Anchors)
+    A.Match = [Inner(std::move(A.Match)), Protection](const Decl *D) {
+      return D->getAccess() == Protection && Inner(D);
+    };
+  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});
+  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();
+      Code = CodeBuffer;
+    }
+  }
+  if (!SM.isWrittenInSameFile(Loc, InClass.getLocation()))
+    return error("Class body in wrong file: {0}", Loc.printToString(SM));
+  return tooling::Replacement(SM, Loc, 0, Code);
+}
+
+} // namespace clangd
+} // namespace clang

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

@@ -0,0 +1,53 @@
+//===--- InsertionPoint.h - Where should we add new code? --------*- C++-*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "clang/AST/DeclCXX.h"
+#include "clang/Basic/Specifiers.h"
+#include "clang/Tooling/Core/Replacement.h"
+
+namespace clang {
+namespace clangd {
+
+// An anchor describes where to insert code into a decl sequence.
+//
+// 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;
+};
+
+// 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.
+// Location is chosen according to Anchors, falling back to the end of Ctx.
+// Fails if the chosen insertion point is in a different file than Ctx itself.
+llvm::Expected<tooling::Replacement> insertDecl(llvm::StringRef Code,
+                                                const DeclContext &Ctx,
+                                                llvm::ArrayRef<Anchor> Anchors);
+
+// Variant for C++ classes that ensures the right access control.
+SourceLocation insertionPoint(const CXXRecordDecl &InClass,
+                              std::vector<Anchor> Anchors,
+                              AccessSpecifier Protection);
+
+// Variant for C++ classes that ensures the right access control.
+// May insert a new access specifier if needed.
+llvm::Expected<tooling::Replacement> insertDecl(llvm::StringRef Code,
+                                                const CXXRecordDecl &InClass,
+                                                std::vector<Anchor> Anchors,
+                                                AccessSpecifier Protection);
+
+} // namespace clangd
+} // namespace clang

clang-tools-extra/clangd/unittests/CMakeLists.txt

@@ -62,6 +62,7 @@ add_unittest(ClangdUnitTests ClangdTests
   IndexActionTests.cpp
   IndexTests.cpp
   InlayHintTests.cpp
+  InsertionPointTests.cpp
   JSONTransportTests.cpp
   LoggerTests.cpp
   LSPBinderTests.cpp

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

@@ -0,0 +1,210 @@
+//===-- InsertionPointTess.cpp  -------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "Annotations.h"
+#include "Protocol.h"
+#include "SourceCode.h"
+#include "TestTU.h"
+#include "TestWorkspace.h"
+#include "XRefs.h"
+#include "refactor/InsertionPoint.h"
+#include "clang/AST/DeclBase.h"
+#include "llvm/Testing/Support/Error.h"
+#include "gmock/gmock.h"
+#include "gtest/gtest.h"
+
+namespace clang {
+namespace clangd {
+namespace {
+using llvm::HasValue;
+
+TEST(InsertionPointTests, Generic) {
+  Annotations Code(R"cpp(
+  namespace ns {
+    $a^int a1;
+    $b^// leading comment
+    int b;
+    $c^int c1; // trailing comment
+    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}});
+    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("", Anchor::Above), Code.point("a"));
+  EXPECT_EQ(Point("", Anchor::Below), Code.point("end"));
+  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"));
+}
+
+// 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();
+    $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({IsMethod, Anchor::Above}, AS_private), Code.point("method"));
+  EXPECT_EQ(Point({IsMethod, Anchor::Below}, AS_private), Code.point("end"));
+  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({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());
+  EXPECT_EQ(offsetToPosition(Code.code(), Edit->getOffset()),
+            Code.point("private"));
+
+  Edit = insertDecl("x", C, {}, AS_private);
+  ASSERT_THAT_EXPECTED(Edit, llvm::Succeeded());
+  EXPECT_EQ(offsetToPosition(Code.code(), Edit->getOffset()),
+            Code.point("end"));
+
+  Edit = insertDecl("x", C, {}, AS_protected);
+  ASSERT_THAT_EXPECTED(Edit, llvm::Succeeded());
+  EXPECT_EQ(offsetToPosition(Code.code(), Edit->getOffset()),
+            Code.point("end"));
+  EXPECT_EQ(Edit->getReplacementText(), "protected:\nx");
+}
+
+MATCHER_P(replacementText, Text, "") {
+  if (arg.getReplacementText() != Text) {
+    *result_listener << "replacement is " << arg.getReplacementText().str();
+    return false;
+  }
+  return true;
+}
+
+TEST(InsertionPointTests, CXXAccessProtection) {
+  // Empty class uses default access.
+  auto AST = TestTU::withCode("struct S{};").build();
+  const CXXRecordDecl &S = cast<CXXRecordDecl>(findDecl(AST, "S"));
+  ASSERT_THAT_EXPECTED(insertDecl("x", S, {}, AS_public),
+                       HasValue(replacementText("x")));
+  ASSERT_THAT_EXPECTED(insertDecl("x", S, {}, AS_private),
+                       HasValue(replacementText("private:\nx")));
+
+  // We won't insert above the first access specifier if there's nothing there.
+  AST = TestTU::withCode("struct T{private:};").build();
+  const CXXRecordDecl &T = cast<CXXRecordDecl>(findDecl(AST, "T"));
+  ASSERT_THAT_EXPECTED(insertDecl("x", T, {}, AS_public),
+                       HasValue(replacementText("public:\nx")));
+  ASSERT_THAT_EXPECTED(insertDecl("x", T, {}, AS_private),
+                       HasValue(replacementText("x")));
+
+  // But we will if there are declarations.
+  AST = TestTU::withCode("struct U{int i;private:};").build();
+  const CXXRecordDecl &U = cast<CXXRecordDecl>(findDecl(AST, "U"));
+  ASSERT_THAT_EXPECTED(insertDecl("x", U, {}, AS_public),
+                       HasValue(replacementText("x")));
+  ASSERT_THAT_EXPECTED(insertDecl("x", U, {}, AS_private),
+                       HasValue(replacementText("x")));
+}
+
+// In ObjC we need to take care to get the @end fallback right.
+TEST(InsertionPointTests, ObjC) {
+  Annotations Code(R"objc(
+    @interface Foo
+     -(void) v;
+    $endIface^@end
+    @implementation Foo
+     -(void) v {}
+    $endImpl^@end
+  )objc");
+  auto TU = TestTU::withCode(Code.code());
+  TU.Filename = "TestTU.m";
+  auto AST = TU.build();
+
+  auto &Impl =
+      cast<ObjCImplementationDecl>(findDecl(AST, [&](const NamedDecl &D) {
+        return llvm::isa<ObjCImplementationDecl>(D);
+      }));
+  auto &Iface = *Impl.getClassInterface();
+  Anchor End{[](const Decl *) { return true; }, Anchor::Below};
+
+  const auto &SM = AST.getSourceManager();
+  EXPECT_EQ(sourceLocToPosition(SM, insertionPoint(Iface, {End})),
+            Code.point("endIface"));
+  EXPECT_EQ(sourceLocToPosition(SM, insertionPoint(Impl, {End})),
+            Code.point("endImpl"));
+}
+
+} // namespace
+} // namespace clangd
+} // namespace clang