[SourceKit] Use recored #if regions in "active regions" request

* Record each IfConfig clause location info in SourceFile
* Update SILProfiler to handle them
* Update SwiftLangSupport::findActiveRegionsInFile() to use the recorded
  regions instead of walking into AST to find #if regions

rdar://118082146

Author: rintaro

include/swift/AST/SourceFile.h

@@ -49,36 +49,41 @@ enum class RestrictedImportKind {
 using ImportAccessLevel = std::optional<AttributedImport<ImportedModule>>;
 
 /// Stores range information for a \c #if block in a SourceFile.
-class IfConfigRangeInfo final {
-  /// The range of the entire \c #if block, including \c #else and \c #endif.
-  CharSourceRange WholeRange;
+class IfConfigClauseRangeInfo final {
+public:
+  enum ClauseKind {
+    // Active '#if', '#elseif', or '#else' clause.
+    ActiveClause,
+    // Inactive '#if', '#elseif', or '#else' clause.
+    InactiveClause,
+    // '#endif' directive.
+    EndDirective,
+  };
+
+private:
+  /// Source location of '#if', '#elseif', etc.
+  SourceLoc DirectiveLoc;
+  /// Character source location of body starts.
+  SourceLoc BodyLoc;
+  /// Location of the end of the body.
+  SourceLoc EndLoc;
 
-  /// The range of the active selected body, if there is one. This does not
-  /// include the outer syntax of the \c #if. This may be invalid, which
-  /// indicates there is no active body.
-  CharSourceRange ActiveBodyRange;
+  ClauseKind Kind;
 
 public:
-  IfConfigRangeInfo(CharSourceRange wholeRange, CharSourceRange activeBodyRange)
-      : WholeRange(wholeRange), ActiveBodyRange(activeBodyRange) {
-    assert(wholeRange.getByteLength() > 0 && "Range must be non-empty");
-    assert(activeBodyRange.isInvalid() || wholeRange.contains(activeBodyRange));
+  IfConfigClauseRangeInfo(SourceLoc DirectiveLoc, SourceLoc BodyLoc,
+                          SourceLoc EndLoc, ClauseKind Kind)
+      : DirectiveLoc(DirectiveLoc), BodyLoc(BodyLoc), EndLoc(EndLoc),
+        Kind(Kind) {
+    assert(DirectiveLoc.isValid() && BodyLoc.isValid() && EndLoc.isValid());
   }
 
-  CharSourceRange getWholeRange() const { return WholeRange; }
-  SourceLoc getStartLoc() const { return WholeRange.getStart(); }
+  SourceLoc getStartLoc() const { return DirectiveLoc; }
+  CharSourceRange getDirectiveRange(const SourceManager &SM) const;
+  CharSourceRange getWholeRange(const SourceManager &SM) const;
+  CharSourceRange getBodyRange(const SourceManager &SM) const;
 
-  friend bool operator==(const IfConfigRangeInfo &lhs,
-                         const IfConfigRangeInfo &rhs) {
-    return lhs.WholeRange == rhs.WholeRange &&
-           lhs.ActiveBodyRange == rhs.ActiveBodyRange;
-  }
-
-  /// Retrieve the ranges produced by subtracting the active body range from
-  /// the whole range. This includes both inactive branches as well as the
-  /// other syntax of the \c #if.
-  SmallVector<CharSourceRange, 2>
-  getRangesWithoutActiveBody(const SourceManager &SM) const;
+  ClauseKind getKind() const { return Kind; }
 };
 
 /// A file containing Swift source code.
@@ -244,9 +249,9 @@ class SourceFile final : public FileUnit {
   ParserStatePtr DelayedParserState =
       ParserStatePtr(/*ptr*/ nullptr, /*deleter*/ nullptr);
 
-  struct IfConfigRangesData {
+  struct IfConfigClauseRangesData {
     /// All the \c #if source ranges in this file.
-    std::vector<IfConfigRangeInfo> Ranges;
+    std::vector<IfConfigClauseRangeInfo> Ranges;
 
     /// Whether the elemnts in \c Ranges are sorted in source order within
     /// this file. We flip this to \c false any time a new range gets recorded,
@@ -255,7 +260,7 @@ class SourceFile final : public FileUnit {
   };
 
   /// Stores all the \c #if source range info in this file.
-  mutable IfConfigRangesData IfConfigRanges;
+  mutable IfConfigClauseRangesData IfConfigClauseRanges;
 
   friend class HasImportsMatchingFlagRequest;
 
@@ -500,12 +505,16 @@ class SourceFile final : public FileUnit {
      const_cast<SourceFile *>(this)->MissingImportedModules.insert(module);
   }
 
-  /// Record the source range info for a parsed \c #if block.
-  void recordIfConfigRangeInfo(IfConfigRangeInfo ranges);
+  /// Record the source range info for a parsed \c #if clause.
+  void recordIfConfigClauseRangeInfo(const IfConfigClauseRangeInfo &range);
+
+  /// Retrieve the source range info for any \c #if clauses in the file.
+  ArrayRef<IfConfigClauseRangeInfo> getIfConfigClauseRanges() const;
 
-  /// Retrieve the source range infos for any \c #if blocks contained within a
+  /// Retrieve the source range infos for any \c #if clauses contained within a
   /// given source range of this file.
-  ArrayRef<IfConfigRangeInfo> getIfConfigsWithin(SourceRange outer) const;
+  ArrayRef<IfConfigClauseRangeInfo>
+  getIfConfigClausesWithin(SourceRange outer) const;
 
   void getMissingImportedModules(
          SmallVectorImpl<ImportedModule> &imports) const override;

lib/AST/Module.cpp

@@ -2778,56 +2778,66 @@ SourceFile::getImportAccessLevel(const ModuleDecl *targetModule) const {
   return restrictiveImport;
 }
 
-SmallVector<CharSourceRange, 2>
-IfConfigRangeInfo::getRangesWithoutActiveBody(const SourceManager &SM) const {
-  SmallVector<CharSourceRange, 2> result;
-  if (ActiveBodyRange.isValid()) {
-    // Split the whole range by the active range.
-    result.emplace_back(SM, WholeRange.getStart(), ActiveBodyRange.getStart());
-    result.emplace_back(SM, ActiveBodyRange.getEnd(), WholeRange.getEnd());
-  } else {
-    // No active body, we just return the whole range.
-    result.push_back(WholeRange);
-  }
-  return result;
+CharSourceRange
+IfConfigClauseRangeInfo::getDirectiveRange(const SourceManager &SM) const {
+  return CharSourceRange(SM, DirectiveLoc, BodyLoc);
 }
 
-void SourceFile::recordIfConfigRangeInfo(IfConfigRangeInfo ranges) {
-  IfConfigRanges.Ranges.push_back(ranges);
-  IfConfigRanges.IsSorted = false;
+CharSourceRange
+IfConfigClauseRangeInfo::getBodyRange(const SourceManager &SM) const {
+  return CharSourceRange(SM, BodyLoc, EndLoc);
 }
 
-ArrayRef<IfConfigRangeInfo>
-SourceFile::getIfConfigsWithin(SourceRange outer) const {
-  auto &SM = getASTContext().SourceMgr;
-  assert(SM.getRangeForBuffer(BufferID).contains(outer.Start) &&
-         "Range not within this file?");
+CharSourceRange
+IfConfigClauseRangeInfo::getWholeRange(const SourceManager &SM) const {
+  return CharSourceRange(SM, DirectiveLoc, EndLoc);
+}
+
+void SourceFile::recordIfConfigClauseRangeInfo(
+    const IfConfigClauseRangeInfo &range) {
+  IfConfigClauseRanges.Ranges.push_back(range);
+  IfConfigClauseRanges.IsSorted = false;
+}
 
-  if (!IfConfigRanges.IsSorted) {
+ArrayRef<IfConfigClauseRangeInfo> SourceFile::getIfConfigClauseRanges() const {
+  if (!IfConfigClauseRanges.IsSorted) {
+    auto &SM = getASTContext().SourceMgr;
     // Sort the ranges if we need to.
-    llvm::sort(IfConfigRanges.Ranges, [&](IfConfigRangeInfo lhs,
-                                          IfConfigRangeInfo rhs) {
-      return SM.isBeforeInBuffer(lhs.getStartLoc(), rhs.getStartLoc());
-    });
+    llvm::sort(
+        IfConfigClauseRanges.Ranges, [&](const IfConfigClauseRangeInfo &lhs,
+                                         const IfConfigClauseRangeInfo &rhs) {
+          return SM.isBeforeInBuffer(lhs.getStartLoc(), rhs.getStartLoc());
+        });
 
     // Be defensive and eliminate duplicates in case we've parsed twice.
-    auto newEnd = std::unique(
-        IfConfigRanges.Ranges.begin(), IfConfigRanges.Ranges.end(),
-        [&](const IfConfigRangeInfo &lhs, const IfConfigRangeInfo &rhs) {
-          if (lhs.getWholeRange() != rhs.getWholeRange())
+    auto newEnd = llvm::unique(
+        IfConfigClauseRanges.Ranges, [&](const IfConfigClauseRangeInfo &lhs,
+                                         const IfConfigClauseRangeInfo &rhs) {
+          if (lhs.getStartLoc() != rhs.getStartLoc())
             return false;
-
-          assert(lhs == rhs && "Active ranges changed on a re-parse?");
+          assert(lhs.getBodyRange(SM) == rhs.getBodyRange(SM) &&
+                 "range changed on a re-parse?");
           return true;
         });
-    IfConfigRanges.Ranges.erase(newEnd, IfConfigRanges.Ranges.end());
-    IfConfigRanges.IsSorted = true;
+    IfConfigClauseRanges.Ranges.erase(newEnd,
+                                      IfConfigClauseRanges.Ranges.end());
+    IfConfigClauseRanges.IsSorted = true;
   }
 
+  return IfConfigClauseRanges.Ranges;
+}
+
+ArrayRef<IfConfigClauseRangeInfo>
+SourceFile::getIfConfigClausesWithin(SourceRange outer) const {
+  auto &SM = getASTContext().SourceMgr;
+  assert(SM.getRangeForBuffer(BufferID).contains(outer.Start) &&
+         "Range not within this file?");
+
   // First let's find the first #if that is after the outer start loc.
-  auto ranges = llvm::ArrayRef(IfConfigRanges.Ranges);
+  auto ranges = getIfConfigClauseRanges();
   auto lower = llvm::lower_bound(
-      ranges, outer.Start, [&](IfConfigRangeInfo range, SourceLoc loc) {
+      ranges, outer.Start,
+      [&](const IfConfigClauseRangeInfo &range, SourceLoc loc) {
         return SM.isBeforeInBuffer(range.getStartLoc(), loc);
       });
   if (lower == ranges.end() ||
@@ -2836,7 +2846,8 @@ SourceFile::getIfConfigsWithin(SourceRange outer) const {
   }
   // Next let's find the first #if that's after the outer end loc.
   auto upper = llvm::upper_bound(
-      ranges, outer.End, [&](SourceLoc loc, IfConfigRangeInfo range) {
+      ranges, outer.End,
+      [&](SourceLoc loc, const IfConfigClauseRangeInfo &range) {
         return SM.isBeforeInBuffer(loc, range.getStartLoc());
       });
   return llvm::ArrayRef(lower, upper - lower);

lib/Parse/ParseIfConfig.cpp

@@ -889,10 +889,12 @@ Result Parser::parseIfConfigRaw(
           ClauseLoc, Condition, isActive, IfConfigElementsRole::Skipped);
     }
 
-    // Record the active body range for the SourceManager.
-    if (shouldEvaluate && isActive) {
-      assert(!activeBodyRange.isValid() && "Multiple active regions?");
-      activeBodyRange = CharSourceRange(SourceMgr, bodyStart, Tok.getLoc());
+    // Record the clause range info in SourceFile.
+    if (shouldEvaluate) {
+      auto kind = isActive ? IfConfigClauseRangeInfo::ActiveClause
+                           : IfConfigClauseRangeInfo::InactiveClause;
+      SF.recordIfConfigClauseRangeInfo(
+          {ClauseLoc, bodyStart, Tok.getLoc(), kind});
     }
 
     if (Tok.isNot(tok::pound_elseif, tok::pound_else))
@@ -905,11 +907,11 @@ Result Parser::parseIfConfigRaw(
   SourceLoc EndLoc;
   bool HadMissingEnd = parseEndIfDirective(EndLoc);
 
-  // Record the #if ranges on the SourceManager.
+  // Record the '#end' ranges in SourceFile.
   if (!HadMissingEnd && shouldEvaluate) {
-    auto wholeRange = Lexer::getCharSourceRangeFromSourceRange(
-        SourceMgr, SourceRange(startLoc, EndLoc));
-    SF.recordIfConfigRangeInfo({wholeRange, activeBodyRange});
+    SourceLoc EndOfEndLoc = getEndOfPreviousLoc();
+    SF.recordIfConfigClauseRangeInfo({EndLoc, EndOfEndLoc, EndOfEndLoc,
+                                      IfConfigClauseRangeInfo::EndDirective});
   }
   return finish(EndLoc, HadMissingEnd);
 }

lib/SIL/IR/SILProfiler.cpp

@@ -1165,14 +1165,37 @@ struct CoverageMapping : public ASTWalker {
                                            Counter.getLLVMCounter()));
     }
     // Add any skipped regions present in the outer range.
-    for (auto IfConfig : SF->getIfConfigsWithin(OuterRange)) {
-      for (auto SkipRange : IfConfig.getRangesWithoutActiveBody(SM)) {
-        auto Start = SM.getLineAndColumnInBuffer(SkipRange.getStart());
-        auto End = SM.getLineAndColumnInBuffer(SkipRange.getEnd());
-        assert(Start.first <= End.first && "region start and end out of order");
-        Regions.push_back(MappedRegion::skipped(Start.first, Start.second,
-                                                End.first, End.second));
+    for (auto clause : SF->getIfConfigClausesWithin(OuterRange)) {
+      CharSourceRange SkipRange;
+      switch (clause.getKind()) {
+      case IfConfigClauseRangeInfo::ActiveClause:
+      case IfConfigClauseRangeInfo::EndDirective:
+        SkipRange = clause.getDirectiveRange(SM);
+        break;
+      case IfConfigClauseRangeInfo::InactiveClause:
+        SkipRange = clause.getWholeRange(SM);
+        break;
+      }
+      if (SkipRange.getByteLength() == 0)
+        continue;
+
+      auto Start = SM.getLineAndColumnInBuffer(SkipRange.getStart());
+      auto End = SM.getLineAndColumnInBuffer(SkipRange.getEnd());
+      assert(Start.first <= End.first && "region start and end out of order");
+
+      // If this is consecutive with the last one, expand it.
+      if (!Regions.empty()) {
+        auto &last = Regions.back();
+        if (last.RegionKind == MappedRegion::Kind::Skipped &&
+            last.EndLine == Start.first && last.EndCol == Start.second) {
+          last.EndLine = End.first;
+          last.EndCol = End.second;
+          continue;
+        }
       }
+
+      Regions.push_back(MappedRegion::skipped(Start.first, Start.second,
+                                              End.first, End.second));
     }
     return SILCoverageMap::create(M, SF, Filename, Name, PGOFuncName, Hash,
                                   Regions, CounterBuilder.getExpressions());

test/SourceKit/ActiveRegions/nested2.swift

@@ -0,0 +1,19 @@
+func foo() -> Int {
+  #if true
+  #if true
+  return 1
+  #else
+  return 2
+  #endif
+  #else
+  return 3
+  #endif
+}
+
+// RUN: %sourcekitd-test -req=active-regions %s -- -D FLAG_1 -module-name active_regions %s | %FileCheck -check-prefix=CHECK1 %s
+// CHECK1: START IF CONFIGS
+// CHECK1-NEXT: 2:3 - active
+// CHECK1-NEXT: 3:3 - active
+// CHECK1-NEXT: 5:3 - inactive
+// CHECK1-NEXT: 8:3 - inactive
+// CHECK1-NEXT: END IF CONFIGS

test/SourceKit/ActiveRegions/postfix.swift

@@ -0,0 +1,18 @@
+struct Value {
+  var prop1: Int = 0
+  var prop2: Int = 0
+}
+func test(value: Value) {
+  let _ = value
+#if false
+    .prop1
+#else
+    .prop2
+#endif
+}
+
+// RUN: %sourcekitd-test -req=active-regions %s -- -module-name active_regions %s | %FileCheck -check-prefix=CHECK1 %s
+// CHECK1: START IF CONFIGS
+// CHECK1-NEXT: 7:1 - inactive
+// CHECK1-NEXT: 9:1 - active
+// CHECK1-NEXT: END IF CONFIGS