[clangd] Support #pragma mark in the outline

Xcode uses `#pragma mark -` to draw a divider in the outline view
and `#pragma mark Note` to add `Note` in the outline view. For more
information, see https://nshipster.com/pragma/.

Since the LSP spec doesn't contain dividers for the symbol outline,
instead we treat `#pragma mark -` as a group with children - the
decls that come after it, implicitly terminating when the symbol's
parent ends.

The following code:

```
@implementation MyClass

- (id)init {}

- (int)foo;
@EnD
```

Would give an outline like

```
MyClass
        > Overrides
                    > init
        > Public Accessors
                    > foo
```

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

Author: DavidGoldman

Diff for: clang-tools-extra/clangd/CollectMacros.cpp

@@ -30,5 +30,33 @@ void CollectMainFileMacros::add(const Token &MacroNameTok, const MacroInfo *MI,
   else
     Out.UnknownMacros.push_back({Range, IsDefinition});
 }
+
+class CollectPragmaMarks : public PPCallbacks {
+public:
+  explicit CollectPragmaMarks(const SourceManager &SM,
+                              std::vector<clangd::PragmaMark> &Out)
+      : SM(SM), Out(Out) {}
+
+  void PragmaMark(SourceLocation Loc, StringRef Trivia) override {
+    if (isInsideMainFile(Loc, SM)) {
+      // FIXME: This range should just cover `XX` in `#pragma mark XX` and
+      // `- XX` in `#pragma mark - XX`.
+      Position Start = sourceLocToPosition(SM, Loc);
+      Position End = {Start.line + 1, 0};
+      Out.emplace_back(clangd::PragmaMark{{Start, End}, Trivia.str()});
+    }
+  }
+
+private:
+  const SourceManager &SM;
+  std::vector<clangd::PragmaMark> &Out;
+};
+
+std::unique_ptr<PPCallbacks>
+collectPragmaMarksCallback(const SourceManager &SM,
+                           std::vector<PragmaMark> &Out) {
+  return std::make_unique<CollectPragmaMarks>(SM, Out);
+}
+
 } // namespace clangd
 } // namespace clang

Diff for: clang-tools-extra/clangd/CollectMacros.h

@@ -99,6 +99,18 @@ class CollectMainFileMacros : public PPCallbacks {
   MainFileMacros &Out;
 };
 
+/// Represents a `#pragma mark` in the main file.
+///
+/// There can be at most one pragma mark per line.
+struct PragmaMark {
+  Range Rng;
+  std::string Trivia;
+};
+
+/// Collect all pragma marks from the main file.
+std::unique_ptr<PPCallbacks>
+collectPragmaMarksCallback(const SourceManager &, std::vector<PragmaMark> &Out);
+
 } // namespace clangd
 } // namespace clang
 

Diff for: clang-tools-extra/clangd/FindSymbols.cpp

@@ -523,9 +523,135 @@ class DocumentOutline {
   ParsedAST *
 };
 
+struct PragmaMarkSymbol {
+  DocumentSymbol DocSym;
+  bool IsGroup;
+};
+
+/// Merge in `PragmaMarkSymbols`, sorted ascending by range, into the given
+/// `DocumentSymbol` tree.
+void mergePragmas(DocumentSymbol &Root, ArrayRef<PragmaMarkSymbol> Pragmas) {
+  while (!Pragmas.empty()) {
+    // We'll figure out where the Pragmas.front() should go.
+    PragmaMarkSymbol P = std::move(Pragmas.front());
+    Pragmas = Pragmas.drop_front();
+    DocumentSymbol *Cur = &Root;
+    while (Cur->range.contains(P.DocSym.range)) {
+      bool Swapped = false;
+      for (auto &C : Cur->children) {
+        // We assume at most 1 child can contain the pragma (as pragmas are on
+        // a single line, and children have disjoint ranges).
+        if (C.range.contains(P.DocSym.range)) {
+          Cur = &C;
+          Swapped = true;
+          break;
+        }
+      }
+      // Cur is the parent of P since none of the children contain P.
+      if (!Swapped)
+        break;
+    }
+    // Pragma isn't a group so we can just insert it and we are done.
+    if (!P.IsGroup) {
+      Cur->children.emplace_back(std::move(P.DocSym));
+      continue;
+    }
+    // Pragma is a group, so we need to figure out where it terminates:
+    // - If the next Pragma is not contained in Cur, P owns all of its
+    //   parent's children which occur after P.
+    // - If the next pragma is contained in Cur but actually belongs to one
+    //   of the parent's children, we temporarily skip over it and look at
+    //   the next pragma to decide where we end.
+    // - Otherwise nest all of its parent's children which occur after P but
+    //   before the next pragma.
+    bool TerminatedByNextPragma = false;
+    for (auto &NextPragma : Pragmas) {
+      // If we hit a pragma outside of Cur, the rest will be outside as well.
+      if (!Cur->range.contains(NextPragma.DocSym.range))
+        break;
+
+      // NextPragma cannot terminate P if it is nested inside a child, look for
+      // the next one.
+      if (llvm::any_of(Cur->children, [&NextPragma](const auto &Child) {
+            return Child.range.contains(NextPragma.DocSym.range);
+          }))
+        continue;
+
+      // Pragma owns all the children between P and NextPragma
+      auto It = llvm::partition(Cur->children,
+                                [&P, &NextPragma](const auto &S) -> bool {
+                                  return !(P.DocSym.range < S.range &&
+                                           S.range < NextPragma.DocSym.range);
+                                });
+      P.DocSym.children.assign(make_move_iterator(It),
+                               make_move_iterator(Cur->children.end()));
+      Cur->children.erase(It, Cur->children.end());
+      TerminatedByNextPragma = true;
+      break;
+    }
+    if (!TerminatedByNextPragma) {
+      // P is terminated by the end of current symbol, hence it owns all the
+      // children after P.
+      auto It = llvm::partition(Cur->children, [&P](const auto &S) -> bool {
+        return !(P.DocSym.range < S.range);
+      });
+      P.DocSym.children.assign(make_move_iterator(It),
+                               make_move_iterator(Cur->children.end()));
+      Cur->children.erase(It, Cur->children.end());
+    }
+    // Update the range for P to cover children and append to Cur.
+    for (DocumentSymbol &Sym : P.DocSym.children)
+      unionRanges(P.DocSym.range, Sym.range);
+    Cur->children.emplace_back(std::move(P.DocSym));
+  }
+}
+
+PragmaMarkSymbol markToSymbol(const PragmaMark &P) {
+  StringRef Name = StringRef(P.Trivia).trim();
+  bool IsGroup = false;
+  // "-\s+<group name>" or "<name>" after an initial trim. The former is
+  // considered a group, the latter just a mark. Like Xcode, we don't consider
+  // `-Foo` to be a group (space(s) after the `-` is required).
+  //
+  // We need to include a name here, otherwise editors won't properly render the
+  // symbol.
+  StringRef MaybeGroupName = Name;
+  if (MaybeGroupName.consume_front("-") &&
+      (MaybeGroupName.ltrim() != MaybeGroupName || MaybeGroupName.empty())) {
+    Name = MaybeGroupName.empty() ? "(unnamed group)" : MaybeGroupName.ltrim();
+    IsGroup = true;
+  } else if (Name.empty()) {
+    Name = "(unnamed mark)";
+  }
+  DocumentSymbol Sym;
+  Sym.name = Name.str();
+  Sym.kind = SymbolKind::File;
+  Sym.range = P.Rng;
+  Sym.selectionRange = P.Rng;
+  return {Sym, IsGroup};
+}
+
 std::vector<DocumentSymbol> collectDocSymbols(ParsedAST &AST) {
-  return DocumentOutline(AST).build();
+  std::vector<DocumentSymbol> Syms = DocumentOutline(AST).build();
+
+  const auto &PragmaMarks = AST.getMarks();
+  if (PragmaMarks.empty())
+    return Syms;
+
+  std::vector<PragmaMarkSymbol> Pragmas;
+  Pragmas.reserve(PragmaMarks.size());
+  for (const auto &P : PragmaMarks)
+    Pragmas.push_back(markToSymbol(P));
+  Range EntireFile = {
+      {0, 0},
+      {std::numeric_limits<int>::max(), std::numeric_limits<int>::max()}};
+  DocumentSymbol Root;
+  Root.children = std::move(Syms);
+  Root.range = EntireFile;
+  mergePragmas(Root, llvm::makeArrayRef(Pragmas));
+  return Root.children;
 }
+
 } // namespace
 
 llvm::Expected<std::vector<DocumentSymbol>> getDocumentSymbols(ParsedAST &AST) {

Diff for: clang-tools-extra/clangd/ParsedAST.cpp

@@ -449,6 +449,13 @@ ParsedAST::build(llvm::StringRef Filename, const ParseInputs &Inputs,
       std::make_unique<CollectMainFileMacros>(Clang->getSourceManager(),
                                               Macros));
 
+  std::vector<PragmaMark> Marks;
+  // FIXME: We need to patch the marks for stale preambles.
+  if (Preamble)
+    Marks = Preamble->Marks;
+  Clang->getPreprocessor().addPPCallbacks(
+      collectPragmaMarksCallback(Clang->getSourceManager(), Marks));
+
   // Copy over the includes from the preamble, then combine with the
   // non-preamble includes below.
   CanonicalIncludes CanonIncludes;
@@ -510,7 +517,7 @@ ParsedAST::build(llvm::StringRef Filename, const ParseInputs &Inputs,
   }
   return ParsedAST(Inputs.Version, std::move(Preamble), std::move(Clang),
                    std::move(Action), std::move(Tokens), std::move(Macros),
-                   std::move(ParsedDecls), std::move(Diags),
+                   std::move(Marks), std::move(ParsedDecls), std::move(Diags),
                    std::move(Includes), std::move(CanonIncludes));
 }
 
@@ -550,6 +557,7 @@ llvm::ArrayRef<Decl *> ParsedAST::getLocalTopLevelDecls() {
 }
 
 const MainFileMacros &ParsedAST::getMacros() const { return Macros; }
+const std::vector<PragmaMark> &ParsedAST::getMarks() const { return Marks; }
 
 std::size_t ParsedAST::getUsedBytes() const {
   auto &AST = getASTContext();
@@ -596,12 +604,14 @@ ParsedAST::ParsedAST(llvm::StringRef Version,
                      std::unique_ptr<CompilerInstance> Clang,
                      std::unique_ptr<FrontendAction> Action,
                      syntax::TokenBuffer Tokens, MainFileMacros Macros,
+                     std::vector<PragmaMark> Marks,
                      std::vector<Decl *> LocalTopLevelDecls,
                      llvm::Optional<std::vector<Diag>> Diags,
                      IncludeStructure Includes, CanonicalIncludes CanonIncludes)
     : Version(Version), Preamble(std::move(Preamble)), Clang(std::move(Clang)),
       Action(std::move(Action)), Tokens(std::move(Tokens)),
-      Macros(std::move(Macros)), Diags(std::move(Diags)),
+      Macros(std::move(Macros)), Marks(std::move(Marks)),
+      Diags(std::move(Diags)),
       LocalTopLevelDecls(std::move(LocalTopLevelDecls)),
       Includes(std::move(Includes)), CanonIncludes(std::move(CanonIncludes)) {
   Resolver = std::make_unique<HeuristicResolver>(getASTContext());

Diff for: clang-tools-extra/clangd/ParsedAST.h

@@ -101,6 +101,8 @@ class ParsedAST {
   /// Gets all macro references (definition, expansions) present in the main
   /// file, including those in the preamble region.
   const MainFileMacros &getMacros() const;
+  /// Gets all pragma marks in the main file.
+  const std::vector<PragmaMark> &getMarks() const;
   /// Tokens recorded while parsing the main file.
   /// (!) does not have tokens from the preamble.
   const syntax::TokenBuffer &getTokens() const { return Tokens; }
@@ -121,7 +123,8 @@ class ParsedAST {
             std::shared_ptr<const PreambleData> Preamble,
             std::unique_ptr<CompilerInstance> Clang,
             std::unique_ptr<FrontendAction> Action, syntax::TokenBuffer Tokens,
-            MainFileMacros Macros, std::vector<Decl *> LocalTopLevelDecls,
+            MainFileMacros Macros, std::vector<PragmaMark> Marks,
+            std::vector<Decl *> LocalTopLevelDecls,
             llvm::Optional<std::vector<Diag>> Diags, IncludeStructure Includes,
             CanonicalIncludes CanonIncludes);
 
@@ -144,6 +147,8 @@ class ParsedAST {
 
   /// All macro definitions and expansions in the main file.
   MainFileMacros Macros;
+  // Pragma marks in the main file.
+  std::vector<PragmaMark> Marks;
   // Data, stored after parsing. None if AST was built with a stale preamble.
   llvm::Optional<std::vector<Diag>> Diags;
   // Top-level decls inside the current file. Not that this does not include

Diff for: clang-tools-extra/clangd/Preamble.cpp

@@ -73,6 +73,8 @@ class CppFilePreambleCallbacks : public PreambleCallbacks {
 
   MainFileMacros takeMacros() { return std::move(Macros); }
 
+  std::vector<PragmaMark> takeMarks() { return std::move(Marks); }
+
   CanonicalIncludes takeCanonicalIncludes() { return std::move(CanonIncludes); }
 
   bool isMainFileIncludeGuarded() const { return IsMainFileIncludeGuarded; }
@@ -103,7 +105,9 @@ class CppFilePreambleCallbacks : public PreambleCallbacks {
 
     return std::make_unique<PPChainedCallbacks>(
         collectIncludeStructureCallback(*SourceMgr, &Includes),
-        std::make_unique<CollectMainFileMacros>(*SourceMgr, Macros));
+        std::make_unique<PPChainedCallbacks>(
+            std::make_unique<CollectMainFileMacros>(*SourceMgr, Macros),
+            collectPragmaMarksCallback(*SourceMgr, Marks)));
   }
 
   CommentHandler *getCommentHandler() override {
@@ -130,6 +134,7 @@ class CppFilePreambleCallbacks : public PreambleCallbacks {
   IncludeStructure Includes;
   CanonicalIncludes CanonIncludes;
   MainFileMacros Macros;
+  std::vector<PragmaMark> Marks;
   bool IsMainFileIncludeGuarded = false;
   std::unique_ptr<CommentHandler> IWYUHandler = nullptr;
   const clang::LangOptions *LangOpts = nullptr;
@@ -389,6 +394,7 @@ buildPreamble(PathRef FileName, CompilerInvocation CI,
     Result->Diags = std::move(Diags);
     Result->Includes = CapturedInfo.takeIncludes();
     Result->Macros = CapturedInfo.takeMacros();
+    Result->Marks = CapturedInfo.takeMarks();
     Result->CanonIncludes = CapturedInfo.takeCanonicalIncludes();
     Result->StatCache = std::move(StatCache);
     Result->MainIsIncludeGuarded = CapturedInfo.isMainFileIncludeGuarded();

Diff for: clang-tools-extra/clangd/Preamble.h

@@ -61,6 +61,8 @@ struct PreambleData {
   // Users care about headers vs main-file, not preamble vs non-preamble.
   // These should be treated as main-file entities e.g. for code completion.
   MainFileMacros Macros;
+  // Pragma marks defined in the preamble section of the main file.
+  std::vector<PragmaMark> Marks;
   // Cache of FS operations performed when building the preamble.
   // When reusing a preamble, this cache can be consumed to save IO.
   std::unique_ptr<PreambleFileStatusCache> StatCache;

Diff for: clang-tools-extra/clangd/SourceCode.cpp

@@ -471,6 +471,13 @@ Range halfOpenToRange(const SourceManager &SM, CharSourceRange R) {
   return {Begin, End};
 }
 
+void unionRanges(Range &A, Range B) {
+  if (B.start < A.start)
+    A.start = B.start;
+  if (A.end < B.end)
+    A.end = B.end;
+}
+
 std::pair<size_t, size_t> offsetToClangLineColumn(llvm::StringRef Code,
                                                   size_t Offset) {
   Offset = std::min(Code.size(), Offset);

Diff for: clang-tools-extra/clangd/SourceCode.h

@@ -129,6 +129,9 @@ llvm::StringRef toSourceCode(const SourceManager &SM, SourceRange R);
 // Note that clang also uses closed source ranges, which this can't handle!
 Range halfOpenToRange(const SourceManager &SM, CharSourceRange R);
 
+// Expand range `A` to also contain `B`.
+void unionRanges(Range &A, Range B);
+
 // Converts an offset to a clang line/column (1-based, columns are bytes).
 // The offset must be in range [0, Code.size()].
 // Prefer to use SourceManager if one is available.