[clangd] Implement simple folding for preprocessor branches (#140959)

This pull request builds on
#121449 by sr-tream with
a unit test and a bug fix.

---------

Co-authored-by: Ruihua Dong <[email protected]>

Author: aketchum15

clang-tools-extra/clangd/SemanticSelection.cpp

@@ -175,9 +175,8 @@ llvm::Expected<std::vector<FoldingRange>> getFoldingRanges(ParsedAST &AST) {
   return collectFoldingRanges(SyntaxTree, TM);
 }
 
-// FIXME( usaxena95): Collect PP conditional regions, includes and other code
-// regions (e.g. public/private/protected sections of classes, control flow
-// statement bodies).
+// FIXME( usaxena95): Collect includes 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) {
@@ -186,12 +185,6 @@ getFoldingRanges(const std::string &Code, bool LineFoldingOnly) {
   auto DirectiveStructure = DirectiveTree::parse(OrigStream);
   chooseConditionalBranches(DirectiveStructure, OrigStream);
 
-  // FIXME: Provide ranges in the disabled-PP regions as well.
-  auto Preprocessed = DirectiveStructure.stripDirectives(OrigStream);
-
-  auto ParseableStream = cook(Preprocessed, genericLangOpts());
-  pairBrackets(ParseableStream);
-
   std::vector<FoldingRange> Result;
   auto AddFoldingRange = [&](Position Start, Position End,
                              llvm::StringLiteral Kind) {
@@ -220,7 +213,32 @@ getFoldingRanges(const std::string &Code, bool LineFoldingOnly) {
   auto EndPosition = [&](const Token &T) {
     return offsetToPosition(Code, EndOffset(T));
   };
+
+  // Preprocessor directives
+  auto PPRanges = pairDirectiveRanges(DirectiveStructure, OrigStream);
+  for (const auto &R : PPRanges) {
+    auto BTok = OrigStream.tokens()[R.Begin];
+    auto ETok = OrigStream.tokens()[R.End];
+    if (ETok.Kind == tok::eof)
+      continue;
+    if (BTok.Line >= ETok.Line)
+      continue;
+
+    Position Start = EndPosition(BTok);
+    Position End = StartPosition(ETok);
+    if (LineFoldingOnly)
+      End.line--;
+    AddFoldingRange(Start, End, FoldingRange::REGION_KIND);
+  }
+
+  // FIXME: Provide ranges in the disabled-PP regions as well.
+  auto Preprocessed = DirectiveStructure.stripDirectives(OrigStream);
+
+  auto ParseableStream = cook(Preprocessed, genericLangOpts());
+  pairBrackets(ParseableStream);
+
   auto Tokens = ParseableStream.tokens();
+
   // Brackets.
   for (const auto &Tok : Tokens) {
     if (auto *Paired = Tok.pair()) {
@@ -240,6 +258,7 @@ getFoldingRanges(const std::string &Code, bool LineFoldingOnly) {
     return OriginalToken(T).Length >= 2 &&
            Code.substr(StartOffset(T), 2) == "/*";
   };
+
   // Multi-line comments.
   for (auto *T = Tokens.begin(); T != Tokens.end();) {
     if (T->Kind != tok::comment) {

clang-tools-extra/clangd/support/DirectiveTree.cpp

@@ -356,5 +356,62 @@ TokenStream DirectiveTree::stripDirectives(const TokenStream &In) const {
   return Out;
 }
 
+namespace {
+class RangePairer {
+  std::vector<Token::Range> &Ranges;
+
+public:
+  RangePairer(std::vector<Token::Range> &Ranges) : Ranges(Ranges) {}
+
+  void walk(const DirectiveTree &T) {
+    for (const auto &C : T.Chunks)
+      std::visit(*this, C);
+  }
+
+  void operator()(const DirectiveTree::Code &C) {}
+
+  void operator()(const DirectiveTree::Directive &) {}
+
+  void operator()(const DirectiveTree::Conditional &C) {
+    Token::Range Range;
+    Token::Index Last;
+    auto First = true;
+    for (const auto &[Directive, _] : C.Branches) {
+      if (First) {
+        First = false;
+      } else {
+        Range = {Last, Directive.Tokens.Begin};
+        Ranges.push_back(Range);
+      }
+      Last = Directive.Tokens.Begin;
+    }
+
+    if (C.End.Kind != tok::pp_not_keyword) {
+      Range = {Last, C.End.Tokens.Begin};
+      Ranges.push_back(Range);
+    }
+
+    for (const auto &[_, SubTree] : C.Branches)
+      walk(SubTree);
+  }
+};
+} // namespace
+
+std::vector<Token::Range> pairDirectiveRanges(const DirectiveTree &Tree,
+                                              const TokenStream &Code) {
+  std::vector<Token::Range> Ranges;
+  RangePairer(Ranges).walk(Tree);
+
+  // Transform paired ranges to start with last token in its logical line
+  for (auto &R : Ranges) {
+    const Token *Tok = &Code.tokens()[R.Begin + 1];
+    while (Tok->Kind != tok::eof && !Tok->flag(LexFlags::StartsPPLine))
+      ++Tok;
+    Tok = Tok - 1;
+    R.Begin = Tok->OriginalIndex;
+  }
+  return Ranges;
+}
+
 } // namespace clangd
 } // namespace clang

clang-tools-extra/clangd/support/DirectiveTree.h

@@ -124,6 +124,10 @@ llvm::raw_ostream &operator<<(llvm::raw_ostream &,
 /// The choices are stored in Conditional::Taken nodes.
 void chooseConditionalBranches(DirectiveTree &, const TokenStream &Code);
 
+/// Pairs preprocessor conditional directives and computes their token ranges.
+std::vector<Token::Range> pairDirectiveRanges(const DirectiveTree &Tree,
+                                              const TokenStream &Code);
+
 } // namespace clangd
 } // namespace clang
 

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

@@ -371,6 +371,45 @@ TEST(FoldingRanges, PseudoParserWithoutLineFoldings) {
         //[[ foo
         /* bar */]]
       )cpp",
+      R"cpp(
+        //Ignore non-conditional directives
+        #define A 1
+
+        void func() {[[
+          int Variable = 100;
+
+          #ifdef FOO[[
+            Variable = 1;
+            #if 1[[
+                Variable = 4;
+            ]]#endif
+          ]]#else[[
+            Variable = 2;
+            //handle nested directives
+            #if 1[[
+              Variable = 3;
+            ]]#endif
+          ]]#endif
+
+
+          ]]}
+      )cpp",
+      R"cpp(
+        int Variable = 0;
+        #if defined(WALDO)
+            Variable = 1;
+        #
+      )cpp",
+      R"cpp(
+        int Variable = 0;
+        #if defined(WALDO)[[
+            Variable = 1;
+        ]]#elif 1[[
+            Variable = 2;
+        ]]#else
+            Variable = 3;
+        #
+      )cpp",
   };
   for (const char *Test : Tests) {
     auto T = Annotations(Test);