[clangd] support folding multi-line `#include`

clang-tools-extra/clangd/SemanticSelection.cpp

 llvm::Expected<std::vector<FoldingRange>> getFoldingRanges(ParsedAST &AST) {
   syntax::Arena A;
   syntax::TokenBufferTokenManager TM(AST.getTokens(), AST.getLangOpts(),
                                      AST.getSourceManager());
   const auto *SyntaxTree = syntax::buildSyntaxTree(A, TM, AST.getASTContext());
   return collectFoldingRanges(SyntaxTree, TM);
 }
 
-// FIXME( usaxena95): Collect PP conditional regions, includes and other code
+namespace {
+struct IncludeCollector {
+  explicit IncludeCollector(const pseudo::DirectiveTree &T) { walk(T); }
+
+  void walk(const pseudo::DirectiveTree &T) {
+    for (const auto &C : T.Chunks)
+      std::visit(*this, C);
+  }
+
+  void operator()(const pseudo::DirectiveTree::Code &) {}
+
+  void operator()(const pseudo::DirectiveTree::Directive &D) {
+    if (D.Kind == tok::pp_include)
+      Results.push_back(D);
+  }
+
+  void operator()(const pseudo::DirectiveTree::Conditional &C) {
+    if (C.Taken)
+      walk(C.Branches[*C.Taken].second);
+  }
+
+  std::vector<pseudo::DirectiveTree::Directive> Results;
+};
+} // namespace
+
+// FIXME( usaxena95): Collect PP conditional regions and other code
 // regions (e.g. public/private/protected sections of classes, control flow
 // statement bodies).
 // Related issue: https://github.com/clangd/clangd/issues/310
 llvm::Expected<std::vector<FoldingRange>>
 getFoldingRanges(const std::string &Code, bool LineFoldingOnly) {
   auto OrigStream = pseudo::lex(Code, clang::pseudo::genericLangOpts());
 
   auto DirectiveStructure = pseudo::DirectiveTree::parse(OrigStream);
if (IsBlockComment(*FirstComment)) {
         // Show last line of a block comment.
         End.line--;
       if (IsBlockComment(*LastComment))
         // Show ending sentinal "*/" of the block comment.
         End.character -= 2;
     }
     AddFoldingRange(Start, End, FoldingRange::COMMENT_KIND);
   }
+  // Multi-line `#include`
+  auto Includes = IncludeCollector{DirectiveStructure}.Results;
+  for (auto It = Includes.begin(); It != Includes.end();) {
+    Position Start = StartPosition(OrigStream.tokens(It->Tokens).front());
+    Position End = EndPosition(OrigStream.tokens(It->Tokens).back());
+    It++;
+    while (It != Includes.end() &&
+           StartPosition(OrigStream.tokens(It->Tokens).front()).line ==
+               End.line + 1) {
+      End = EndPosition(OrigStream.tokens(It->Tokens).back());
+      It++;
+    }
+    AddFoldingRange(Start, End, FoldingRange::IMPORT_KIND);
+  }
   return Result;
 }
 
 } // namespace clangd
 } // namespace clang

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

for (const char *Test : Tests) {
     auto T = Annotations(Test);
     EXPECT_THAT(
         StripColumns(gatherFoldingRanges(llvm::cantFail(
             getFoldingRanges(T.code().str(), /*LineFoldingsOnly=*/true)))),
         UnorderedElementsAreArray(StripColumns(T.ranges())))
         << Test;
   }
 }
+
+TEST(FoldingRanges, Includes) {
+  const char *Test = R"cpp(
+    [[#include <cmath>
+    #include <cstdio>
+    #include <iostream>]]
+
+    #include <vector>
+
+    [[#include "external/Logger.h"
+    #include "external/Vector.h"
+    #include "project/DataStructures/Trie.h"]]
+
+    [[#include "project/File.h"
+    #include <map>
+    #include <set>]]
+
+    #include "math.h"
+  )cpp";
+  auto T = Annotations{Test};
+  EXPECT_THAT(gatherFoldingRanges(
+                  llvm::cantFail(getFoldingRanges(T.code().str(), false))),
+              UnorderedElementsAreArray(T.ranges()))
+      << Test;
+}
 } // namespace
 } // namespace clangd
 } // namespace clang